推理加速性能超越英伟达FasterTransformer 50%，开源方案打通大模型落地关键路径

性能测试

张量并行可扩展性测试结果展示。硬件环境：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);
}

更多特性

本次发布的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也广受关注，荣登热榜第一。

        

项目团队

潞晨技术团队的核心成员均来自美国加州大学伯克利分校，斯坦福大学，清华大学，北京大学，新加坡国立大学，新加坡南洋理工大学等国内外知名高校；拥有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名厂商工作经历。公司成立即获得创新工场、真格基金等多家顶尖VC机构种子轮投资。

       

△潞晨科技创始人尤洋教授：加州大学伯克利分校博士、IPDPS/ICPP最佳论文、ACM/IEEE George Michael HPC Fellowship、福布斯30岁以下精英(亚洲 2021)、IEEE-CS超算杰出新人奖、UC伯克利EECS Lotfi A. Zadeh优秀毕业生奖      

△潞晨CSO Prof. James Demmel：加州大学伯克利分校杰出教授、ACM/IEEE Fellow，美国科学院、工程院、艺术与科学院三院院士

传送门

项目地址： https://github.com/hpcaitech/ColossalAI

 

 

 

参考链接：https://medium.com/@hpcaitech/6139c5bc7790