推理加速效能超越英偉達FasterTransformer 50%，開源方案打通大模型落地關鍵路徑

語言: CN / TW / HK

時間 2022-06-01 09:02:32 潞晨科技

伴隨著深度學習模型規模的指數型增長，常見的單卡推理解決方案已然無法滿足前沿AI大模型的推理需求。例如1750億引數的GPT-3模型，僅僅是載入模型引數就需要數百GB的儲存空間，遠超單個GPU的容納能力。因此，對於AI大模型使用多卡並行的方式進行推理已成為必然選擇。

針對現有推理系統的這一痛點，Colossal-AI團隊以“高效能、高可用、可伸縮”的理念，深入單例項多裝置推理場景，開發了大模型推理系統Energon-AI，在效能和易用性上兼具優勢：

僅需對現有專案進行極 少量修改 ， 使用者 便可 完成自定義大模型 的 推理部署，獲得 並行擴充套件的超線性加速，對於 AI 大模型 分散式 推理加速，相比英偉達FasterTransformer可提升50%以上。

相比現有推理方案，Energon-AI不再需要使用者對通訊、記憶體等各部分協作進行手動管理，也無需額外編譯，大幅降低了使用者的使用門檻。

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

AI大模型推理部署的困難

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

近年來，計算裝置（如GPU）的平行計算能力、記憶體容量，記憶體速度等都得到了極大的增強，然而，單裝置縱向擴充套件（scale up）的效能增益在面對指數型增長的模型規模時，仍難以滿足大模型的記憶體與效能需求。而當前的深度學習推理系統，主要面向多例項單裝置以及單例項單裝置的簡單推理場景，忽視了AI大模型推理所需要的單例項多裝置的挑戰與機遇，Energon-AI系統正是為了解決這一痛點而生。

Energon-AI系統設計

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

面向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。

我們模擬真實場景下多使用者同時傳送大量變長推理請求的情況，將我們的動態batch規劃方法與傳統的FIFO(先入先出)佇列打包方法進行了吞吐量對比。由於dynamic batching的演算法緩解了直接padding造成的大量冗餘計算問題，在該策略下dynamic batching的吞吐量實現了34.7%的提升。

易用性

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

energonai service init --config_file=gpt_config.py

在追求效能的同時，Energon-AI希望保持系統使用的靈活度與易用性，使用者僅需自定義【並行模型】、【並行引數】以及【服務請求邏輯】加入到配置檔案中，即可啟動推理服務。目前，我們提供了最常見的GPT、BERT和ViT模型作為示例，更詳盡的教程將會在近期完善。

在構建新的並行模型時，Energon-AI使用Python，且使用方式與Pytorch相似，有層的概念且初始化與執行邏輯清晰，使用者無需考慮記憶體管理，並行通訊等行為。如下程式碼展示了兩層Linear層組成的模型並行執行的完整程式碼。

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++程式碼還需要額外編譯工作。這些都對使用者的並行知識與程式設計能力提出了嚴峻挑戰。

// 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);
}

構建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也廣受關注，榮登熱榜第一。

專案團隊

潞晨技術團隊的核心成員均來自美國加州大學伯克利分校，斯坦福大學，清華大學，北京大學，新加坡國立大學，新加坡南洋理工大學等國內外知名高校；擁有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名廠商工作經歷。公司成立即獲得創新工場、真格基金等多家頂尖VC機構種子輪投資。