推理加速GPT-3超越英偉達(dá)方案50%！最新大模型推理系統(tǒng)Energon-AI開(kāi)源，來(lái)自Colossal-AI團(tuán)隊(duì)

在前沿AI大模型規(guī)模呈指數(shù)級(jí)增長(zhǎng)的趨勢(shì)下，僅憑單卡已經(jīng)很難滿足當(dāng)下的推理需求。

就拿擁有1750億參數(shù)的GPT-3來(lái)說(shuō)。

僅僅是加載模型參數(shù)就需要數(shù)百GB的存儲(chǔ)空間，遠(yuǎn)超單個(gè)GPU的容納能力。

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

但現(xiàn)有的推理系統(tǒng)仍舊存在不少弊端。

比如需要用戶(hù)對(duì)通信、內(nèi)存等各部分協(xié)作進(jìn)行手動(dòng)管理，需要額外編譯等……導(dǎo)致用戶(hù)使用門(mén)檻居高不下。

為此，大規(guī)模并行AI訓(xùn)練系統(tǒng)Colossal-AI團(tuán)隊(duì)提出了大模型推理系統(tǒng)Energon-AI。

以“高性能、高可用、可伸縮”的理念，深入單實(shí)例多設(shè)備推理場(chǎng)景，Energon-AI在性能和易用性上兼具優(yōu)勢(shì)。

僅需對(duì)現(xiàn)有項(xiàng)目進(jìn)行極少量修改，用戶(hù)就能完成自定義大模型的推理部署，獲得并行擴(kuò)展的超線性加速。

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

開(kāi)源地址：https://github.com/hpcaitech/ColossalAI

AI大模型推理部署的困難

近年來(lái)，計(jì)算設(shè)備（如GPU）的并行計(jì)算能力、內(nèi)存容量，內(nèi)存速度等都得到了極大的增強(qiáng)。

然而，單設(shè)備縱向擴(kuò)展（scale up）的性能增益在面對(duì)指數(shù)型增長(zhǎng)的模型規(guī)模時(shí)，仍難以滿足大模型的內(nèi)存與性能需求。

而當(dāng)前的深度學(xué)習(xí)推理系統(tǒng)，主要面向多實(shí)例單設(shè)備以及單實(shí)例單設(shè)備的簡(jiǎn)單推理場(chǎng)景，忽視了AI大模型推理所需要的單實(shí)例多設(shè)備的挑戰(zhàn)與機(jī)遇，Energon-AI系統(tǒng)正是為了解決這一痛點(diǎn)而生。

Energon-AI系統(tǒng)設(shè)計(jì)

面向AI大模型部署，Colossal-AI團(tuán)隊(duì)設(shè)計(jì)了單實(shí)例多設(shè)備推理系統(tǒng)Energon-AI。

Energon-AI系統(tǒng)設(shè)計(jì)分為三個(gè)層次，即運(yùn)行時(shí)系統(tǒng)（Runtime）、分布式推理實(shí)例（Engine）以及前端服務(wù)系統(tǒng)（Serving）：

• Runtime

在運(yùn)行時(shí)系統(tǒng)設(shè)計(jì)過(guò)程中，團(tuán)隊(duì)發(fā)現(xiàn)當(dāng)模型規(guī)模不斷增大，通用矩陣乘的時(shí)間占比逐漸增大。

而訪存密集型算子與Kernel Launch的時(shí)間占比則逐漸降低，推理過(guò)程進(jìn)一步從訪存密集型向計(jì)算密集型方向遷移，TensorRT以及專(zhuān)用推理系統(tǒng)對(duì)訪存密集型操作的優(yōu)化效果被極大削減。

Energon-AI Runtime依賴(lài)于Colossal-AI實(shí)現(xiàn)張量并行，同時(shí)設(shè)計(jì)了流水線并行包裝方法用于顯存不足的情況。

此外，團(tuán)隊(duì)引入了大量推理專(zhuān)用算子及方法。

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

• Engine

單設(shè)備推理中程序有相同的數(shù)據(jù)入口與出口，分布式訓(xùn)練的主要目標(biāo)是模型參數(shù)，因此無(wú)須對(duì)多個(gè)進(jìn)程的輸入輸出進(jìn)行管理，而多設(shè)備推理則不同。

團(tuán)隊(duì)希望通過(guò)良好的封裝使得Engine具有與單設(shè)備推理完全相同的行為。

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

同時(shí)，團(tuán)隊(duì)在每個(gè)進(jìn)程中設(shè)計(jì)并維護(hù)了分布式消息隊(duì)列，用以保證多個(gè)進(jìn)程中多線程調(diào)用執(zhí)行的一致性。

• Serving：

針對(duì)用戶(hù)請(qǐng)求分散和變長(zhǎng)的特點(diǎn)及大模型推理對(duì)GPU并行運(yùn)算的依賴(lài)之間的矛盾，Energon-AI引入了動(dòng)態(tài)Batching機(jī)制。

將請(qǐng)求隊(duì)列中的請(qǐng)求按照機(jī)器性能進(jìn)行最優(yōu)打包后，根據(jù)等候時(shí)間、batch大小、batch的擴(kuò)展可能性（根據(jù)padding后的句子長(zhǎng)度）等挑選優(yōu)先級(jí)最高的batch處理。

