来源：DataFunTalk
导读 本次分享题目为超大模型高效训练开源分布式框架 EPL（ Easy Parallel Library）。
主要介绍：
1. 模型训练的趋势和挑战
2. 分布式框架 EPL 的设计
3. EPL 最佳实践和效果
分享嘉宾｜贾贤艳 阿里云 技术专家
编辑整理｜胡俊琪
出品社区｜DataFun
模型训练的趋势和挑战
1. 模型训练发展趋势
随着深度学习的发展，模型的参数规模迅速增长，从早期模型算力两年增长一倍，到每三四个月翻一倍，这个速度远超硬件的发展速度。同时随着模型规模的增大， Transformer 类模型在多个场景上都被验证能给模型质量带来极大提升。大模型在提升模型效果的同时也为模型训练带来了诸多的挑战。
2. 训练方式的变迁
一方面随着模型参数规模和数据量的增加，单 GPU 的训练速度已经不能满足很多业务场景的需求。比较常见的加速策略是采用数据并行，通过多个模型副本来并行计算不同的输入，并在每个训练迭代进行参数的同步，以此来加速海量数据的训练。但是随着参数规模的进一步扩大，单 GPU 已经没有办法放下一个完整的模型副本。这个时候我们就不能用单纯的数据并行来加速。
一种常见的做法是采用模型并行来对模型进行拆分，把模型的子图放置在不同的 GPU 卡上进行扩展。常见的模型并行又包括流水并行和算子拆分并行。
流水并行也叫 pipeline 并行，它是将模型 layer 间切分成几个 stage 然后放置在不同的 GPU 上。同时我们会将数据的 mini batch 切分成多个 micro batch，然后通过 pipeline 的形式来交替执行 micro batch 的前向与反向进行 overlap，从而实现高效的分布式模型并行。在上图所示例子中，我们将模型切分成两个 stage，并放置在两张 GPU 卡上进行前向和反向的交替执行，从而提高整体 GPU 的利用率。
与 pipeline 并行不同的是，算子拆分并行是将一个算子内部进行拆分，然后每一个 GPU 负责一部分的计算，随后通过插入一些通信的算子将拆分后的结果进行聚合，从而实现等价的计算。
我们还发现在一些场景下，随着参数规模的进一步增长，单一的并行策略无法满足实际业务对训练性能的需求，并不是一个最优的选择。我们可以同时采用混合多种并行策略的方式来进一步对模型做分布式的加速，比如嵌套流水并行和数据并行、组合算子拆分并行和数据并行等等。通过这种混合策略的方式，我们可以从多个维度对模型的训练做加速。
3. 训练框架的现状和挑战
接下来总结一下当前训练框架的现状和挑战。
（1）首先，当前的一些训练框架，支持的并行策略比较单一，只支持少量的并行策略，缺乏一个统一的抽象来支持所有的并行策略及其组合。比如 Horovod 常用来做数据并行的加速，Gpipe 是一个流水并行的框架，Mesh TensorFlow 只支持算子拆分。
（2）其次，前面提到的这些并行框架，其使用难度较高。用户需要一些领域专家经验才能实现一个高效的分布式模型。
（3）最后，已有的工作会将并行策略实现耦合到模型代码中。当用户想要将这个策略应用到新的模型或者在同一个模型选择其他并行策略的时候，都需要做大量的代码重构，这也增加了这类并行框架的使用成本以及迁移成本。
接下来我们来介绍一下分布式框架 EPL 是怎么解决上述的这些挑战的。
02
分布式框架 EPL 介绍
EPL 是一个统一多种并行策略、易用的分布式深度学习训练框架，它将不同的并行策略进行了统一抽象。在一套分布式训练框架中，支持多种并行策略，包括数据并行、流水并行和算子拆分并行，并支持不同策略的组合和嵌套使用。同时 EPL 提供了灵活应用的接口，用户只需要添加几行代码就可以实现丰富的并行化策略。模型侧不需要去做任何的代码改动。除了用户手动标记并行策略之外，EPL 也支持自动的并行策略探索，包括自动的算子拆分策略以及流水并行中的自动 layer 切分策略等等。EPL 在框架层面也提供了全方位的优化，包括多维度的显存优化、计算优化、通信优化等，从而实现高效的分布式训练性能。
上图是 EPL 的整体架构，整个系统主要分为四个模块：
（1）首先是接口层：EPL 定义了简洁易用的接口，用户可以非常容易地去组合使用各种并行策略。EPL 现在的模型编程接口和原生 TensorFlow 是兼容的。用户可以通过类似 Parallel Annotation 的方式来表达并行策略。
（2）中间表达层是将用户的模型和并行策略转换成一个内部的中间表达。通过 TaskGraph、ParallelStrategy、VirtualDevice 等抽象来表达各种并行策略。其中 Task Graph 是内部模型子图的一个抽象，它是 EPL 的并行化的基本单元。ParallelStrategy 是并行化策略的抽象。VirtualDevice 是对硬件计算资源的抽象。
（3）并行化引擎层会基于 EPL 的中间表达，对计算图做策略探索，进行显存计算、通信优化等等，并自动生成分布式计算图。
（4）Runtime 执行引擎层会将刚才的分布式执行图转换成分布式的 TensorFlow graph 然后再调用 TF 的 runtime 来执行最终的计算。
接下来介绍一下 EPL 的并行化原语。EPL 是通过 strategy annotation 的方式来将模型划分成多个子图。这里的单个子图在 EPL 内部表达为一个 TaskGraph， 然后我们会在 TaskGraph 的基础上去做并行化。当前 EPL 定义了两种并行化原语 replicate 和 split，这两个原语可以用来表达目前所有的并行化策略和它们的组合场景。
（1）replicate 表示的是一个被复制的 TaskGraph，其中 device_count 用于表示计算一个 TaskGraph 副本所需要的 GPU 数量。replicate 配置的模型能表达一个数据并行的计算，它会把 replicate scope 下定义的模型部分去复制多份，然后每一个副本会消费不同的数据，从而来实现这样的计算并行化。
（2）split 则配置了模型的 tensor 拆分计算，它表达了内部的算子拆分。这里device_count 表示的是在 split scope 下的 TaskGraph 需要拆分的数量。拆分后的子图会被放置在多张计算卡上去进行计算。
接下来通过几个例子来详细讲解如何应用这两个并行化原语来表达不同的分布式策略。
第一个是简单的数据并行的例子。在这个例子中，我们只需要标记模型为 replicate 即可实现数据并行的策略。另外一个例子是稍微复杂一点的嵌套流水并行和数据并行的组合策略。在这个例子中，我们将模型通过 replicate scope 划分成三个 stage ，用户可以通过 EPL 的 config 来配置当前的 micro batch 的数量为 5。这三个 TaskGraph 副本一共需要三张卡去计算。如果当前用户申请了六张卡，EPL 会自动地在这个 pipeline 外面嵌套一个数据并行实现 pipeline 并行和数据并行的嵌套。
除了手动地去标记这种并行化策略，EPL 还支持自动的策略的探索。比如在第一个例子中是一个自动的流水并行和数据并行的组合。用户只要去将 auto parallel 配置为 true， 然后设置一个 pipeline 的 stage 数，EPL就会自动地去对模型进行分析，然后对模型执行 stage 的划分以及并行化。
第二个例子是算子拆分和数据并行的组合。用户只需要给定模型拆分份数， EPL 就会自动地去对模型进行拆分探索，然后将模型切分到不同的卡上去做并行计算。
接下来看一下 EPL 如何通过 Parallel Planner 生成一个分布式的 plan，它会结合分布式的标记生成一个并行执行的策略 plan 。
（1）首先 Parallel Planner 的输入是一个 Local Model ，用户可以在 Local Model 上做一些像 replicate 和 split 的标记。通过这样的标记以及给定的一些计算资源、训练配置等等，EPL 会将这些并行元语标记转换成一个或多个 TaskGraph。
（2）第二步是 Virtual Device 生成的流程。Virtual device 是对实际的物理硬件和网络拓扑结构的虚拟化。用户无需关心内部是如何去做划分和映射的。EPL 会自动地生成一个 Virtual device 的 plan 来决定资源应该如何划分以及映射到每一个 TaskGraph 上。
（3）接下来我们会对 TaskGraph 去做对应的复制拷贝，拆分调度通信节点的插入来生成一个分布式的执行 plan。因为在 EPL 中我们是允许不同的 TaskGraph 应用不同的并行化策略，所以 TaskGraph 之间有可能会出现输入和输出不匹配的情况。这个时候 EPL 会自动地去插入一个桥接层，在不同的 TaskGraph 之间去做一个连接和转换。
接下来介绍几个 EPL 中的框架优化点，包括显存优化和通信优化。
其中重算技术 Gradient Checkpoint 是大模型训练中很常见的显存优化手段。它通过保留前向计算中的部分 activation ，反向中通过重算被释放的 activation 用计算换显存的方式来节省显存。因此如何选择一个合适的 checkpoint tensor 对于整体的显存优化效果有很大的影响。在 EPL 中我们提供了一个基于图分析的技术自动选取 checkpoint tensor 来最大化显存的节省。同时我们也会通过对非幂等算子，比如随机性算子 dropout 的重算，防止随机性算子重算引发的收敛性问题。为了进一步提高训练性能，EPL 也会尽量地去减少像模型中穿插的一些通信算子这种开销大的算子的重算来提高整体的训练性能。在 EPL 中如果要开启显存优化功能，用户并不需要在模型侧去做干预，通过一个简单的配置就可以开启 auto GC 的功能。
除了刚才提到的 Gradient checkpoint 技术，EPL 中还实现了类似 DeepSpeed 的 ZeRO 的功能。ZeRO 显存优化技术是通过减少数据并行场景中的模型副本的冗余信息，从而减少整体的峰值显存利用。我们可以通过开启不同级别的 ZeRO 来得到不同程度的显存优化和性能的 trade off。其中 V0 是将 optimizer states 分片存储在不同的卡上。V1 是将 optimizer states 以及 Gradients 都分片存储；V2 会将 optimizer states gradients 以及 weight 都做一个分片存储。EPL 中实现的 ZeRO 也是在基于计算图的分析来自动实现，无需用户的干预。用户只要通过设置 EPL 的 config 中的 zero level 就可以开启不同程度的 ZeRO 的功能。
除了刚才提到的两种技术，EPL 还支持 CPU offload， 这个技术是将训练的存储空间从显存扩展到内存。因为随着模型参数的进一步扩大，如果在有限的资源上想要去训练一个超大的模型，我们就需要通过扩展内存的方式来支持在少量的机器上去训练一个超大的模型。EPL 支持将 weight 和 optimizer states offload 到内存中，并且支持控制 offload 的粒度，将部分的 weight offload 到 CPU 中，这样可以充分地利用显存和内存来达到计算和显存的一个 trade off 。在 EPL 中，我们也可以通过简单的配置就可以开启这个 offload 功能，同时也可以修改一些配置来控制 offload 粒度，或者做一些手动的控制。
除了显存优化，EPL 中也做了比较多的通信优化。一个常见的优化是基于静态图去分析梯度，然后对梯度做一个细粒度的分组融合来实现通信效率的最大化。通过通信分组以及和 gradient 反向计算的 overlap， EPL 可以减少 AllReduce 通信在计算过程中带来的通信的开销。
03
EPL 最佳实践和效果
接下来通过几个案例来展示 EPL 的最佳实践和实际的落地效果。
1. 最佳实践-大规模分类
首先来分享一下大规模图像分类的案例。这个案例中模型由 ResNet50 特征提取层和一个 10 万类的分类层组成。这个模型中，特征提取层 ResNet 部分的参数大概占 90M 左右；分类层主要是由全连接层组成，参数会达到 782M，占整体模型的 90% 左右。
如果我们采取一个单纯的数据并行，从 timeline 里可以看到全连接层的通信会成为整个模型训练的瓶颈。在这里我们采取的混合并行策略是对特征提取层去做数据并行，对超大的全连接层去做算子拆分来消除全连接层的通信瓶颈。通过这样三行代码的改造，我们就可以为这个大规模分类模型实现一个混合的分布式策略。在实测中，我们会发现整体的梯度同步开销可以减少 90% 左右，从而得到一个更好的分布式性能。
采取混合并行，我们在 64 卡的规模下，相比于纯粹的数据并行能够得到 14.8 倍的性能提升。如果后续用户需要去提升他们整体模型的规模，比如说将分类数提高到 1 亿分类，用户只需要去增加 GPU 而无需做任何的代码修改，就可以扩展到一个更大的规模上去。
接下来我们来看一下 BERT large 模型的最佳实践。
2. 最佳实践-Bert large
如果我们对 BERT large 做单纯的数据并行。EPL 相对于 Horovod，64 卡的加速比比 Horovod 的高 1.74 倍，这得益于 EPL 的通信优化。但是相对于单卡来说，数据并行整体的加速比还是不够理想因为 BERT large 参数量大，整体的 AllReduce 通信占比比较高。
我们可以通过用流水并行和数据并行的组合来进行一个分布式的扩展。EPL 通过三行代码就可以完成这样的分布式改造。通过混合 pipeline 并行和数据并行，我们可以将加速比进一步提升 2.32 倍。
3. 最佳实践-中文多模态模型 M6 预训练
最后再分享一下，我们在一个最大的中文多模态模型 M6 预训练模型中的最佳实践。训练 M6 模型存在非常多的挑战，一是算力和资源的需求非常的大，业界去训练一个万亿规模的模型，像 NVIDIA Google 动辄就需要用到几千张卡来做这样的训练。如果我们想要去训练一个十万亿规模的模型，它的算力就要进一步扩大 10 倍。同时训练一个大模型，它的显存需求量也是非常大的。以一个十万亿规模的模型来看，它的模型参数梯度能到 40 TB ，仅考虑存储就需要几千张卡。
在 EPL 中，我们引入了 MOE 稀疏化的计算结构来提升整体的参数和计算的利用率。EPL 支持 MoE layer 的自动专家并行。同时为了提升资源的利用率，降低整体的训练成本，我们也引入了大量的显存优化技术。比如自动 gradient checkpoint 技术，来自动地选取最优的 checkpoint 点来节约整体的 activation 显存开销。同时我们也引入了 CPU offload 技术来支持细粒度 offload 的控制，优化整体显存需求。我们也引入了 Group-wise Apply 的技术来优化 optimizer apply 阶段的显存，以及 Zero 去优化 optimizer states 的显存。我们的通信池化技术可以去控制每一个通信块的大小和并发度来节省通信过程中带来的显存开销。通过这种非常细粒度的显存优化，我们可以进一步地扩大模型的规模。同时我们还采取了一系列的计算通信加速技术，比如分组通信融合半精度通信、拓扑感知的 All2All 通信算子等技术来提高通信效率，结合混合精度、编译优化等技术来提高整体的训练效率。
训练一个这么大规模的模型，EPL 仅仅需要增加 4 行的代码，就可以实现 MoE 算子拆分并行结合数据并行的混合并行模式。在万亿规模的模型上，我们在 480 张 V100 卡，在三天内完成这样的训练。在十万亿的规模的模型上，我们结合各种显存优化技术，在 512 张 V100 卡完成训练。
EPL（EasyParallelLibrary）现在已经开源，大家有兴趣的话可以参照这里面的示例代码使用不同的并行策略。同时我们的论文也在 ATC 22 上发表，供大家阅读。也欢迎大家加入到我们的钉钉群中，随时与我们进行交流。