阿里开源支持10万亿模型的自研分布式训练框架EPL

一 导读

最近阿里云机器学习PAI平台和达摩院智能计算实验室一起发布“低碳版”巨模型M6-10T，模型参数已经从万亿跃迁到10万亿，规模远超业界此前发布的万亿级模型，成为当前全球最大的AI预训练模型。同时，做到了业内极致的低碳高效，使用512 GPU在10天内即训练出具有可用水平的10万亿模型。相比之前发布的大模型GPT-3，M6实现同等参数规模，能耗仅为其1%。

M6模型训练使用的正是阿里云机器学习PAI平台自研的分布式训练框架EPL(Easy Parallel Library，原名whale)。EPL通过对不同并行化策略进行统一抽象、封装，在一套分布式训练框架中支持多种并行策略，并进行显存、计算、通信等全方位优化来提供易用、高效的分布式训练框架。

EPL背后的技术框架是如何设计的？开发者可以怎么使用EPL？EPL未来有哪些规划？今天一起来深入了解。

二 EPL是什么

EPL(Easy Parallel Library)是阿里巴巴最近开源的，统一了多种并行策略、灵活易用的自研分布式深度学习训练框架。

1 项目背景

近些年随着深度学习的火爆，模型的参数规模也飞速增长，OpenAI数据显示：

• 2012年以前，模型计算耗时每2年增长一倍，和摩尔定律保持一致；
• 2012年后，模型计算耗时每3.4个月翻一倍，远超硬件发展速度；

近一年来，百亿、千亿级的参数模型陆续面世，谷歌、英伟达、阿里、智源研究院更是发布了万亿参数模型。随着模型参数规模的增大，模型效果逐步提高，但同时也为训练框架带来更大的挑战。当前已经有一些分布式训练框架Horovod、Tensorflow Estimator、PyTorch DDP等支持数据并行，Gpipe、PipeDream、PipeMare等支持流水并行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore等支持算子拆分，但当训练一个超大规模的模型时还是会面临一些挑战：

• 如何简洁易用：
  • 接入门槛高：用户实现模型分布式版本难度大、成本高，需要有领域专家经验才能实现高效的分布式并行策略；
• 最优策略难：随着研究人员设计出越来越灵活的模型以及越来越多的并行加速方法，如果没有自动并行策略探索支持，用户很难找到最适合自身的并行策略；
• 迁移代价大：不同模型适合不同的混合并行策略，但切换并行策略时可能需要切换不同的框架，迁移成本高；
• 如何提高性价比：
• 业界训练万亿规模模型需要的资源：英伟达 3072 A100、谷歌 2048 TPU v3，资源成本非常高；
• 如何降本增效，组合使用各种技术和方法来减少需要的资源，提高训练的速度；

为了应对当前分布式训练的挑战，阿里云机器学习PAI团队自主研发了分布式训练框架EPL，将不同并行化策略进行统一抽象、封装，在一套分布式训练框架中支持多种并行策略。同时，EPL提供简洁易用的接口，用户只需添加几行annotation（注释）即可完成并行策略的配置，不需要改动模型代码。EPL也可以在用户无感的情况下，通过进行显存、计算、通信等全方位优化，打造高效的分布式训练框架。

2 主要特性

• 多种并行策略统一：在一套分布式训练框架中支持多种并行策略（数据/流水/算子/专家并行）和其各种组合嵌套使用；
• 接口灵活易用：用户只需添加几行代码就可以使用EPL丰富的分布式并行策略，模型代码无需修改；
• 自动并行策略探索：算子拆分时自动探索拆分策略，流水并行时自动探索模型切分策略；
• 分布式性能更优：提供了多维度的显存优化、计算优化，同时结合模型结构和网络拓扑进行调度和通信优化，提供高效的分布式训练。

3 开源地址见文末

三 EPL主要技术特点

• EPL通过丰富并行化策略、简单易用的接口、多维度的显存优化技术和优化的计算通信加速技术，让每一位算法工程师都能轻松训练分布式大模型任务。

丰富的并行化策略：EPL提供了多种并行化策略及其组合策略，包含数据并行、流水并行、算子拆分并行及并行策略的组合嵌套。丰富的策略选择使得不同的模型结构都能找到最适合自己的分布式训练方式。

• 易用性：用户的模型编程接口和训练接口均基于TensorFlow，用户只需在已有的单机单卡模型上做简单的标记，即可实现不同的分布式策略。EPL设计了两种简单的策略接口(replicate/split)来表达分布式策略及混合并行。分布式策略标记的方式让用户无需学习新的模型编程接口，仅需几行代码即可实现和转换分布式策略，极大降低了分布式框架的使用门槛。
• 显存优化：EPL提供了多维度的显存优化技术，包含自动重算技术(Gradient Checkpoint)，ZeRO数据并行显存优化技术，CPU Offload技术等，帮助用户用更少的资源训练更大的模型。
• 通信优化技术：EPL深度优化了分布式通信库，包括硬件拓扑感知、通信线程池、梯度分组融合、混合精度通信、梯度压缩等技术。

