关于分布式计算:阿里开源-支持10万亿模型的自研分布式训练框架EPLEasy-Parallel-Library

作者：王林、飒洋

导读

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

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，资源老本十分高；
如何降本增效，组合应用各种技术和办法来缩小须要的资源，进步训练的速度；
为了应答以后分布式训练的挑战，咱们研发了分布式训练框架 EPL，将不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略。EPL 同时提供简洁易用的接口，用户只需增加几行 annotation 即可实现并行策略的配置，模型代码不须要改变。在用户无感的状况下，EPL 通过进行各种显存、计算、通信优化，打造高效的分布式训练框架。

2. 次要个性

• 多种并行策略对立：在一套分布式训练框架中反对多种并行策略（数据 / 流水 / 算子 / 专家并行）和其各种组合、嵌套应用；
• 接口灵便易用：用户只需增加几行代码就能够应用 EPL 丰盛的分布式并行策略，模型代码无需批改；
• 主动并行策略摸索：算子拆分时主动摸索拆分策略，流水并行时主动摸索模型切分策略；
• 分布式性能更优：提供了多维度的显存优化、计算优化，同时联合模型构造和网络拓扑进行调度和通信优化，提供高效的分布式训练；

3. 开源地址

EPL(Easy Parallel Library)的开源地址是：
https://github.com/alibaba/Ea…

咱们同时提供了 model zoo，欢送大家试用：
https://github.com/alibaba/Fa…

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。通过这两种并行化接口，能够表白出各种不同的并行化策略，例如：

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

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

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

同时 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 反对各种显存优化技术的组合应用，达到显存的极致优化。咱们在 T5 模型上开启 GC+ZeRO+ 显存优化的 AMP 技术，在性能放弃不变的状况下，显存升高 2.6 倍。

利用场景

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 的流水并行 + 数据并行，相较于数据并行失去了 66% 的晋升。

Roadmap

咱们决定建设开源生态次要有如下的思考：

• EPL 发祥于阿里云外部的业务需要，在服务外部业务的过程中，咱们积攒了大量的教训和了解，很好的反对了大规模、多样性的业务场景，在咱们本身随着业务需要的迭代逐步欠缺的同时，也心愿可能开源给社区，将本身积攒的教训和了解也同样回馈给社区，心愿和深度学习训练框架的开发者或深度学习从业者之间有更多更好的交换和共建，为这个行业奉献咱们的技术力量。
• 咱们心愿可能借助开源的工作，收到更多实在业务场景下的用户反馈，以帮忙咱们继续欠缺和迭代，并为后续的工作投入方向提供输出。
• 同时咱们心愿借助开源的工作，能吸引一些气味相投的同学、公司或组织来参加共建，继续欠缺深度学习生态。

后续咱们打算以两个月为单位公布 Release 版本。EPL 近期的 Roadmap 如下：

• 继续的性能优化和稳定性改良；
• 通用算子拆分性能；
• 主动拆分策略摸索的根底版；
• 主动流水并行策略摸索；

此外，在中长期，咱们在上面几个探索性的方向上会继续投入精力，也欢送各种维度的反馈和改良倡议以及技术探讨，同时咱们非常欢送和期待对开源社区建设感兴趣的同行一起参加共建。

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

欢送退出 EPL 用户交换群，钉钉群号：34761216。

参考文献

1. Whale: Scaling Deep Learning Model Training to the Trillions https://arxiv.org/abs/2011.09208
2. https://openai.com/blog/ai-an…
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