TUTEL: ADAPTIVE MIXTURE-OF-EXPERTS AT SCALE

发表时间: 2022-06 · arXiv:2206.03382

原文: https://arxiv.org/abs/2206.03382

Changho Hwang * 1, Wei Cui * 1, Yifan Xiong * 1, Ziyue Yang * 1, Ze Liu 1, Han Hu 1, Zilong Wang 2, Rafael Salas 2, Jithin Jose 2, Prabhat Ram 2, Joe Chau 2, Peng Cheng 1, Fan Yang 1, Mao Yang 1, Yongqiang Xiong 1

1 Microsoft Research Asia & Azure
2 Microsoft Azure

速读

一句话结论 本文针对大规模混合专家模型（MoE）中动态路由导致的负载不均与计算低效问题，提出了一个名为 TUTEL 的全栈自适应加速系统，通过零成本的动态并行切换和自适应流水线技术，在 2048 张 A100 显卡上实现了单层 MoE 5.75 倍的加速，并显著提升了 SwinV2-MoE 等真实视觉模型的端到端速度与精度。

要解决什么问题 稀疏门控混合专家模型（MoE）的性能高度依赖于令牌路由机制。在运行时，门控网络会动态地将输入令牌分发给最合适的专家子模型，这导致每个专家在不同迭代中需要处理的令牌数量（即专家容量）不断剧烈变化。然而，现有的深度学习系统大多采用静态的并行策略和流水线策略。由于最优的执行策略会随动态工作负载而改变，静态策略无法在所有情况下保持高效，常常导致计算与通信无法有效重叠、GPU 利用率低下。如果在运行时强行切换并行策略，又会因为重新切分数据和迁移模型权重而引入巨大的通信开销。此外，现有的 All-to-All 集合通信在扩展到大规模集群时，由于节点间传输的数据块被切得很小，无法充分利用高速网络带宽；且传统的密集矩阵运算在处理 MoE 的数据分发与合并时会产生大量无效的零乘法计算，严重拖慢了非专家计算的延迟并占用了过多显存。

怎么做的 TUTEL 的核心思路是设计一种统一的数据与权重布局，使得系统能在每次迭代时根据当前工作负载，以零成本动态切换到最优的并行和流水线策略。关键设计由四个部件构成。首先是自适应并行切换，系统将复杂的并行策略简化为数据并行（DP）和专家加数据加模型并行（EP+DP+MP）的组合。通过引入控制参数 $r$，将所有 GPU 划分为大小为 $\lceil(W/E)/r\rceil$ 的组（$W$ 为 GPU 总数，$E$ 为全局专家数）。系统在组内执行 DP，在组间执行 MP。由于底层统一采用了 ZeRO-DP Stage-3 的权重分区机制，每个设备始终持有固定的权重切片，切换 $r$ 值时完全无需迁移模型参数，从而实现了零成本切换。其次是自适应流水线，TUTEL 不对整个 MoE 层进行切分，而是仅对 All-to-All 通信和专家计算进行虚拟分区，利用多流并发执行。系统在运行时会根据当前的专家容量，查表联合搜索出最优的流水线深度和通信算法。第三是二维分层（2DH）All-to-All 算法，专门解决大规模下的小消息传输问题。它通过高效的跨步内存复制，先在本地将发往同一远程 GPU 的非连续小数据块对齐并聚合成大数据块，再进行节点间通信，从而保持了与 GPU 数量无关的恒定低延迟。最后是快速编码与解码，TUTEL 摒弃了传统的高复杂度密集矩阵乘法，开发了 SIMT 高效的 GPU 稀疏内核，直接处理令牌的分发与合并，将时间复杂度从密集的 $O(T \cdot E \cdot Cg \cdot D)$ 大幅降至稀疏的 $O(T \cdot k \cdot D)$（其中 $T$ 为令牌数，$k$ 为路由数量，$D$ 为隐藏层维度），彻底消除了无效的零乘法计算，显著降低了延迟并节省了大量显存。

