Energon-AI

A Large-scale Model Inference System. Energon-AI provides 3 levels of abstraction for enabling the large-scale model inference:

• Runtime - tensor parallel operations, pipeline parallel wrapper, distributed message queue, distributed checkpoint loading, customized CUDA kernels.
• Engine - encapsulate the single instance multiple devices (SIMD) execution with the remote procedure call, which acts as the single instance single device (SISD) execution.
• Serving - batching requests, managing engines.

For models trained by Colossal-AI, they can be seamlessly transferred to Energon-AI. For single-device models, they require manual coding works to introduce tensor parallelism and pipeline parallelism.

At present, we pre-build distributed Bert, GPT, and ViT models.
For GPT, it extends to at most 175B parameters, which is called GPT3.
For Bert, Google reports a super-large Bert with 481B parameters in MLPerf-Training v1.1 open, indicating that Bert can also extend to large-scale.

Installation

$ git clone [email protected]:hpcaitech/EnergonAI.git
$ pip install -r requirements.txt
$ pip install .

Huggingface GPT2 Generation Task Case

# Download checkpoint
$ wget https://huggingface.co/gpt2/blob/main/pytorch_model.bin
# Download files for tokenizer
$ wget https://huggingface.co/gpt2/blob/main/tokenizer.json
$ wget https://huggingface.co/gpt2/blob/main/vocab.json
$ wget https://huggingface.co/gpt2/blob/main/merges.txt

# Launch the service
export PYTHONPATH=~/ColossalAI-Inference/examples/hf_gpt2
energonai service init --config_file=~/ColossalAI-Inference/hf_gpt2/hf_gpt2_config.py

# Request for the service
Method 1: 
    FastAPI provides an automatic API docs, you can forward 
    http://127.0.0.1:8005/docs and make request with the graphical interface.
Method 2:
    curl -X 'GET' \
    'http://127.0.0.1:8005/run_hf_gpt2/I%20do%20not?max_seq_length=16' \
    -H 'accept: application/json' 

Large-scale Model Inference Performance

Scaling Ability

Here GPT3-12-layers in FP16 is adopted.
Here a node with 8 A100 80 GB GPUs is adopted. GPUs are fully connected with NvLink.
Energon-AI adopts the redundant computation elimination method. The method is first raised in EffectiveTransformer, and our implementation refers to TurboTransformer.
Here the sequence length is set the half of the padding length.

Latency

Here GPT3 in FP16 is adopted.
Here a node with 8 A100 80 GB GPUs is adopted. Every two GPUs are connected with NvLink.
Here the sequence length is set the half of the padding length when using redundant computation elimination method, which is the Energon-AI(RM).
Here FasterTransformer is adopted in comparison and it does not support the redundant computation elimination method in the distributed execution.

Batching

Energon-AI dynamically selects the batch processing with the highest priority regarding the waiting time, batch size, batch expansion possibility (based on the sentence length after padding). Our dynamic batching method is inspired by the DP algorithm from TurboTransformer.
Here FIFO batching is selected in comparison.

Contributing

If interested in making your own contribution to the project, please refer to Contributing for guidance.

Thanks so much!

Technical Overview