1 技术架构

EPL框架如下图所示，主要分为以下几个模块：

• 接口层：用户的模型编程接口基于TensorFlow，同时EPL提供了易用的并行化策略表达接口，让用户可以组合使用各种混合并行策略；
• 中间表达层：将用户模型和并行策略转化成内部表达，通过TaskGraph、VirtualDevices和策略抽象来表达各种并行策略；
• 并行化引擎层：基于中间表达，EPL会对计算图做策略探索，进行显存/计算/通信优化，并自动生成分布式计算图；
• Runtime执行引擎：将分布式执行图转成TFGraph，再调用TF 的Runtime来执行；

2 并行化策略表达

EPL通过strategy annotation的方式将模型划分为多个TaskGraph，并在此基础上进行并行化。EPL有两类strategy：replicate 和 split。通过这两种并行化接口，可以表达出各种不同的并行化策略，例如：

1、数据并行: 下面这个例子是一个数据并行的例子，每个模型副本用一张卡来计算。如果用户申请了8张卡，就是一个并行度为8的数据并行任务。

2、流水并行：在下面的例子里，模型被切分成2个 TaskGraph, "stage0"和"stage1"，用户可以通过配置pipeline.num_micro_batch参数来设定pipeline的micro batch数量。在这个例子里，"stage_0"和"stage_1"组成一个模型副本，共需要2张GPU卡。如果用户申请了8张卡，EPL会自动在pipeline外嵌套一层并行度为4的数据并行（4个pipeline副本并行执行）。

3、算子拆分并行：在以下例子中，EPL会对split scope下的模型定义做拆分，并放置在不同的GPU卡上做并行计算。

4、同时，EPL支持对上述并行策略进行组合和嵌套，组成各种混合并行策略，更多示例可以参考开源代码的文档和示例。

3 显存优化

当模型增长，GPU的显存常常成为训练大模型的瓶颈。EPL提供了多维度的显存优化技术，极大优化了训练显存消化。

• 重算 Recomputation (Gradient Checkpoint)：正常的DNN前向过程中会生成activation，这部分activation会在后向过程中用于梯度计算。因此，在梯度生成之前，前向的activation会一直存留在显存中。activation大小和模型结构以及batch size相关，通常占比都非常高。Gradient Checkpoint (GC) 通过保留前向传播过程中的部分activation，在反向传播中重算被释放的activation，用时间换空间。GC中比较重要的一部分是如何选择合适的checkpoint点，在节省显存、保证性能的同时，又不影响收敛性。EPL提供了自动GC功能，用户可以一键开启GC优化功能。
• ZeRO：在数据并行的场景下，每个卡上会存放一个模型副本，optimizer state等，这些信息在每张卡上都是一样，存在很大的冗余量。当模型变大，很容易超出单卡的显存限制。在分布式场景下，可以通过类似DeepSpeed ZeRO的思路，将optimizer state和gradient分片存在不同的卡上，从而减少单卡的persistent memory占用。
• 显存优化的AMP(Auto Mixed Precision)：在常规的AMP里，需要维护一个FP16的weight buffer，对于参数量比较大的模型，也是不小的开销。EPL提供了一个显存优化的AMP版本，FP16只有在用的时候才cast，从而节约显存。
• Offload: Offload将训练的存储空间从显存扩展到内存甚至磁盘，可以用有限的资源训练大模型。

同时，EPL支持各种显存优化技术的组合使用，达到显存的极致优化。阿里云机器学习PAI团队在T5模型上开启了GC+ZeRO+显存优化的AMP技术，在性能保持不变的情况下，显存降低2.6倍。

四 应用场景

EPL适合不同场景的模型，在阿里巴巴内部已经支持图像、推荐、语音、视频、自然语言、多模态等业务场景。同时，EPL也支持不同规模的模型，最大完成了10万亿规模的M6模型训练，下面以M6和Bert模型为例进行介绍。

1 万亿/10万亿 M6模型预训练

训练一个万亿/10万亿参数模型，算力需求非常大。为了降低算力需求，EPL中实现了MoE(Mixture-of-Experts)结构，MoE的主要特点是稀疏激活，使用Gating(Router)来为输入选择Top-k的expert进行计算（k常用取值1、2），从而大大减少算力需求。

EPL支持专家并行(Expert Parallelism, EP)，将experts拆分到多个devices上，降低单个device的显存和算力需求。同时，数据并行有利于提升训练的并发度，因此，采用数据并行+专家并行的混合并行策略来训练M6模型：MoE layer采用专家并行，其他layer采用数据并行。

