本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

在前沿AI大模型規模呈指數級增長的趨勢下，僅憑單卡已經很難滿足當下的推理需求。

就拿擁有1750億參數的GPT-3來說。

僅僅是加載模型參數就需要數百GB的存儲空間，遠超單個GPU的容納能力。

因此，多卡并行被視為AI大模型推理的必然選擇。

但現有的推理系統仍舊存在不少弊端。

比如需要用戶對通信、內存等各部分協作進行手動管理，需要額外編譯等……導致用戶使用門檻居高不下。

為此，大規模并行AI訓練系統Colossal-AI團隊提出了大模型推理系統Energon-AI。

以“高性能、高可用、可伸縮”的理念，深入單實例多設備推理場景，Energon-AI在性能和易用性上兼具優勢。

僅需對現有項目進行極少量修改，用戶就能完成自定義大模型的推理部署，獲得并行擴展的超線性加速。

對于AI大模型分布式推理加速，相比英偉達FasterTransformer可提升50%以上。

開源地址：https://github.com/hpcaitech/ColossalAI

AI大模型推理部署的困難

近年來，計算設備（如GPU）的并行計算能力、內存容量，內存速度等都得到了極大的增強。

然而，單設備縱向擴展（scale up）的性能增益在面對指數型增長的模型規模時，仍難以滿足大模型的內存與性能需求。

而當前的深度學習推理系統，主要面向多實例單設備以及單實例單設備的簡單推理場景，忽視了AI大模型推理所需要的單實例多設備的挑戰與機遇，Energon-AI系統正是為了解決這一痛點而生。