這樣一來(lái)，可以在最大化GPU使用率的同時(shí)規(guī)避饑餓問(wèn)題，減小平均請(qǐng)求時(shí)延。

性能測(cè)試

并行推理超線性擴(kuò)展

硬件環(huán)境：8 * A100 GPU 80GB。

由于單設(shè)備顯存無(wú)法滿足GPT-3推理需求，此處為GPT-3 12層的測(cè)試結(jié)果，設(shè)置句長(zhǎng)為Padding的1/2。

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

運(yùn)行時(shí)推理性能提升50%

硬件環(huán)境：8 * A100 GPU 80GB。

設(shè)置句長(zhǎng)為Padding的1/2。GPT-3-24-Layers for TP=2, GPT-3-48-Layers for TP=4。

以高度優(yōu)化的英偉達(dá)FasterTransformer GPT-3作為對(duì)比方案。

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

此外，對(duì)于NLP推理輸入句長(zhǎng)不同的特點(diǎn)，其分布式推理無(wú)冗余計(jì)算消除功能。

對(duì)于GPT-3模型，Energon-AI的運(yùn)行時(shí)系統(tǒng)在Batch Size為1時(shí)性能略低于FasterTransformer，而在Batch Size較大時(shí)能夠?qū)崿F(xiàn)超過(guò)50%的性能提升。

Dynamic Batching吞吐量增加30%

硬件環(huán)境：8 * A100 GPU 80GB。

測(cè)試使用的模型為GPT-3, 測(cè)試句長(zhǎng)為256以?xún)?nèi)隨機(jī)生成，padding策略為batch內(nèi)最長(zhǎng)padding。

模擬真實(shí)場(chǎng)景下多用戶(hù)同時(shí)發(fā)送大量變長(zhǎng)推理請(qǐng)求的情況，將Energon-AI的動(dòng)態(tài)batch規(guī)劃方法與傳統(tǒng)的FIFO(先入先出)隊(duì)列打包方法進(jìn)行了吞吐量對(duì)比。

由于dynamic batching的算法緩解了直接padding造成的大量冗余計(jì)算問(wèn)題，在該策略下dynamic batching的吞吐量實(shí)現(xiàn)了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

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

目前，已經(jīng)提供了最常見(jiàn)的GPT、BERT和ViT模型作為示例，更詳盡的教程將會(huì)在近期完善。

在構(gòu)建新的并行模型時(shí)，Energon-AI使用Python，且使用方式與Pytorch相似，有層的概念且初始化與執(zhí)行邏輯清晰，用戶(hù)無(wú)需考慮內(nèi)存管理，并行通信等行為。

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

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

與之相對(duì)，在構(gòu)建新的并行模型時(shí)，F(xiàn)asterTransformer需要使用C++代碼并且需要用戶(hù)自行進(jìn)行內(nèi)存管理，定義通信等底層行為組織。

受篇幅限制，如下代碼展示兩層Linear層模型并行運(yùn)行的內(nèi)存管理，具體執(zhí)行，通信的部分代碼。

除此之外，用戶(hù)想要代碼正確執(zhí)行，還需要花費(fèi)大量時(shí)間精力對(duì)內(nèi)存管理、執(zhí)行邏輯、通信行為之間的配合進(jìn)行調(diào)試，C++代碼還需要額外編譯工作。

這些都對(duì)用戶(hù)的并行知識(shí)與編程能力提出了嚴(yán)峻挑戰(zhàn)。

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

更多特性

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

構(gòu)建AI大模型生態(tài)系統(tǒng)

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

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

對(duì)于蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)應(yīng)用AlphaFold，基于Colossal-AI的加速方案的FastFold，成功超越谷歌和哥倫比亞大學(xué)的方案，將AlphaFold訓(xùn)練時(shí)間從11天減少到67小時(shí)。

且總成本更低，在長(zhǎng)序列推理中也實(shí)現(xiàn)9.3～11.6倍的速度提升。

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

相比現(xiàn)有主流方案，可提升參數(shù)容量十余倍，大幅度降低了AI大模型微調(diào)和推理等下游任務(wù)和應(yīng)用部署的門(mén)檻。

Colossal-AI注重開(kāi)源社區(qū)建設(shè)，提供中文教程，開(kāi)放用戶(hù)社群及論壇，對(duì)于用戶(hù)反饋進(jìn)行高效交流與迭代更新，不斷添加等前沿特性。

自開(kāi)源以來(lái)，Colossal-AI已經(jīng)多次登上GitHub熱榜Python方向世界第一，與眾多已有數(shù)萬(wàn)star的明星開(kāi)源項(xiàng)目一起受到海內(nèi)外關(guān)注！

在反映機(jī)器學(xué)習(xí)領(lǐng)域熱點(diǎn)的Papers With Code網(wǎng)站上，Colossal-AI也廣受關(guān)注，登上熱榜第一。

傳送門(mén)

GitHub地址：

https://github.com/hpcaitech/ColossalAI

參考鏈接：https://medium.com/@hpcaitech/6139c5bc7790