效果如何 实验在最多包含 2048 张 A100 GPU 的 Azure 集群上进行，对比基线为代表主流静态执行路线的 Fairseq 和 DeepSpeed MoE 框架。在单层 MoE 的扩展性测试中，TUTEL 在 16 张和 2048 张 GPU 上分别比基线提速了 4.96 倍和 5.75 倍。其中，快速编码与解码在小规模下贡献了高达 3.52 倍的加速，而 2DH All-to-All 则在大规模下带来了 4.25 倍的核心增益。在端到端真实场景中，作者训练了视觉模型 SwinV2-MoE。在 8 到 128 张 GPU 的规模下，TUTEL 相比 Fairseq 实现了 1.14 至 1.55 倍的训练提速和 1.95 至 2.11 倍的推理提速。此外，SwinV2-MoE-B 模型在 ImageNet-22K 预训练和 COCO 目标检测等下游任务中的准确率均超越了同等规模的密集模型 SwinV2-B，证明了稀疏 MoE 模型在下游视觉任务上的优越性。不过，该方法也存在一定代价与局限：2DH 算法由于引入了额外的内存复制，在处理大尺寸消息时延迟反而更高，因此必须与传统线性算法进行动态适配；此外，直接对 MoE 模型进行下游微调会导致性能下降，必须在微调期间固定所有 MoE 层的权重才能获得正向收益。

A1 主要贡献

本文针对稀疏门控混合专家（MoE）模型在扩展至万亿参数规模时面临的系统挑战，提出了一种名为 TUTEL 的高度可扩展的堆栈设计与实现。

核心问题: MoE 模型的性能依赖于其令牌路由机制，该机制在运行时动态地将输入令牌（token）转发给最合适的子模型（专家）。这种动态性导致每个专家的工作负载在运行时不断变化。然而，现有的深度学习系统（包括最新的 MoE 框架）大多采用静态执行方式，即静态的并行策略和流水线策略，无法适应 MoE 动态变化的工作负载，从而导致计算效率低下。具体而言，由于最优并行策略随动态工作负载而变化，静态并行常常无法发挥最佳性能，而在运行时切换并行策略会引入巨大的数据重分布开销。

研究目标: 设计一个能够适应 MoE 动态工作负载的系统，以在任何规模下都能高效运行 MoE 模型。该系统需要能够动态调整并行和流水线策略，且切换成本要尽可能低，同时保持算法的数学等价性。

核心创新点:
1. 自适应并行切换 (Adaptive Parallelism Switching)：TUTEL 的关键机制是设计了一种统一的模型参数和输入数据的分布布局，该布局能够同时支持多种并行策略。这使得系统可以在每次迭代中根据当前的工作负载动态地切换到最优的并行策略，而无需重新格式化数据或迁移模型权重，从而实现了零成本切换。
2. 全面的 MoE 加速技术栈：在自适应并行切换的基础上，TUTEL 实现了一系列 MoE 加速技术，包括：
* 自适应流水线 (Adaptive Pipelining)：动态调整流水线深度，以最优地重叠计算与通信。
* 二维分层 (2DH) All-to-All 算法：一种新颖的集合通信算法，用于在大规模场景下减少小消息通信开销，提高网络带宽利用率。
* 灵活 All-to-All (Flexible All-to-All)：一种抽象接口，确保专家计算（矩阵乘法）在不同规模下都能保持高效。
* 快速编码/解码 (Fast Encode/Decode)：使用 SIMT 高效的稀疏操作优化 MoE 的数据分发和合并过程，显著降低了非专家计算的延迟并节省了 GPU 内存。

主要成果:
* 单层 MoE 性能：与现有最优系统相比，TUTEL 在 16 个和 2,048 个 A100 GPU 上分别将单个 MoE 层的速度提升了 4.96 倍和 5.75 倍。
* 端到端模型性能：TUTEL 被用于实现和运行一个名为 SwinV2-MoE 的真实世界视觉模型。与 Fairseq 相比，SwinV2-MoE 的训练和推理速度分别提升了 1.55 倍和 2.11 倍。
* 模型效果：SwinV2-MoE 模型在预训练和下游计算机视觉任务（如 COCO 目标检测）中的准确性均优于其对应的密集模型，证明了 TUTEL 在训练真实世界 AI 模型方面的有效性和准备就绪。


图 1. Swin Transformer V2 (Liu et al., 2021; 2022) 的 MoE 版本的 thin-tiny (左) 和 base (右) 模型在端到端训练过程中 MoE 层的动态变化工作负载。y 轴表示运行时所需的专家容量，它指示了工作负载的大小（详见 2.1 节）。为了视图清晰，图中仅显示了模型总共 10 个 MoE 层中的第 1、4、10 层。


