一、简介近日，阿里云机器学习PAI平台与达摩院智能计算实验室联合发布“低碳版本”巨型模型M6-10T，模型参数从万亿级跃升至10万亿级，远超业界此前发布的万亿级模型，成为全球最大的AI预训练模型。同时实现了业界极致的低碳高效，使用512个GPU在10天内训练出10万亿个可用级别的模型。与此前发布的大型号GPT-3相比，M6实现了同等参数规模，能耗仅为1%。M6模型训练采用了阿里云机器学习PAI平台开发的分布式训练框架EPL（EasyParallelLibrary，原whale）。EPL通过统一抽象和封装不同的并行化策略，在一个分布式训练框架中支持多种并行策略，对内存、计算、通信进行全方位优化，提供了一个易用高效的分布式训练框架。EPL背后的技术框架是如何设计的？开发人员如何使用EPL？EPL未来有什么计划？今天让我们详细了解一下。2、什么是EPLEPL（EasyParallelLibrary）是阿里巴巴近期自主研发的统一多种并行策略、灵活易用的分布式深度学习训练框架。一、项目背景随着近几年深度学习的普及，模型的参数规模也快速增长。OpenAI数据显示，2012年之前，模型计算时间每2年翻一番，符合摩尔定律；2012年后，花在模型计算上的时间每3.4个月翻一番，远超硬件发展速度；十亿参数模型。随着模型参数规模的增大，模型的效果逐渐提升，但也给训练框架带来了更大的挑战。目前Horovod、TensorflowEstimator、PyTorchDDP等分布式训练框架支持数据并行；Gpipe、PipeDream、PipeMare等支持流水线并行；MeshTensorflow、FlexFlow、OneFlow、MindSpore等支持算子拆分。但是，大规模模型仍然存在一些挑战：如何简单易用：接入门槛高：用户实现模型的分布式版本难度大，成本高，需要领域专家经验实现高效的分布式并行策略；最优策略难：随着研究者设计的模型越来越灵活，并行加速方法越来越多，如果没有自动并行策略探索的支持，用户很难找到最适合自己的并行策略；迁移成本高：不同的模型适用于不同的，但在切换并行策略时，可能需要在不同的框架之间切换，迁移成本高；如何提升性价比：行业需要训练万亿级模型的资源：Nvidia3072A100、Google2048TPUv3，资源成本非常高；如何降低成本和提高效率，结合使用各种技术和方法来减少所需资源并提高训练速度；为了应对当前分布式训练的挑战，阿里云机器学习PAI团队自主研发了分布式训练框架EPL，将不同的并行泛化策略统一抽象和封装，在一个分布式训练框架中支持多种并行策略.同时，EPL提供了简单易用的界面。用户只需添加几行注解即可完成并行策略的配置，无需更改模型代码。EPL还可以在用户无感知的情况下，通过对显存、计算、通信的全方位优化，打造高效的分布式训练框架。2.主要特点多种并行策略的统一：支持多种并行策略（数据/管道/算子/专家并行）和分布式训练框架中各种组合的嵌套使用；灵活易用的接口：用户只需添加几行代码即可使用EPL丰富的分布式并行策略，无需修改模型代码；自动并行策略探索：operator拆分时自动探索拆分策略，pipeline并行时自动探索模型拆分策略；分布式性能更好：提供多维内存优化和计算优化，结合模型结构和网络拓扑进行调度和通信优化，提供高效的分布式训练。3、开源地址见文末3、EPL的主要技术特点EPL让每一位算法工程师都能轻松训练分布式大模型任务。丰富的并行化策略：EPL提供了多种并行化策略及其组合策略，包括数据并行、流水线并行、算子拆分并行以及并行策略的组合嵌套。丰富的策略选项使得不同的模型结构能够找到最适合自己的分布式训练方式。易用性：用户的模型编程界面和训练界面均基于TensorFlow。用户只需在现有的单机单卡模型上做简单的标记，即可实现不同的分布式策略。EPL设计了两个简单的策略接口（replicate/split）来表达分布式策略和混合并行。分布式策略标记的方式使得用户无需学习新的模型编程接口，只需几行代码就可以实现和转换分布式策略，大大降低了分布式框架的使用门槛。内存优化：EPL提供多维内存优化技术，包括自动重计算技术（GradientCheckpoint）、ZeRO数据并行内存优化技术、CPUOffload技术等，帮助用户用更少的资源训练更大的模型。通信优化技术：EPL对分布式通信库进行深度优化，包括硬件拓扑感知、通信线程池、梯度包融合、混合精度通信、梯度压缩等技术。1.技术架构EPL框架如下图所示，主要分为以下模块：接口层：用户模型编程接口基于TensorFlow，EPL提供简单易用的并行策略表达接口，允许用户组合各种混合并行策略；中间表达层：将用户模型和并行策略转化为内部表达，通过TaskGraph、VirtualDevices和策略抽象来表达各种并行策略；并行引擎层：EPL基于中间表达式对计算图进行策略探索，进行内存/计算/通信优化，自动生成分布式计算图；Runtime执行引擎：将分布式执行图转化为TFGraph，然后调用TFRuntime执行；2.Parallelizationstrategy通过strategyannotation表示EPL将模型划分为多个TaskGraphs，并在此基础上进行并行化。EPL有两种类型的策略：复制和拆分。通过这两个并行化接口，可以表达出多种并行化策略，例如：(1)数据并行化：下面的例子是数据并行化的一个例子，每个模型副本由一张卡计算。如果用户申请8张卡，则为并行度为8的数据并行任务。(2)Pipelineparallelism：在下面的例子中，模型被分成两个TaskGraphs，“stage0”和“stage1”。用户可以通过配置pipeline.num_micro_batch参数来设置pipeline的微批数量。在这个例子中，“stage_0”和“stage_1”形成模型的副本，总共需要2张GPU卡。如果用户申请8张卡，EPL会自动在管道外嵌套一层并行度为4的数据并行（4个管道副本并行执行）。(3)Operatorsplitparallelism：在下面的例子中，EPL将splitscope下的modeldefinition进行拆分，放到不同的GPU卡上进行并行计算。(4)同时，EPL支持上述并行策略的组合和嵌套，形成多种混合并行策略。更多示例请参考开源代码的文档和示例。3.内存优化当模型增长时，GPU内存往往成为训练大型模型的瓶颈。EPL提供多维显存优化技术，极大优化训练显存消化。RecalculationRecomputation(GradientCheckpoint)：Activation在DNN正常的forward过程中产生，这部分activation将用于backward过程中的梯度计算。因此，前向激活将保留在显存中，直到生成梯度。activationsize与模型结构和batchsize有关，通常占比很高。GradientCheckpoint（GC）通过保留前向传播过程中的部分激活，重新计算反向传播中释放的激活，以时间换取空间。GC的一个比较重要的部分就是如何选择一个合适的checkpoint点，既节省显存又保证性能，同时又不影响收敛。EPL提供了自动GC功能，用户可以一键开启GC优化功能。ZeRO：在数据并行的场景下，每张卡都会存储一份模型，优化器状态等，这些信息在每张卡上都是一样的，冗余度很大。当模型变大时，很容易超出单卡的内存限制。在分布式场景下，优化器状态和梯度分片可以通过类似DeepSpeedZeRO的思路存储在不同的卡上，从而减少单卡的持久内存占用。Memory-optimizedAMP（AutoMixedPrecision）：在传统的AMP中，需要维护一个FP16的weightbuffer，这对于参数量比较大的模型来说是一个不小的开销。EPL提供显存优化的AMP版本，FP16只在使用时才投，从而节省显存。Offload：Offload将训练存储空间从显存扩展到内存甚至磁盘，可以在资源有限的情况下训练大型模型。同时，EPL支持多种显存优化技术的组合使用，实现显存的极致优化。阿里云机器学习PAI团队在T5模型上开启了GC+ZeRO+内存优化的AMP技术。在性能不变的情况下，内存减少了2.6倍。4.应用场景EPL是一个适用于不同场景的模型。在阿里巴巴内部，已经支持图像、推荐、语音、视频、自然语言、多模态等业务场景。同时EPL还支持不同规模的模型，已经完成了最大规模10万亿的M6模型的训练。下面将以M6和Bert模型为例进行介绍。1.万亿/10万亿M6模型预训练训练一个万亿/10万亿参数模型需要非常大的计算能力。为了降低对计算能力的需求，EPL中实现了MoE（Mixture-of-Experts）结构。MoE的主要特征是稀疏激活。Gating（Router）用于选择Top-k专家作为输入进行计算（k通常取1、2的值），从而大大降低了对算力的需求。EPL支持专家并行（ExpertParallelism，EP），将专家拆分到多个设备中，降低了单个设备的内存和计算能力需求。同时，数据并行有利于提高训练的并发性。因此采用数据并行+专家并行的混合并行策略来训练M6模型：MoE层采用专家并行，其他层采用数据并行。EPL为模型的混合并行训练提供了一个简单易用的界面。它只需要添加几行注释（注释）来配置并行策略，而无需对模型本身进行任何修改。例如M6模型采用了数据并行+专家并行的策略，只需要添加如下图所示的标注即可：同时，为了节省训练资源，提高训练效率，我们使用EPL的显存优化技术和计算通信加速技术，包括automaticGradientCheckpointing节省激活显存占用，CPUOffload技术用于优化optimizerstates和weight的显存占用，采用DP+EP混合并行策略降低算力需求，结合混合精度、编译优化等技术提升训练效率等。借助EPL框架，首次在480V100上3天完成万亿M6模型的预训练。与业界以往训练的同尺寸模型相比，此次仅使用480颗V10032GGPU就成功训练出TrillionModelM6，节省超过80%的计算资源，训练效率提升近11倍。进一步使用512GPU在10天内训练出一个10万亿级的模型，达到可用水平。2.BertLarge模型训练的流水线并行加速对于BertLarge模型，结构图如下图所示：由于BertLarge模型显存消耗较大，在NvidiaV10016G显卡上batchsize通常为只有2-8左右（具体值和Embeddingsize，SequenceLength等）。batchsize过小会导致算法收敛波动较大，收敛效果差。同时，通过数据并行模式训练，通信占比高，分布式加速效果不理想。分析BertLarge模型，该模型由24层具有重复结构的编码器组成，可以使用流水线并行进行加速。这里，我们将BertLarge中的EncoderLayer1~8、EncoderLayer9~16、EncoderLayer17~24分别放在不同的卡片上进行训练。并行化后的计算图如下图所示：，每张卡片训练的内存开销会减少，从而可以增加batchsize来提高收敛速度。另外，针对模型过大无法训练模型，单卡内存放不下的场景，通过并行拆分层间模型进行分布式训练。通过epl.replicate接口可以实现模型的stage划分，通过pipeline并行执行调度可以提高并行化性能，如下图所示：上例是pipelinemicrobatchmember为5的情况.从流水线并行优化后的时间轴可以看出，多卡可以同时进行并行计算。当5个microbatches结束后，每张卡片都会在本地累加梯度，然后更新。与纯模型并行相比，通过流水线的交替执行提高了GPU的利用率。EPL还使用Backward-Preferred调度优化策略来提高流水线并行性能并减少GPU空闲时间和显存开销。为了实现更高的横向扩展，EPL还支持流水线并行之外的嵌套数据并行，以提高训练吞吐量。EPL自动导出嵌套数据并行度的并行度。最新测试结果显示，在32个GPU的规模下，使用EPL的流水线+数据并行优化BertLarge模型，相比数据并行，训练速度提升了66%。5.Roadmap我们决定构建开源生态主要出于以下考虑：EPL源于阿里云内部业务需求，很好地支撑了大规模、多样化的业务场景，在这个过程中积累了大量的资源。服务于内部业务。经验，在EPL本身随着业务需求的迭代而逐渐完善的同时，我们也希望能够开源给社区，将我们积累的经验和理解回馈给社区。我们之间有更多更好的交流与共建，为这个行业贡献我们的技术力量。我们希望通过开源工作，在真实的业务场景中得到更多的用户反馈，帮助我们不断完善和迭代，为后续工作提供输入。同时，我们希望借助开源工作，吸引一些志同道合的同学、企业或组织参与共建，不断完善深度学习生态。未来我们计划以两个月为单位发布Release版本。EPL近期路线图如下：持续性能优化和稳定性提升；通用算子分裂函数；自动拆分策略探索基础版；自动流水线并行策略探索；此外，在中长期，我们将进行软硬件优化、全自动策略探索等探索方向，也欢迎来自各个维度的反馈、改进建议和技术讨论。同时，我们欢迎并期待有志于开源社区建设的同仁参与共建。全自动模型并行策略探索；高效的策略探索算法和准确的CostModel评估；eager模型下的并行策略探索；更多新硬件支持、适配和协同优化；高效算子优化集成、极致显存优化、软硬件一体化通信优化；EPL（EasyParallelLibrary）开源地址：https://github.com/alibaba/EasyParallelLibrary我们也提供了模型动物园，欢迎大家试用：https://github.com/alibaba/FastNN