EPL中提供简洁易用的接口来进行模型的混合并行训练，只需要增加几行annotation（注释）来配置并行策略，不需要对模型本身做任何修改。例如，M6模型采用数据并行+专家并行的策略，只需要增加如下图的annotation：

同时，为了节约训练资源、提高训练效率，我们采用了EPL的显存优化技术和计算通信加速技术，包含自动 Gradient Checkpointing节省activation显存占用，CPU Offload技术用于优化optimizer states和weight的显存占用，采用DP+EP混合并行策略降低算力需求，结合混合精度、编译优化等技术提高训练效率等。

借助EPL框架，首次在480 V100 上，3天内完成万亿M6模型的预训练。相比此前业界训练同等规模的模型，此次仅使用480张V100 32G GPU就成功训练出万亿模型M6，节省算力资源超80%，且训练效率提升近11倍。进一步使用512 GPU在10天内即训练出具有可用水平的10万亿模型。

2 流水并行加速Bert Large模型训练

对于Bert Large模型，结构图如下图所示：

由于Bert Large模型对显存消耗较大，Nvidia V100 16G显卡上batch size常常只有2-8左右（具体值和Embedding大小、Sequence Length等有关）。Batch size太小会导致算法收敛波动大、收敛效果差的问题。同时，通过数据并行模式训练，通信占比较高，分布式加速效果不理想。

分析Bert Large模型，由24层重复结构的encoder组成，可以使用流水并行进行加速。这里，我们将Bert Large中的Encoder Layer 1~8层、Encoder Layer 9~16层，Encoder Layer 17~24层分别放在不同的卡上进行训练，并行化后的计算图如下图所示：

如此，每张卡训练时的显存开销会减少，从而可以增大batch size以提升收敛加速。另外，针对因模型过大、单卡显存无法放下所导致的无法训练的场景，通过Layer间拆分的模型并行方式来进行分布式训练。通过epl.replicate接口可以实现模型的stage划分，同时通过流水并行的执行调度来提升并行化性能，如下图所示：

上述例子是一个流水micro batch mumber为5的情况。通过流水并行优化后的时间轴可以看出，在同一个时间上，多张卡可以并行计算。当5个micro batch结束后，每张卡会将梯度进行本地的累计之后再进行update。与单纯的模型并行相比，通过流水的交替执行，提高了GPU的利用率。EPL还通过采用Backward-Preferred调度优化策略来提升流水并行性能，降低GPU空闲时间和显存开销。

为能够获得更高的水平扩展，EPL还支持在流水并行外嵌套数据并行来提升训练吞吐。EPL会自动推导嵌套的数据并行的并行度。最新测试结果显示，在32卡GPU规模下，使用EPL的流水+数据并行对Bert Large模型进行优化，相比于数据并行，训练速度提升了66%。

五 Roadmap

我们决定建设开源生态主要有如下的考虑：

• EPL发源于阿里云内部的业务需求，很好地支持了大规模、多样性的业务场景，在服务内部业务的过程中也积累了大量的经验，在EPL自身随着业务需求的迭代逐渐完善的同时，我们也希望能够开源给社区，将自身积累的经验和理解回馈给社区，希望和深度学习训练框架的开发者或深度学习从业者之间有更多更好的交流和共建，为这个行业贡献我们的技术力量。
• 我们希望能够借助开源的工作，收到更多真实业务场景下的用户反馈，以帮助我们持续完善和迭代，并为后续的工作投入方向提供输入。
• 同时，我们希望借助开源的工作，能吸引一些志同道合的同学、公司或组织来参与共建，持续完善深度学习生态。

后续，我们计划以两个月为单位发布Release版本。EPL近期的Roadmap如下：

• 持续的性能优化和稳定性改进；
• 通用算子拆分功能；
• 自动拆分策略探索的基础版；
• 自动流水并行策略探索；

此外，在中长期，我们将在软硬件一体优化、全自动策略探索等几个探索性的方向上持续投入精力，也欢迎各种维度的反馈和改进建议以及技术讨论，同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建。

• 全自动的模型并行策略探索；
• 高效的策略探索算法和精准的CostModel评估；
• eager model下的并行策略探索；
• 更多新硬件的支持、适配和协同优化；
• 高效的算子优化和集成、极致的显存优化、软硬一体的通信优化；

EPL(Easy Parallel Library)的开源地址：https://github.com/alibaba/EasyParallelLibrary
我们同时提供了model zoo，欢迎大家试用：https://github.com/alibaba/FastNN

作者 | 王林、飒洋

原文链接

本文为阿里云原创内容，未经允许不得转载。