[SIGMOD 2023]大规模稀疏预训练模型的训练加速

关键词

Deep Learning System

Distributed Computing

Sparse Model

导 读

大规模稀疏预训练模型的训练加速

近年来，可用的训练数据不断增加，深度学习研究员们倾向于使用更大参数量的模型来拟合这些知识，也取得了更好的效果。然而，随着模型参数量的不断扩大，模型训练所需要的计算量也在不断增长，这大大提高了模型训练的成本。因此，很多研究人员转向使用混合专家结构 (Mixture-of-Experts，MoEs) 来高效地扩大模型参数，并通过稀疏门控网络引入模型稀疏性，在不显著增加计算量的情况下增加模型的参数，在各类下游任务上取得了很好的效果。然而，由于这种稀疏的模式面临着动态路由和负载不均衡的问题， 现有系统无法高效地处理此类计算模式。

在这篇论文中，我们提出了FlexMoE，一种面向大规模稀疏预训练模型的深度学习框架，从系统设计的角度解决了MoE模型中由动态数据流引起的低效问题。首先，我们观察了不同MoE模型训练过程中专家的负载情况，这驱动我们通过动态专家管理和设备放置机制来克服动态路由和负载不均的问题。然后，我们在现有的DNN系统上增加了一个新的调度模块，用来监控训练时的数据流，制定专家调度计划，动态调整模型与硬件的映射关系。此外，我们利用了启发式算法来动态优化训练期间的设备放置。我们在自然语言处理模型（例如 BERT 和 GPT）和视觉模型（例如 Swin）上进行了实验。结果表明，与现有系统相比，FlexMoE 在实际工作负载上可以实现更优秀的性能—— FlexMoE 的平均性能优于 DeepSpeed 1.70 倍，最高可达 2.10 倍；平均性能优于 FasterMoE 1.30 倍，最高可达 1.45 倍。

论文链接：https://dl.acm.org/doi/abs/10.1145/3588964 

01

问 题 背 景

大规模预训练模型，如 Transformer 模型，已经在各种复杂的任务上推动了深度学习的发展，包括自然语言处理（例如 BERT、GPT、T5）、计算机视觉（例如 ViT、Swin）和广告推荐（例如 M6）。这些模型被称为基础模型，因为它们在数百 GB 的数据上进行训练，仅微调便可以适配各类下游任务。然而，近期的研究表明，模型的表现效果正比于训练数据集规模、参数规模以及训练的计算量。因此，研究员们倾向用更大的模型、更大的数据来拟合目标任务。当前表现最好的基础模型已经增长到万亿参数，训练变得极其昂贵且耗时。

为了降低模型训练的成本，通常会采用稀疏门控混合专家结构（Mixture-of-Experts，MoE）。这种稀疏 MoE 架构通过扩展专家网络来增加模型的参数量，同时利用稀疏门控网络(如Top-2路由)来控制计算成本。与此同时，为了满足 MoE 模型训练过程中的显存需求，提出了专家并行的方法，将专家网络分配到不同的 GPU 上，并通过 AllToAll 通信将数据发送至指定的专家，从而实现分布式训练。然而，由于 MoE 层采用条件计算（门控网络）进行专家的稀疏激活，这导致了专家间的负载不平衡，进而引起不同 GPU 间的负载不均衡，降低了 GPU 的计算利用率。更为复杂的是，在训练过程中，由于路由波动现象的存在，专家间的负载会在不同的训练阶段发生动态变化，而现有的训练系统尚未能有效解决这一问题。

02

挑 战 和 机 遇

挑   战

在MoE模型的训练过程中，面临着专家间负载不平衡和路由波动的挑战。由于当前轮的专家负载是由输入数据和门控网络参数共同决定的，我们无法在执行前预测最优的专家到 GPU 设备的映射方案。因此，主要的难题在于如何在 GPU 高速计算的同时，应对变化的专家负载，有效地设计和调整专家到设备的映射方案，并且高效地实施这些调整。