△模型參數的迅速增長[https://arxiv.org/abs/2111.14247]

Energon-AI系統設計

面向AI大模型部署，Colossal-AI團隊設計了單實例多設備推理系統Energon-AI。

△Energon-AI超大模型推理系統示意圖

Energon-AI系統設計分為三個層次，即運行時系統（Runtime）、分布式推理實例（Engine）以及前端服務系統（Serving）：

• Runtime

在運行時系統設計過程中，團隊發現當模型規模不斷增大，通用矩陣乘的時間占比逐漸增大。

而訪存密集型算子與Kernel Launch的時間占比則逐漸降低，推理過程進一步從訪存密集型向計算密集型方向遷移，TensorRT以及專用推理系統對訪存密集型操作的優化效果被極大削減。

Energon-AI Runtime依賴于Colossal-AI實現張量并行，同時設計了流水線并行包裝方法用于顯存不足的情況。

此外，團隊引入了大量推理專用算子及方法。

如面對NLP中輸入變長的特點，引入transpose_padding_rebulid與transpose_padding_remove等算子用以高效支持Encoder和Decoder模型中MLP層的冗余計算消除方法。

• Engine

單設備推理中程序有相同的數據入口與出口，分布式訓練的主要目標是模型參數，因此無須對多個進程的輸入輸出進行管理，而多設備推理則不同。

團隊希望通過良好的封裝使得Engine具有與單設備推理完全相同的行為。

采用半中心化方法，主進程中使用RPC在每個設備調用初始化或推理方法，分布式推理就可以得到中心化的控制，同時每個設備則保有自己的Tensor Parallel與Pipeline Parallel通信邏輯。

同時，團隊在每個進程中設計并維護了分布式消息隊列，用以保證多個進程中多線程調用執行的一致性。

• Serving：

針對用戶請求分散和變長的特點及大模型推理對GPU并行運算的依賴之間的矛盾，Energon-AI引入了動態Batching機制。

將請求隊列中的請求按照機器性能進行最優打包后，根據等候時間、batch大小、batch的擴展可能性（根據padding后的句子長度）等挑選優先級最高的batch處理。

這樣一來，可以在最大化GPU使用率的同時規避饑餓問題，減小平均請求時延。

△Batch管理流程示意圖

性能測試

并行推理超線性擴展

△張量并行可擴展性測試結果展示

硬件環境：8 * A100 GPU 80GB。

由于單設備顯存無法滿足GPT-3推理需求，此處為GPT-3 12層的測試結果，設置句長為Padding的1/2。

Energon-AI八卡并行推理在Batch Size為32時，相比于單卡Pytorch直接推理，可獲得8.5倍的超線性加速。

運行時推理性能提升50%

△張量并行運行時系統推理時延對比

硬件環境：8 * A100 GPU 80GB。

設置句長為Padding的1/2。GPT-3-24-Layers for TP=2, GPT-3-48-Layers for TP=4。

以高度優化的英偉達FasterTransformer GPT-3作為對比方案。

FasterTransformer在其4.0版本中推出了分布式推理特性，目前支持GPT-3模型的分布式推理，但由于其純C++代碼高度耦合的特點，靈活度與易用性相對較低。

此外，對于NLP推理輸入句長不同的特點，其分布式推理無冗余計算消除功能。

對于GPT-3模型，Energon-AI的運行時系統在Batch Size為1時性能略低于FasterTransformer，而在Batch Size較大時能夠實現超過50%的性能提升。

Dynamic Batching吞吐量增加30%

△Dynamic batching與直接打包batch吞吐量對比

硬件環境：8 * A100 GPU 80GB。

測試使用的模型為GPT-3, 測試句長為256以內隨機生成，padding策略為batch內最長padding。

模擬真實場景下多用戶同時發送大量變長推理請求的情況，將Energon-AI的動態batch規劃方法與傳統的FIFO(先入先出)隊列打包方法進行了吞吐量對比。

由于dynamic batching的算法緩解了直接padding造成的大量冗余計算問題，在該策略下dynamic batching的吞吐量實現了34.7%的提升。

易用性

Python
from gpt import gpt3
from gpt_server import launch_engine
# for engine
model_class = gpt3
model_type = "gpt"
host = "127.0.0.1"
port = 29400
half = True
backend = "nccl"
# for parallel
tp_init_size = 4
pp_init_size = 2
# for server
engine_server = launch_engine
server_host = "127.0.0.1"
server_port = 8020
rm_padding = True

Python
energonai service init --config_file=gpt_config.py

在追求性能的同時，Energon-AI希望保持系統使用的靈活度與易用性，用戶僅需自定義并行模型、并行參數以及服務請求邏輯加入到配置文件中，即可啟動推理服務。

目前，已經提供了最常見的GPT、BERT和ViT模型作為示例，更詳盡的教程將會在近期完善。

在構建新的并行模型時，Energon-AI使用Python，且使用方式與Pytorch相似，有層的概念且初始化與執行邏輯清晰，用戶無需考慮內存管理，并行通信等行為。

如下代碼展示了兩層Linear層組成的模型并行運行的完整代碼。

Python
class MLP(nn.Module):
def __init__(self, dim, dtype, bias):
super().__init__()
self.dense_0 = Linear1D_Col(dim, dim, dtype=dtype, bias=bias, gather_output=False)
self.dense_1 = Linear1D_Row(dim, dim, dtype=dtype, bias=bias, parallel_input=True)
def forward(self, x):
x = self.dense_0(x)
x = self.dense_1(x)
return x

與之相對，在構建新的并行模型時，FasterTransformer需要使用C++代碼并且需要用戶自行進行內存管理，定義通信等底層行為組織。

受篇幅限制，如下代碼展示兩層Linear層模型并行運行的內存管理，具體執行，通信的部分代碼。

除此之外，用戶想要代碼正確執行，還需要花費大量時間精力對內存管理、執行邏輯、通信行為之間的配合進行調試，C++代碼還需要額外編譯工作。

這些都對用戶的并行知識與編程能力提出了嚴峻挑戰。

C++
// Memory Allocation (only for a single paramerter).
T *d_inter_kernel = NULL
param_.ffn.intermediate_weight.kernel = d_inter_kernel;
device_malloc(&d_inter_kernel, dim * dim);
// Two MLP Layers
cublasMM_cublasLtMM_wrapper(param_.cublaslt_handle, param_.cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, param_.ffn.intermediate_weight.kernel, AType_, n, attr_matmul_buf_, BType_, k, &beta, (DataType_ *)inter_matmul_buf_, CType_, n, param_.stream, cublasAlgoMap_, sm_, cublas_workspace_);
add_bias_act_kernelLauncher<DataType_>(inter_matmul_buf_, param_.ffn.intermediate_weight.bias, m, n, ActivationType::GELU, param_.stream);
n = k;
cublasMM_cublasLtMM_wrapper(param_.cublaslt_handle, param_.cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, param_.ffn.output_weight.kernel, AType_, n, inter_matmul_buf_, BType_, k, &beta, (DataType_ *)(param_.transformer_out), CType_, n, param_.stream, cublasAlgoMap_, sm_, cublas_workspace_);
add_bias_input_layernorm_kernelLauncher<DataType_>(param_.transformer_out, attr_matmul_buf_, param_.ffn.output_weight.bias, param_.ffn_layernorm.gamma, param_.ffn_layernorm.beta, m, n, param_.stream);
// Communication
if(t_parallel_param_.world_size > 1)
{
all2all_gather(nccl_logits_buf_, nccl_logits_buf_, local_batch * n, t_parallel_param_, decoding_params.stream);
}

更多特性

本次發布的Energon-AI子系統為beta版，近期會根據用戶反饋與既定計劃，進行密集的迭代更新，盡早為用戶提供正式版，充分滿足用戶的不同推理部署需求，歡迎向Energon-AI提出您的需求與建議。

構建AI大模型生態系統

面對AI大模型的時代浪潮，除了本次新增的推理部署特性，針對現有大模型訓練方案并行維度有限、效率不高、通用性差、部署困難、缺乏維護等痛點，Colossal-AI通過高效多維并行和異構并行等技術，讓用戶僅需極少量修改，即可高效快速部署AI大模型訓練。

例如對于GPT-3這樣的超大AI模型，相比英偉達方案，Colossal-AI僅需一半的計算資源，即可啟動訓練；若使用相同計算資源，則能提速11%，可降低GPT-3訓練成本超百萬美元。

對于蛋白質結構預測應用AlphaFold，基于Colossal-AI的加速方案的FastFold，成功超越谷歌和哥倫比亞大學的方案，將AlphaFold訓練時間從11天減少到67小時。

且總成本更低，在長序列推理中也實現9.3～11.6倍的速度提升。

Colossal-AI兼容低端設備，在僅有一塊GPU的個人PC上便能訓練高達180億參數GPT；普通的筆記本電腦，也能訓練十幾億參數的模型。

相比現有主流方案，可提升參數容量十余倍，大幅度降低了AI大模型微調和推理等下游任務和應用部署的門檻。

Colossal-AI注重開源社區建設，提供中文教程，開放用戶社群及論壇，對于用戶反饋進行高效交流與迭代更新，不斷添加等前沿特性。

自開源以來，Colossal-AI已經多次登上GitHub熱榜Python方向世界第一，與眾多已有數萬star的明星開源項目一起受到海內外關注！

在反映機器學習領域熱點的Papers With Code網站上，Colossal-AI也廣受關注，登上熱榜第一。

傳送門

GitHub地址：

https://github.com/hpcaitech/ColossalAI