本文经AI新媒体量子位（公众号ID：QbitAI）授权转载，转载请联系出处。随着前沿AI大规模模型的指数级增长，仅靠单卡难以满足当前的推理需求。以具有1750亿个参数的GPT-3为例。仅加载模型参数就需要数百GB的存储空间，远远超过单个GPU的容量。因此，多卡并行被视为AI大模型推理的必然选择。但是现有的推理系统仍然存在很多不足。比如需要用户手动管理通信、内存等各个部分的协同，需要额外编译……导致用户使用门槛高。为此，大规模并行AI训练系统Colossal-AI团队提出了大型模型推理系统Energon-AI。Energon-AI以“高性能、高可用、可扩展性”为理念，深入单实例多设备推理场景，在性能和易用性上均有优势。用户只需对现有项目进行极小的改动，即可完成自定义大模型的推理部署，获得并行扩展的超线性加速。对于AI大模型的分布式推理加速，相比NVIDIAFasterTransformer可提升50%以上。开源地址：https://github.com/hpcaitech/ColossalAIAI大模型推理部署困难近年来，计算设备（如GPU）的并行计算能力、内存容量和内存速度都得到了极大的提升。然而，在面对呈指数级增长的模型规模时，单个设备向上扩展的性能增益仍然难以满足大型模型对内存和性能的要求。然而，目前的深度学习推理系统主要面向多实例单机和单实例单机的简单推理场景，忽视了AI大模型推理所需的单实例多机的挑战和机遇。Energon-AI系统就是为了解决这个痛点而诞生的。△模型参数的快速增长[https://arxiv.org/abs/2111.14247]Energon-AI系统设计面向AI大模型部署，Colossal-AI团队设计了单实例多设备推理系统Energon-AI。△Energon-AI超大模型推理系统示意图Energon-AI系统设计分为三个层次，即运行时系统（Runtime）、分布式推理实例（Engine）和前端服务系统（Serving）：运行时在在运行时系统设计过程中，团队发现，当模型规模不断增大时，一般矩阵乘法的时间占比逐渐增加。但是，内存密集型算子与KernelLaunch的时间比逐渐降低，推理过程进一步从内存密集型迁移到计算密集型。TensorRT和专用推理系统对内存密集型操作的优化效果大大降低。减少。Energon-AIRuntime依托Colossal-AI实现张量并行，同时针对显存不足的情况设计了流水线并行封装方式。此外，团队还引入了大量推理专用算子和方法。例如，面对NLP中输入的可变长度，引入transpose_padding_rebulid和transpose_padding_remove等算子，高效支持Encoder和Decoder模型中MLP层的冗余计算剔除方法。Engine单机推理中的程序，数据入口和出口是一样的。分布式训练的主要目标是模型参数，因此不需要管理多个进程的输入输出，但多设备推理不同。团队希望通过良好的封装，让Engine具有与单机推理完全相同的行为。采用半集中方式，在主进程中使用RPC调用各个设备上的初始化或推理方法，分布式推理可以集中控制。同时，每个设备维护自己的TensorParallel和PipelineParallel通信逻辑。同时，团队在每个进程中设计并维护了一个分布式消息队列，以保证多进程中多线程调用执行的一致性。Serving：Energon-AI引入动态Batching机制，解决用户请求分散、长度不定与大模型推理依赖GPU并行计算的矛盾。请求队列中的请求根据机器性能进行优化打包后，根据等待时间、batchsize、batch扩展的可能性（根据padding后的句子长度）等选择优先级最高的batch进行处理.这样可以在最大化GPU利用率的同时避免饥饿问题，降低平均请求延迟。△Batch管理流程示意图性能测试并行推理超线性扩展△Tensor并行可扩展性测试结果展示硬件环境：8*A100GPU80GB。由于单个设备的内存无法满足GPT-3的推理要求，这里是GPT-3第12层的测试结果，句子长度设置为Padding的1/2。Energon-AI八卡并行推理在batchsize为32时可以获得比单卡Pytorch直接推理8.5倍的超线性加速，运行时推理性能提升50%△Tensor并行运行时系统推理时延对比硬件环境：8*A100GPU80GB。将句子长度设置为Padding的1/2。TP=2的GPT-3-24-Layers，TP=4的GPT-3-48-Layers。使用高度优化的NVIDIAFasterTransformerGPT-3作为对比程序。FasterTransformer在其4.0版本中引入了分布式推理特性，目前支持GPT-3模型的分布式推理，但由于其纯C++代码的高耦合性，其灵活性和易用性相对较低。另外，针对NLP推理输入句长不同的特点，其分布式推理没有冗余计算剔除功能。对于GPT-3模型，当BatchSize为1时，Energon-AI的runtime系统性能略低于FasterTransformer，当BatchSize较大时可实现50%以上的性能提升。DynamicBatching吞吐量提升30%△Dynamicbatching和直接打包的batch吞吐量对比硬件环境：8*A100GPU80GB。测试使用的模型是GPT-3，测试句子长度在256以内随机生成，padding策略是batch中最长的padding。模拟真实场景中多个用户同时发送大量变长推理请求的情况，将Energon-AI的动态批计划方法的吞吐量与传统的FIFO（先进先出）进行对比)队列封装方法。由于动态批处理算法缓解了直接填充带来的大量冗余计算问题，动态批处理的吞吐量在该策略下实现了34.7%的提升。易于使用Pythonfromgptimportgpt3fromgpt_serverimportlaunch_engine#forenginemodel_class=gpt3model_type="gpt"host="127.0.0.1"port=29400half=Truebackend="nccl"#forparalleltp_init_size=4pp_init_size=2#forserverengine_server1=launch_27host.0.1"server_port=8020rm_padding=TruePythonenergonaiserviceinit--config_file=gpt_config.pyEnergon-AI在追求性能的同时，希望保持系统的灵活性和易用性，用户只需自定义并行模型、并行参数和服务添加请求逻辑添加到配置文件启动推理服务，目前最常见的GPT、BERT、ViT模型已经提供示例，更详细的教程将在近期完成。在构建新的并行模型时，Energon-AIPython的使用方式与Pytorch类似，有层次的概念，初始化和执行逻辑清晰，用户无需考虑内存漫游等行为代理和并行通信。以下代码显示了运行由两个并行线性层组成的模型的完整代码。PythonclassMLP(nn.Module):def__init__(self,dim,dtype,bias):super().__init__()self.dense_0=Linear1D_Col(dim,dim,dtype=dtype,bias=bias,gather_output=False)自我.dense_1=Linear1D_Row(dim,dim,dtype=dtype,bias=bias,parallel_input=True)defforward(self,x):x=self.dense_0(x)x=self.dense_1(x)returnxrelativeto,在构建新的并行模型时，FasterTransformer需要使用C++代码，并需要用户进行内存管理和定义通信等底层行为组织。限于篇幅，以下代码展示了并行运行的两层Linear层模型的内存管理、具体执行、通信部分。此外，用户要想代码正确执行，还需要花费大量的时间和精力去调试内存管理、执行逻辑、通信行为之间的配合。C++代码还需要额外的编译工作。这些都对用户的并行知识和编程能力提出了严峻的挑战。C++//内存分配（仅针对单个参数）。T*d_inter_kernel=NULLparam_.ffn.intermediate_weight.kernel=d_inter_kernel;device_malloc(&d_inter_kernel,dim*dim);//两个MLP层scublasMM_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
