作者：王林、飒洋

导读

最近阿里云机器学习 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(Easy Parallel Library)是阿里最近开源的，统一多种并行策略的、灵活易用的自研分布式深度学习训练框架。

项目背景

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

• 2012 年以前，模型计算耗时每 2 年增长一倍，和摩尔定律保持一致；

• 2012 年后，模型计算耗时每 3.4 个月翻一倍，远超硬件发展速度；

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

• 如何简洁易用：

• 接入门槛高：用户实现模型分布式版本难度大、成本高，需要有领域专家经验才能实现高效的分布式并行策略；

• 最优策略难：随着研究人员设计出越来越灵活的模型，以及越来越多的并行加速方法，如果没有自动并行策略探索支持，用户很难找到最适合自身的并行策略；

• 迁移代价大：不同模型适合不同的混合并行策略，但切换并行策略时可能需要切换不同的框架，迁移成本高；

• 如何提高性价比：

• 业界训练万亿规模模型需要的资源：英伟达 3072 A100、谷歌 2048 TPU v3，资源成本非常高；

• 如何降本增效，组合使用各种技术和方法来减少需要的资源，提高训练的速度；

为了应对当前分布式训练的挑战，我们研发了分布式训练框架 EPL，将不同并行化策略进行统一抽象、封装，在一套分布式训练框架中支持多种并行策略。EPL 同时提供简洁易用的接口，用户只需添加几行 annotation 即可完成并行策略的配置，模型代码不需要改动。在用户无感的情况下，EPL 通过进行各种显存、计算、通信优化，打造高效的分布式训练框架。

主要特性

• 多种并行策略统一：在一套分布式训练框架中支持多种并行策略（数据/流水/算子/专家并行）和其各种组合、嵌套使用；

• 接口灵活易用：用户只需添加几行代码就可以使用 EPL 丰富的分布式并行策略，模型代码无需修改；

• 自动并行策略探索：算子拆分时自动探索拆分策略，流水并行时自动探索模型切分策略；

• 分布式性能更优：提供了多维度的显存优化、计算优化，同时结合模型结构和网络拓扑进行调度和通信优化，提供高效的分布式训练；

开源地址

EPL(Easy Parallel Library)的开源地址是：

https://github.com/alibaba/EasyParallelLibrary

我们同时提供了 model zoo，欢迎大家试用：

https://github.com/alibaba/FastNN

EPL 主要技术特点

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

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

• 易用性：用户的模型编程接口和训练接口均基于 TensorFlow，用户只需在已有的单机单卡模型上做简单的标记即可实现不同的分布式策略。EPL 设计了两种简单的策略接口(replicate/split)来表达分布式策略及混合并行。分布式策略标记的方式让用户无需学习新的模型编程接口，仅需几行代码即可实现和转换分布式策略，极大降低了分布式框架的使用门槛。

• 显存优化：EPL 提供了多维度的显存优化技术，包含自动重算技术(Gradient Checkpoint)，ZeRO 数据并行显存优化技术，CPU Offload 技术等，帮助用户用更少的资源训练更大的模型。

• 通信优化技术：EPL 深度优化了分布式通信库，包括硬件拓扑感知，通信线程池，梯度分组融合，混合精度通信、梯度压缩等技术。

技术架构

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

• 接口层：用户的模型编程接口基于 TensorFlow，同时 EPL 提供了易用的并行化策略表达接口，让用户可以组合使用各种混合并行策略；

• 中间表达层：将用户模型和并行策略转成化内部表达，通过 TaskGraph、VirtualDevices 和策略抽象来表达各种并行策略；

• 并行化引擎层：基于中间表达，EPL 会对计算图做策略探索，进行显存/计算/通信优化，并自动生成分布式计算图。

• Runtime 执行引擎：将分布式执行图转成 TFGraph，再调用 TF 的 Runtime 来执行；

并行化策略表达

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

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

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

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

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

显存优化

当模型增长，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 支持各种显存优化技术的组合使用，达到显存的极致优化。我们在 T5 模型上开启 GC+ZeRO+显存优化的 AMP 技术，在性能保持不变的情况下，显存降低 2.6 倍。

应用场景

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

万亿/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 万亿模型。

流水并行加速 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 的流水并行+数据并行，相较于数据并行得到了 66%的提升。

Roadmap

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

• EPL 发源于阿里云内部的业务需求，在服务内部业务的过程中，我们积累了大量的经验和理解，很好的支持了大规模、多样性的业务场景，在我们自身随着业务需求的迭代逐渐完善的同时，也希望能够开源给社区，将自身积累的经验和理解也同样回馈给社区，希望和深度学习训练框架的开发者或深度学习从业者之间有更多更好的交流和共建，为这个行业贡献我们的技术力量。

• 我们希望能够借助开源的工作，收到更多真实业务场景下的用户反馈，以帮助我们持续完善和迭代，并为后续的工作投入方向提供输入。

• 同时我们希望借助开源的工作，能吸引一些志同道合的同学、公司或组织来参与共建，持续完善深度学习生态。

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

• 持续的性能优化和稳定性改进；

• 通用算子拆分功能；

• 自动拆分策略探索的基础版；

• 自动流水并行策略探索；

此外，在中长期，我们在下面几个探索性的方向上会持续投入精力，也欢迎各种维度的反馈和改进建议以及技术讨论，同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建。

• 全自动的模型并行策略探索；

• 高效的策略探索算法和精准的 CostModel 评估；

• eager model 下的并行策略探索；

• 更多新硬件的支持、适配和协同优化；

• 高效的算子优化和集成、极致的显存优化、软硬一体的通信优化；

欢迎加入 EPL 用户交流群：

参考文献

1. Whale: Scaling Deep Learning Model Training to the Trillions https://arxiv.org/abs/2011.09208
   

2. https://openai.com/blog/ai-and-compute/

3. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. https://arxiv.org/abs/1910.02054
   

4. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding https://arxiv.org/abs/1810.04805