为解决负载不均衡的问题，现有方法主要通过修改模型定义，包括：(1) 引入balance loss，驱使门控网络以尽可能均衡的方式分配负载；（2）设定容量（capacity）限制，强制每个专家处理的数据量不超过其最大容量，从而实现负载均衡。我们首先对这些技术的影响进行了探究，结果如图1所示。当balance loss的系数从0增加到0.05时，GPU利用率从18.77%提升至63.30%，但模型的准确率却从94.588%下降到93.981%。这一结果表明，通过调整token的分配来强制实现工作负载的平衡，将不可避免地导致系统效率和模型质量之间的权衡。

图1   Swin-MoE在不同balance loss系数下的Top-5正确率

机   遇

通过观察MoE模型训练时专家间的动态不平衡工作负载（如图2，其中不同颜色区间代表不同专家的负载），我们总结了两个关键特性：（1）skewness：专家间负载的偏斜性，如图中Top 10的专家处理了近75%的数据（总共64个专家）；（2）smoothness and continuousness：每个专家的负载都在连续且缓慢的变化。这些特性意味着在短时间内，最优的专家到GPU设备的映射方案不会显著变化，因此我们可以基于当前的路由决策来调整映射方案，无需预测未来几轮训练的最优方案。与此同时，由于深度学习训练是迭代式的，我们可以对计算和通信进行建模，预估出不同映射方案的收益和开销，从而找到最佳方案。

图2   MoE模型训练时专家间动态不平衡的工作负载

03

FlexMoE 设 计 方 案

3-01

总体流程

为了解决专家的动态负载不平衡问题，我们设计了一种动态的专家管理机制，该机制在训练过程中自适应地调整专家到GPU设备的映射方案和Token的分配。如图3所示，FlexMoE的工作流程除了包括常规的MoE系统训练流程外，还引入了两个新的组件：调度模块（Scheduler）和策略制定模块（Policy Maker）。调度模块负责监控专家的实时负载，并在当前的不平衡指标超过预定阈值时，将负载信息发送给策略制定模块。策略制定模块根据接收到的负载信息和当前的专家摆放方案，计算出最优的专家摆放方案，然后生成修改指令并发送调度模块。调度模块随后与当前深度学习框架的执行器（Executor）进行交互，并在运行时触发专家摆放方案的修改。

图3    FlexMoE系统设计

3-02

动态的专家管理机制

为了优化专家到设备映射策略，我们提出了一个名为vExpert的新抽象。

vExpert定义了专家负载的最小调度单位，与动态的专家管理机制可以很好的结合。vExpert有以下3个特性：

（1） 每个vExpert只能被分配给一个专家

（2） 同一个GPU上相同专家的vExpert共享权重

（3） 对于一个专家，负载在其对应的vExpert间均匀分配

vExpert的抽象简化了复制专家、增加/减少副本数和分配负载的决策过程，有效缩小了优化问题的搜索空间。我们设计了三种修改原语：Expand（增加副本数）、Shrink（减少副本数）和Migrate（交换副本位置）：

（1） Expand：在专家的负载增加时，为专家分配额外的vExpert资源

（2） Shrink：在专家的负载减少时，释放vExpert资源

（3） Migrate：通过在vExperts间交换模型状态，以减少训练时的同步成本（AllReduce）

3-03

优化方案

问 题 建 模

首先，我们对这个优化问题进行了形式化的定义。为了简化问题，我们考虑一个包含门控网络和一组专家E (e∈ E)的单个MoE层。在给定当前的专家到设备的映射P ((e, g) ∈ P，表示GPU g有专家e)，以及当前迭代的Token分配I (Ie,g表示专家e到GPUg的Token数量)的情况下，训练成本可以表示为T (I, P)。我们的目标是最小化这个训练成本，可以表示为公式1。

公式1   优化目标建模

其中第一项代表计算开销，第二项代表all-to-all的通信开销，第三项代表专家副本间的同步开销。为了最小化训练开销同时不影响模型的精度，我们通过在训练过程中根据负载自适应地调整专家到设备的映射P和以及当前迭代的Token分配I。

调 度 模 块

为了描述当前状态下各GPU之间的负载分配情况，我们定义了一个平衡比率（balance ratio），如公式2所示。在训练过程中，调度模块（Scheduler）会实时监控训练负载并计算当前的负载均衡情况，当平衡比率超过预定的阈值时，调度模块将会触发策略制定模块（Policy Maker）来生成映射修改方案。随后，调度模块负责与训练框架的执行器（Executor）进行交互，以执行这些修改，包括Expand和Shrink。具体流程可参见算法1。此外，调度模块会持续地调用Migrate来优化同步代价。