图 2. 一个跨三个 GPU 的 MoE 层示例，专家 Ei 位于 GPU i 上。G0 表示所有 GPU 共享的门控函数。不同的颜色或图案表示不同的样本（输入的列），颜色的不同梯度表示样本内的不同令牌（输入的行）。此示例显示每批次 2 个样本，每个样本 6 个令牌，并采用容量因子为 1.0 的均匀分发的 top-1 路由——详见 2.2 节。

A3 背景知识与动机

本节介绍混合专家模型的动态特性及其在大规模训练中的低效问题。

2.1 背景与相关工作

2.2 静态并行

2.3 静态流水线

A2 方法细节

TUTEL 是一个全栈 MoE 系统，支持完整的 MoE 层并提供自适应优化。所有优化对 DNN 模型开发者透明，不会改变深度学习框架的接口，可以轻松集成到其他框架中。

表 3. 符号说明。

3.1 自适应并行切换

3.1.1 值得进行并行切换的最小子集是什么？
3.1.2 零成本可切换并行的执行流

3.2 针对 Linear 和 2DH All-to-All 的自适应流水线

3.3 最优并行与流水线字典

4 实现

4.1 特性

4.2 优化

A4 实验

实验环境

实验结果

5.1 自适应 MoE 评估
5.2 单个 MoE 层扩展性
5.3 真实世界问题应用：SwinV2-MoE

A5 结论

本文从系统角度分析了 MoE 的关键动态特性，并设计了一个自适应系统 TUTEL 来解决相应问题。TUTEL 主要包含两大方面：用于优化专家执行的自适应并行，以及用于解决 MoE 层中低效且不可扩展的 dispatch/combine 操作的自适应流水线。在 Azure A100 集群上使用 2048 个 GPU 的评估表明，TUTEL 对于单个 MoE 层实现了高达 5.75 倍的加速。TUTEL 能够赋能真实世界 SOTA 深度学习模型的训练和推理。本文以 SwinV2-MoE 为例，展示了 TUTEL 在开发先进模型中的应用，并证明了 MoE 在计算机视觉任务中相对于密集模型的有效性。

A6 附录

A. 二维分层（2DH）All-to-All

1: procedure ALL2ALL_LINEAR(output, input)
2:   n ← ngpus, S ← sizeof input
3:   chunksize ← S / n
4:   for r = 0; r < n; r++ do
5:     loc ← r × chunksize, peer ← r
6:     ncclSend(input[loc], chunksize, peer)
7:     ncclRecv(output[loc], chunksize, peer)
8:   end for
9: end procedure

1: procedure STRIDEMEMCPY(output, input, chunksize, row, col)
2:   for i = 0; i < row × col; i++ do
3:     j ← i % row × col + i / col
4:     output[j × chunksize : (j+1) × chunksize] ← input[i × chunksize : (i+1) × chunksize]
5:   end for
6: end procedure
7: procedure ALL2ALL_2DH(output, input)
8:   // 步骤 1: 节点内 All-to-All
9:   strideMemcpy(buffer, input, chunksize, ngpus_per_node, nnodes)
10:  for g = 0; g < ngpus_per_node; g++ do
11:    loc ← g × nnodes × chunksize, peer ← g + node_rank × ngpus_per_node
12:    ncclSend(buffer[loc], nnodes × chunksize, datatype, peer, comm)
13:    ncclRecv(output[loc], nnodes × chunksize, datatype, peer, comm)
14:  end for
15:  strideMemcpy(buffer, output, chunksize, nnodes, ngpus_per_node)
16:  // 步骤 2: 节点间 All-to-All
17:  for n = 0; n < nnodes; n++ do
18:    loc ← n × ngpus_per_node × chunksize, peer ← local_rank + n × ngpus_per_node
19:    ncclSend(buffer[loc], ngpus_per_node × chunksize, datatype, peer, comm)
20:    ncclRecv(output[loc], ngpus_per_node × chunksize, datatype, peer, comm)

B. SIMT 高效的快速编码和解码

C. 关于 SwinV2-MoE 的更多结果