近日，阿里云机器学习 PAI 对于深度学习模型高效的分布式训练框架的论文《Whale: Efficient Giant Model Training over Heterogeneous GPUs》被计算机系统畛域国内顶级学术会议 USENIX ATC’22 接管。
Whale 是阿里云机器学习 PAI 平台自研的分布式训练框架，开源后的名称是 EPL(Easy Parallel Library)，Whale 通过对不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略，并进行显存、计算、通信等全方位的优化，来提供易用、高效的分布式训练框架。Whale 提供简洁易用的接口，用户只需增加几行 annotation 即可组合各种混合并行策略。同时 Whale 提供了基于硬件感知的自动化分布式并行策略，感知不同硬件的算力、存储等资源，依据这些资源来正当的切分模型，平衡不同硬件上的计算量，最大化计算效率。借助 Whale 的硬件感知负载平衡算法，Bert-Large、Resnet50 和 GNMT 模型 在异构 GPU 训练上提速 1.2 至 1.4 倍。同时，应用 Whale 框架，万亿 M6 模型应用 480 张 V100 在 3 天内实现预训练。相比此前业界训练等同规模的模型，节俭算力资源超 80%，且训练效率晋升近 11 倍。进一步应用 512 GPU 在 10 天内即训练出具备可用程度的 10 万亿模型。
背景和挑战
随着近些年深度学习的火爆，模型的参数规模也增长迅速，OpenAI 数据显示：
● 2012 年以前，模型计算耗时每 2 年增长一倍，和摩尔定律保持一致；
● 2012 年后，模型计算耗时每 3.4 个月翻一倍，远超硬件倒退速度；
特地最近两年模型参数规模飞速增长，谷歌、英伟达、阿里、智源研究院都公布了万亿参数模型，有大厂也公布了百亿、千亿参数模型。随着模型参数规模增大，模型成果也在逐步提高，但同时也为训练框架带来更大的挑战。以后曾经有一些分布式训练框架，例如：Horovod、Tensorflow Estimator、PyTorch DDP 等反对数据并行，Gpipe、PipeDream、PipeMare 等反对流水并行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore 等反对算子拆分，但当咱们要训练一个超大规模的模型时会面临一些挑战：
- 以后的分布式训练框架只反对大量的并行策略，不足一个对立的形象来反对所有的并行策略及其混合策略。
- 实现简单的并行策略须要大量的模型代码革新和对底层零碎的了解，大大增加了用户的应用难度。
- 因为集群中异构 GPU 计算能力和显存的差别，动态的并行化策略无奈充分利用异构资源实现高效训练。
破局
为了应答以后分布式训练的挑战，咱们研发了高效、通用、硬件感知的分布式训练框架 Whale。Whale 形象并定义了两个分布式原语 (replicate 和 split)。用户能够通过在模型上增加几行原语标记，即来表白和转换各种并行化策略及其组合，极大升高了分布式框架的应用门槛。Whale runtime 将用户的标记信息交融到计算图中，主动实现模型的并行化优化。同时 Whale 提供了基于硬件感知的自动化分布式并行策略，优化在异构 GPU 集群上分布式训练性能。Whale 的设计很好地均衡了模型用户的干涉和系统优化机会，让每一个算法工程师都能轻松高效训练分布式大模型工作。
技术架构
Whale 框架如下图所示，次要分为以下几个模块：
● 接口层： 用户的模型编程接口基于 TensorFlow，同时 Whale 提供了易用的并行化策略表白接口，让用户能够组合应用各种混合并行策略；
● 两头表白层： 将用户模型和并行策略转成化外部表白，通过 TaskGraph、VirtualDevices 和策略形象来表白各种并行策略；
● 并行化引擎层： 基于两头表白，Whale 会对计算图做策略摸索，联合硬件资源进行显存 / 计算 / 通信优化，并主动生成分布式计算图。
● Runtime 执行引擎： 将分布式执行图转成 TFGraph，再调用 TF 的 Runtime 来执行；
并行化策略表白
Whale 通过 strategy annotation 的形式来划分模型为多个 TaskGraph，并在此基础上进行并行化。Whale 有两类 strategy：replicate 和 split。通过这两种并行化接口，能够表白出各种不同的并行化策略，例如：
● 数据并行: 上面这个例子是一个数据并行的例子，每个模型正本用一张卡来计算。如果用户申请了 8 张卡，就是一个并行度为 8 的数据并行任务。
● 流水并行：在上面的例子里，模型被切分成 2 个 TaskGraph, “stage0″ 和 ”stage1″，用户能够通过配置 pipeline.num_micro_batch 参数来设定 pipeline 的 micro batch 数量。在这个例子里，”stage_0″ 和 ”stage_1″ 组成一个模型正本，共须要 2 张 GPU 卡。如果用户申请了 8 张卡，Whale 会主动在 pipeline 外嵌套一层并行度为 4 的数据并行（4 个 pipeline 正本并行执行）。
● 算子拆分并行：在以下例子中，Whale 会对 split scope 下的模型定义做拆分，并搁置在不同的 GPU 卡上做并行计算。
● 同时 Whale 反对对上述并行策略进行组合和嵌套，来组成各种混合并行策略，更多示例能够参考开源代码的文档和示例。
Parallel Planner
Paraller Planner 是 Whale runtime 的重要一环，它的职责是生成一个高效的分布式执行 plan。Parallel Planner 的流程蕴含 (a) Paraller Planner 的输出蕴含用户模型、用户标记 (可选)、硬件计算资源和其余用户配置配置。(b) 将物理计算资源映射成 VirtualDevice, 用户无需放心如何将算子搁置在分布式物理设施上。(c) 模型被宰割成 TaskGraph 子图。因为 Whale 容许对不同的 TaskGraph 利用不同的分布式策略，所以在 TaskGraph 之间可能存在输出 / 输入 shape 不匹配。在这种状况下，Paraller Planner 将主动在两个 TaskGraphs 之间插入相应的桥接层。
硬件感知的负载平衡算法
当模型训练资源蕴含异构硬件（比方混合了 V100 和 T4），硬件感知的负载平衡算法能够进步在异构资源下的训练效率。Whale 设计了两种均衡策略：Intra-TaskGraph 和 Inter-TaskGraph 的均衡。
（1）对于 Intra-TaskGraph 的计算均衡，Whale 会 profile 出模型的算力 FLOP，并依照机器的计算能力按比例调配对应的计算负载，以此平衡每个 step 不同卡型上的模型计算工夫。对于数据并行，咱们能够通过调整不同正本上的 batch 大小来实现计算负载的平衡（放弃全局 batch 不变）。对于算子拆分，咱们能够通过不平均的维度拆分来实现不同卡上子模块的计算负载平衡。
（2）当多个 TaskGraph 在多个异构 GPU 上执行流水并行，咱们须要 Inter-TaskGraph 的计算均衡来晋升整体的计算效率。因为流水并行的计算特点，后面的 TaskGraph 会比前面的 TaskGraph 多缓存一些前向计算的后果，因而对于 Transformer 类模型平衡 layer 切分的状况下，后面的 TaskGraph 会有更大的显存压力。因而在 TaskGraph 搁置的时候，Whale 会优先将后面的 TaskGraph 搁置在显存容量更大的 GPU 卡上。与此同时，Whale 会依照 GPU 卡的算力去切分模型，使得切分后的 TaskGraph 计算 load 和 GPU 算力成正比。
利用示例
借助 Whale 框架，咱们 4 行代码实现 M6 模型数据并行 + 专家并行的混合并行训练。
如下图所示，MoElayer 采纳专家并行，其余 layer 采纳数据并行：
并首次在 480 V100 上，3 天内实现万亿 M6 模型的预训练。相比此前业界训练等同规模的模型，此次仅应用 480 张 V100 32G GPU 就胜利训练出万亿模型 M6，节俭算力资源超 80%，且训练效率晋升近 11 倍。进一步应用 512 GPU 在 10 天内即训练出具备可用程度的 10 万亿模型。
结语
Whale 通过对立的形象反对各种并行化策略，并提出简洁的并行化原语来进步框架的易用性。同时，Whale 提供了异构硬件感知的主动分布式计算图优化，实现高效的分布式模型训练。咱们心愿 Whale 可能成为一个大规模深度学习训练的基石，进一步推动模型算法翻新和系统优化，使大模型训练技术可能广泛应用在理论生产环境中。Whale 曾经开源 (https://github.com/alibaba/Ea…)，欢送大家来试用和共建。
● 论文名称：Whale: Efficient Giant Model Training over Heterogeneous GPUs
● 论文作者：贾贤艳，江乐，王昂，肖文聪，石子骥，张杰，李昕元，陈浪石，李永，郑祯，刘小勇，林伟
● 论文链接：https://arxiv.org/pdf/2011.09…
● 开源链接：https://github.com/alibaba/ea…
参考文献：
[1] Junyang Lin, An Yang, Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, et al. M6-10t: A sharing-delinking paradigm for effi- cient multi-trillion parameter pretraining. arXiv preprint arXiv:2110.03888, 2021.
[2] Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, Jie Zhang, Jianwei Zhang, Xu Zou, Zhikang Li, Xiaodong Deng, Jie Liu, Jinbao Xue, Huiling Zhou, Jianxin Ma, Jin Yu, Yong Li, Wei Lin, Jingren Zhou, Jie Tang, and Hongxia Yang. M6: A chinese multimodal pretrainer, 2021.
[3] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
[4] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: gen- eralized pipeline parallelism for dnn training. In Pro- ceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1–15, 2019.