公式2   Balance Ratio定义

算法1

策 略 制 定 模 块

我们基于映射方案的修改原语（Expand, Shrink, Migrate），设计了一种高效的基于vExpert的调度算法，用于生成修改操作序列，具体来说，我们采用了一种基于代价模型（Cost Model）的搜索规划方法，该方法可以预估调整后的训练时间来选择不同的决策方案。

算法2

04

实 验

我们在Azure VMs上进行了实验，每台VM配备了192核的AMD CPU和8个NVIDIA Ampere A100 GPU，VM之间通过8个InfiniBand NICs连接，总带宽达到16Tb/s。我们选取了NLP领域的BERT和GPT，以及CV领域的Swin等常见模型，将我们的系统（FlexMoE）与主流的DeepSpeed-MoE和FasterMoE进行了比较。

4-01

模型效果对比

为了展示不丢弃训练Token的重要性，我们比较了FlexMoE和DeepSpeed在不同模型和任务上的模型效果。如表1所示，FlexMoE在几乎所有任务上都优于DeepSpeed，这表明现有MoE系统（例如，DeepSpeed）的容量（Capacity）限制可能导致模型质量下降。

注：考虑到训练成本，Swin-MoE模型只训练了100个epoch，因此准确率略低于开源基准。

表1   模型效果对比

4-02

训练效率对比

我们还在训练效率方面评估了FlexMoE与其他SOTA系统的性能。为了衡量训练效率，我们记录了达到目标模型质量所需的训练时间，结果如图4所示。实验结果表明，FlexMoE在训练效率方面表现最佳，平均比DeepSpeed高出1.70倍，最高可达2.10倍，平均比FasterMoE高出1.30倍，最高可达1.45倍。正如上文所提到的，尽管DeepSpeed由于其有限的容量而获得了最小的迭代时间，但它通过丢弃Token来跳过专家网络，因此需要更多的迭代次数才能收敛。FasterMoE提出了阴影专家策略来实现负载平衡，该策略将受欢迎的专家复制到所有GPU上。然而，由于其粗粒度的专家管理（即在1个GPU上或在所有GPU上），它回退到了次优解。随着GPU数量的增加，FasterMoE受到专家副本的全局同步的影响。

图4   不同系统的训练效率对比

此外，我们还进行了消融实验，以比较不同的平衡比率（balance ratio）定义和调度策略对训练的影响，并评估了我们的成本建模与实际运行时间的差异。

图5   消融实验

总     结

在本研究中，我们提出了FlexMoE，这是一种解决大规模MoE模型训练过程中遇到的动态不平衡挑战的新颖解决方案。通过在现有的DNN框架上集成调度模块，FlexMoE监控数据流量，制定专家调度计划，并在训练过程中动态调整专家到设备的映射。我们在六个流行的MoE模型上进行的实证结果表明，FlexMoE的性能平均比DeepSpeed高出1.70倍，最高可达2.10倍，同时也比FasterMoE平均高出1.30倍，最高可达1.45倍。

详细了解本工作，请访问下方链接地址：

论文链接：https://dl.acm.org/doi/abs/10.1145/3588964 

References：

[1] Jared Kaplan,Sam McCandlish, Tom Henighan, Tom B Brown,Benjamin Chess, Rewon Child,Scott Gray,Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. CoRR abs/2001.08361 (2020).

[2] William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research 23, 120 (2022), 1–39.

[3] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. 2022. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 120–134.

[4] Xiaonan Nie, Pinxue Zhao, Xupeng Miao, Tong Zhao, and Bin Cui. 2022. HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System. CoRR abs/2203.14685 (2022).

[5] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. 2022. DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale. In International Conference on Machine Learning. PMLR, 18332–18346.

END

欢迎关注本公众号，帮助您更好地了解北京大学数据与智能实验室（PKU-DAIR），第一时间了解PKU-DAIR实验室的最新成果！

实验室简介

北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文100余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。