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

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 背景与相关工作

  • 稀疏门控混合专家模型（MoE）的机制。MoE 模型包含多个专家模型，每个专家负责处理其擅长的子任务，共同解决整个任务。在大规模分布式 DNN 模型中，它通过一个跨 GPU 的层来部分交换来自不同 GPU 的隐藏特征【索引7，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】；【索引14，M6-10T: A sharing-delinking paradigm for efficient multitrillion parameter pretraining，2021，CoRR】；【索引31，Scaling vision with sparse mixture of experts，2021，NeurIPS】。如图 2 所示，首先，模型运行一个门控函数，该函数决定每个输入令牌在接下来的 All-to-All 集合通信中的目标 GPU。在第一次 All-to-All（称为 dispatch）之后，每个 GPU 运行自己的专家（一个前馈网络层），然后进行第二次 All-to-All（称为 combine），将每个令牌对应的输出发送回其来源的 GPU。

  • MoE 作为超大规模深度学习的关键。MoE 与现有的扩大 DNN 模型规模的方法（如增加深度或宽度）不同，其成本效益非常高。具体来说，在 MoE 层中增加更多的模型参数（专家）并不会增加处理每个令牌的计算成本。如今，MoE 被认为是超大规模深度学习的一项关键技术，已在多项工作中展示了其最先进的成果【索引7，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】；【索引31，Scaling vision with sparse mixture of experts，2021，NeurIPS】；【索引12，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】；【索引6，Glam: Efficient scaling of language models with mixture-of-experts，2022，ICML】。目前，许多先进的框架（如 DeepSpeed 和 Fairseq）已经支持 MoE。

  • MoE 的动态工作负载。MoE 动态工作负载的根源在于其令牌路由机制。MoE 层动态地将每个令牌路由到多个专家，而令牌在专家之间的分布通常是不均匀的。这导致每个专家的工作负载在每次迭代中都会动态变化，如图 1 所示。专家容量 (Expert Capacity) 是一个常用指标，用于表示每个专家的工作负载，即一个专家需要处理的令牌数量。专家容量取决于每批次的令牌数 $T$、全局专家数 $E$、top-k 路由（$1 \le k \le E$）以及容量因子 $f$（$f \ge 1$），计算公式如下：


    $f=1$ 是最小值，表示最均匀的令牌分布。更大的 $f$ 值表示更不平衡的令牌路由，意味着一个专家必须处理更多的令牌。大多数现有 MoE 框架简单地将 $f$ 设置为一个静态的上限 $f_{upper}$，但这不仅会引入不必要的计算，还可能因 $f_{upper}$ 设置得不够大而丢弃过多的令牌，从而影响模型精度。本文考虑的系统（如 TUTEL）支持使用最小必需的 $f$ 值进行 MoE 训练，既不产生多余计算也不丢弃令牌。

  • 现有 MoE 框架的局限性。GShard【索引12，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】提供了一种保证 MoE 算法正确性的计算逻辑，但许多流行的 MoE 框架【索引27，fairseq: A fast, extensible toolkit for sequence modeling，2019，NAACL-HLT】；【索引30，Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale，2022，ICML】虽然遵循相同的逻辑，但在大规模场景下性能不佳。Fast/FasterMoE【索引8，Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP】提出了与 GShard 在计算上不等价的不同门控算法，其“影子专家”和“智能调度”只在令牌分布持续不平衡时才带来有条件的收益，否则可能损害吞吐量。TUTEL 则旨在保持与 GShard 相同的计算逻辑，并在任何环境下实现确定性的增益。

  • 负载均衡损失（Load Balancing Loss）的权衡。负载均衡（LB）损失通过鼓励门控函数平衡专家工作负载来调节 MoE 层的训练【索引36，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，ICLR】；【索引7，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】。LB 损失有助于降低和稳定 MoE 的工作负载，但通常不足以解决动态工作负载问题，因为过大的 LB 损失权重会损害模型精度。如表 1 所示，大的 LB 损失权重损害了模型准确性。此外，经验发现 LB 损失并不总能带来更平衡的工作负载。本文仅考虑不受 LB 损失影响的通用系统级解决方案。

    表 1. 严苛的负载均衡会损害 MoE 模型精度。粗体数字突出显示了使用大 LB 损失权重时的精度下降。所有实验均在 ImageNet-22K 图像分类上进行，报告的是 SwinV2-S 模型的 top-1 准确率。超参数：32 个专家，top-1 路由，容量因子 f=infinity。

2.2 静态并行

  • 静态并行策略的局限性。在 MoE 层的动态特性下，使用多个 GPU 加速单个专家以提高吞吐量变得具有挑战性。研究表明，当专家数量很多（>256）时，增加更多专家带来的收益会迅速递减。因此，在大规模训练中，通常专家数量相对 GPU 数量较少，一个专家会分配给多个 GPU。本文考虑三种并行方法：专家并行（EP）、数据并行（DP）和模型并行（MP）。实验表明，静态地采用某种并行方法在动态工作负载下并非总是高效。例如，图 3 比较了 EP+DP 和 EP+MP 两种并行方法的性能，显示最佳并行方法取决于工作负载，两种方法之间存在 7.39%-27.76% 的性能差距。


    图 3. 两种不同并行方法的运行时偏好。Y 轴衡量 EP+MP 相对于 EP+DP 的吞吐率。它比较了在不同容量因子 f（即不同工作负载量）和不同 top-k 配置下的吞吐量，其中 > 1.0 意味着 EP+MP 优于 EP+DP，反之亦然。模型设置：fflayer 隐藏层大小 16K，fflayer 通道大小 2048，批处理大小 4。

  • 并行策略切换的开销。不幸的是，在运行时切换不同的并行方法会产生巨大的开销。现有工作中的一种并行训练方式（如数据并行）并未设计成与另一种（如模型并行）兼容，因为它们对数据切分、权重切分、参数梯度动量管理，甚至启动训练的框架接口都有不同的要求。此外，如图 4 所示，改变并行性时还会产生高昂的参数迁移开销。这就是为什么在现有系统中很少使用并行切换的原因。


    图 4. 在传统 EP+DP 和 MP 之间切换并行性导致的参数迁移。Epi 指的是第 i 个专家的第 p 个切片（以模型并行方式）。EP+DP 在每两个 GPU 上复制每个专家，而 MP 将每个专家分别切分到四个 GPU 上。

2.3 静态流水线

  • 静态流水线的效率问题。如图 2 所示的 MoE 层，由于顺序执行 All-to-All 和 fflayer，常常导致 GPU 未被充分利用。由于 All-to-All 主要由非计算密集的 GPU 间数据复制组成，可以通过与运行数值计算的 fflayer 进行流水线操作来更好地利用 GPU 的计算能力。表 2 显示，通过重叠 All-to-All 和 fflayer 计算，可以获得高达 1.86 倍的潜在加速。

    表 2. 在典型 MoE 设置中，All-to-All 开销的比例以及通过完全重叠 All-to-All 和计算所带来的潜在加速。模型设置：fflayer 隐藏层大小 4K，fflayer 通道大小 4K，每个 GPU 2 个专家，每次迭代 64K 令牌。

  • 动态流水线的必要性。我们观察到，用于 dispatch 和 combine 的静态流水线策略（即静态的 All-to-All 算法和流水线深度）在处理动态工作负载时效率低下。如图 5 所示，根据不同的 MoE 设置和规模，相应的最优流水线策略包含多种 All-to-All 算法（Linear 或 2DH）和流水线深度。这意味着单一的静态策略无法在所有情况下都达到最佳性能，因此需要在运行时采用动态流水线策略以适应变化的设置。

  • 计算与通信的干扰。更糟糕的是，计算和通信之间的干扰使得如果只单独考虑一个方面，很难找到最优的流水线策略。这是因为在同一 GPU 上并发运行 NCCL 内核和计算内核所导致的性能下降难以估计。我们的实验表明，即使两种不同的 All-to-All 算法吞吐量相似，当引入相同的并发计算内核时，它们的吞吐量常常会大相径庭，并且任何一种算法都可能在特定情况下胜出。这表明，动态调整应同时针对计算和通信进行，以获得最优的整体吞吐量。


    图 5. 各种 MoE 工作负载配置下最优流水线策略的分布。每一列表示在该 X 轴描述的策略下表现最佳的配置数量。工作负载配置的细节与第 5.1.2 节中描述的相同。

A2 方法细节

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

表 3. 符号说明。

3.1 自适应并行切换

3.1.1 值得进行并行切换的最小子集是什么？

  • 并行策略的复杂度分析。考虑到 EP、DP 和 MP 可以衍生出 7 种可能的并行方法组合，一种临时的方法是为每种方法设计一个执行流，并使其能与其他所有方法切换。然而，这没有必要，因为问题可以精确地简化为一个更小但效率等价的问题。我们的方法是分析所有并行方法的复杂度，将它们缩小到我们需要为其设计执行流的最小子集。这里只关心通信复杂度，因为所有 GPU 执行相同的计算，因此计算复杂度相同。如表 4 所示，我们分析了所有并行方法的通信复杂度，以排除那些（1）在任何情况下都不是最优的或（2）是另一种方法的特例的策略。通过一系列比较，我们得出结论，该子集可以只包含 DP 和 EP+DP+MP。因此，后续只针对这两种策略设计并行结构，这仍能保证覆盖所有模型配置下的最优并行方法。

    表 4. MoE 并行通信复杂度分析。

3.1.2 零成本可切换并行的执行流

  • 零成本切换的设计原则。如 2.2 节所述，可切换的并行性应保证 MoE 训练具有完全相同的数据布局和执行流。零成本意味着切换并行性是完全免费的，不会因参数/令牌迁移引入任何大于 $O(1)$ 的开销。我们分别解释针对 DP 和 EP+DP+MP 的设计。

  • 可切换 DP（Switchable DP）。如图 6 所示，它遵循传统的 DP 训练，只接收本地令牌作为输入，但权重参数遵循 ZeRO-DP Stage-3 分区机制【索引29，Zero: memory optimizations toward training trillion parameter models，2020，SC】。具体来说，它让每个设备拥有一个独特的权重切片，并在前向传播期间执行一次 all-gather 通信，在后向传播期间执行一次 reduce-scatter 通信，而不是传统训练中在后向传播期间执行一次 all-reduce 通信。这两种方式复杂度等价，因为单个 all-reduce 自然地由一个 reduce-scatter 和一个 all-gather 组成。在图 8 中，r = 0 代表可切换 DP。


    图 6. TUTEL 中 DP 执行流的一个例子。All-gather 在所有 (W) 个 GPU 上执行。

  • 可切换 EP+DP+MP (Switchable EP+DP+MP)。如图 7 所示，该并行方法与可切换 DP 工作方式相同——它们共享相同的输入读取格式和权重切分方式。它不仅确保整个计算在数学上与 DP 等价，还确保所需的计算和网络复杂度在表 4 中 ⑦ 所示的 EP+DP+MP 的预期复杂度范围内。我们定义一个控制参数 $r$，它将所有 GPU 划分为一个或多个大小为 $\lceil(W/E)/r\rceil$ 的组，以便在每个组内执行 DP，并在不同组之间执行 MP。具体来说，它在执行流开始时以 MP 风格重复本地令牌 $r$ 次，并最终对称地执行一个局部求和。DP 仅用于在大小为 $\lceil(W/E)/r\rceil$ 的组内执行 all-gather。当 $r$ 增加到 $W/E$ 时，组大小变为 1，因此组内的 all-gather 通信被优化掉。这就是为什么在 ⑦ 中，$r \ge W/E$ 的情况消除了额外的 $O(P/E/r)$。在图 8 中，从 1 到 $W/E$ 的 $r$ 值代表可切换的 EP+DP+MP，其中 $r=1$ 和 $r=\lceil W/E \rceil$ 是两个特例，分别与 EP+DP 和 EP+MP 完全等价。


    图 7. TUTEL 中 EP+DP+MP 执行流的一个例子。Local repeat 生成 r 个门控函数结果的副本，local sum 归约来自 MoE combine 的 r 个输出，all-gather 在 $\lceil(W/E)/r\rceil$ 个 GPU 上执行。


    图 8. 使用 adaptive:r 指定并行方法，max 代表值 $\lceil W/E \rceil$，所有大于此上限的 r 值都被视为与 $\lceil W/E \rceil$ 相同。

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

  • 联合优化流水线深度与 All-to-All 算法。本节介绍自适应流水线的设计。由于 All-to-All 通信延迟对最优流水线深度有显著影响，我们的自适应流水线同时联合优化流水线深度和 All-to-All 算法（Linear 或 2DH）。本节只解释如何为流水线划分输入令牌，而 3.3 节将描述我们如何联合搜索最优的流水线深度和 All-to-All 通信算法。

  • 用于多流流水线的令牌划分。为了实现更细粒度数据块上的流重叠，需要对令牌进行适当划分，以便计算和通信可以在独立的 GPU 流上提交并并行运行。传统的划分方式（如批处理切分或流水线并行）会对层中的所有操作进行划分，但这在 MoE 中行不通，因为它会放大 MoE 分发的不平衡性并破坏诸如批量优先路由 (Batch Prioritized Routing) 等机器学习特性的正确性。我们提出只对两个 All-to-All 操作及其中间的专家计算进行划分，而不是整个 MoE 层。图 9 展示了一个 2-GPU 的例子。

  • 流水线执行流程。在前向传播中，每个 GPU 上的形状为 (E, Cg, D) 的输入沿着维度 C 被分割成两个形状为 (E, Cg/2, D) 的虚拟分区。这两个虚拟分区被标记为 C0 和 C1。分割后，每个虚拟分区 Ci 按顺序在通信流上异步执行 All-to-All 操作。All-to-All 被定制为接受分离的数据块作为输入，并执行内联数据重排，生成形状为 (Eg, C/2, D) 的输出。接着，一旦各自的 All-to-All 完成，两个 All-to-All 的输出被编程发送到计算流上执行专家计算。专家计算的输出又被编程在各自的专家计算完成后发送到通信流上执行第二次 All-to-All。最后，在第二次 All-to-All 之后设置一个屏障，屏障后各分区的输出被合并以生成最终形状为 (E, Cg, D) 的输出。后向传播过程与前向传播类似。所有划分和重塑操作都是通过定制操作内联完成的，因此与非重叠情况相比没有额外的数据复制开销。


    图 9. 在 2 专家 2 GPU 上进行 All-to-All-Expert 多流重叠的令牌划分概览。Ei 表示数据发送到第 i 个 GPU 并由第 i 个专家处理，Ci 表示数据属于容量维度的第 i 个分区。不同容量分区的 All-to-All 和专家操作可以重叠。

3.3 最优并行与流水线字典

  • 构建最优配置字典。TUTEL 维护一个字典，用于记忆不同专家容量范围下的最优并行和流水线设置。该字典被定义为一个哈希映射：$\lfloor c/R \rfloor \rightarrow \{r^*, d^*, a^*\}$，其中 $c$ 是某次迭代的容量值，$R$ 是将多个相邻 $c$ 值收敛到同一个键的窗口大小（默认为 128），而 $\{r^*, d^*, a^*\}$ 是一个最优设置的元组（分别是 adaptive:r、流水线深度和 All-to-All 算法）。为了预先构建这个字典，我们需要为每个可能的键找到最优设置，这只需要进行少量试验，试验次数计算如下：

    其中，$(\log \lceil W/E \rceil + 2)$ 是通过三分搜索【索引39，Ternary Search，2023，Wikipedia】找到 $r^*$ 所需的试验次数，因为 $r$ 在 $[1, \lceil W/E \rceil - 1]$ 范围内决定了一个凸的最优分布，外加两次对 $r=0$ 和 $r=\lceil W/E \rceil$ 的额外试验。"4" 是搜索 $d^*$ 所需的试验次数，因为我们将流水线深度的搜索空间限制为 $\{1, 2, 4, 8\}$。"2" 指的是 All-to-All 算法的数量（Linear 或 2DH）。

4 实现

4.1 特性

  • 功能全面性。与 DeepSpeed MoE、Fairseq MoE 和 FastMoE 等其他 MoE 框架相比，TUTEL 为不同设备、数据类型和 MoE 相关特性提供了更全面的支持。

  • 动态 Top-ANY MoE 门控。为了给 MoE 训练提供多样的稀疏性选项，TUTEL 支持 top-ANY 路由。$k$ 值也可以在每一步进行定制，以实现动态稀疏性更新，这对于一个 MoE 层的不同迭代使用其偏好的 top-k 设置而不是固定的 $k$ 值非常有用。用户可以利用此功能动态微调 MoE 层的稀疏性。

  • 动态容量因子。为了在令牌不平衡变化的情况下智能地控制容量上限，TUTEL 支持在每次迭代中动态调整容量因子。如图 10 所示，调整行为由传递给 MoE 层 API 的参数 capacity_setting = x 控制。如果 x 为正，该值直接作为 MoE 层的容量因子。如果 x 为零，TUTEL 在每次迭代中自动将容量因子调整为不丢弃任何令牌的最小值。如果 x 为负，其工作方式与 x 为零时相同，只是将 -x 设置为容量因子的上限。


    图 10. 当容量设置分别为 4、0 和 -4 时，动态容量因子适应的示例。

4.2 优化

  • 灵活 All-to-All（Flexible All-to-All）。我们提出了一种在传统 MPI/NCCL All-to-All 接口之上的抽象，称为 Flexible All-to-All，以确保 MoE 专家无论在何种规模下都具有高计算吞吐量。现有的 All-to-All 将张量布局从 (E, Cg, D) 转换为 (W, Eg, Cg, D)，其中 Cg 依赖于 W，这会影响后续专家进行矩阵乘法的效率。相反，我们将输出布局转换为 (Eg, C, D)，确保在任何规模 W 下都进行形状相同的矩阵乘法。图 11 比较了传统 All-to-All 和 Flexible All-to-All 下的专家计算吞吐量。


    图 11. 基于 A2A (All-to-All) 布局和 Flexible A2A 布局的专家计算吞吐量。

  • 内核优化：快速编码和解码。根据 GShard 的描述，现有的 MoE dispatch 和 combine 实现需要多次 einsum 和矩阵乘法操作。TUTEL 通过使用 SIMT 高效的稀疏操作（我们称之为快速编码和解码）对此进行了深度优化。它极大地减少了非专家计算的延迟，如图 15 所示。此优化还节省了 GPU 内存，在大多数情况下实现了 20% 到 90% 的内存节省。更多细节见附录 B。

    表 5. 单个 MoE 层的 GPU 内存成本。(静态设置: D = H = 4096, top-k = 2, Eg = 2)

A4 实验

实验环境

  • 硬件配置：
    • 平台：Azure Standard ND96amsr A100 v4 虚拟机。
    • GPU：每台虚拟机配备 8 块 NVIDIA A100 SXM 80GB GPU。最多使用了 2,048 个 A100 GPU（256 台虚拟机）。
    • 互连：节点内 GPU 通过第三代 NVLink 和 NVSwitch 连接。节点间通过 8 x 200 Gbps HDR InfiniBand 非阻塞网络连接，支持自适应路由。
    • CPU：每台虚拟机配备 96 核第二代 AMD Epyc CPU。
    • 内存：1.9 TiB。
  • 软件配置：
    • 框架：PyTorch 1.8.0，基线使用 Fairseq 的 moe 分支。
    • 通信库：NCCL 2.10.3-1，并使用了 NCCL RDMA SHARP 插件。

实验结果

5.1 自适应 MoE 评估

  • 自适应并行切换

    • 实验内容：在单节点（8 GPU）上评估自适应并行切换的性能，测试了两种 MoE 配置（Base 和 Large）在容量因子 $f$ 从 1.0 变化到 8.0 时的吞吐量。
    • 实验结果：如图 12 所示，最优的并行策略（由参数 $r$ 控制）随专家配置和容量因子（即工作负载）的变化而变化。例如，当专家容量高时，DP（$r=0$）倾向于表现更好；随着容量降低，最优策略逐渐转向 EP+DP（$r=1$）乃至 EP+DP+MP（$r>1$）。这种多样性证明了 TUTEL 通过根据动态变化的 $f$ 选择最优并行方法，有巨大的性能提升潜力。

      图 12. 在 Base (左) 和 Large (右) MoE 配置下，不同容量因子 f 的归一化吞吐量。该图显示了最优并行方法随容量因子 f 的变化而不同。

  • 自适应流水线

    • 实验内容：在 16 到 256 个 GPU 规模下，对 243 种典型的 MoE 模型设置评估自适应流水线的性能，并与不同的静态流水线策略（不同深度和 All-to-All 算法）进行比较。同时，使用不同的容量因子 $f$ 来模拟不同训练迭代中的工作负载模式。
    • 实验结果：
      • 与基线（流水线深度为1，Linear All-to-All）相比，自适应流水线平均取得了 9% 到 101% 的提升（表 6a）。
      • 在最坏情况下，自适应流水线能避免性能衰退，并取得 23% 到 599% 的显著提升（表 6b）。
      • 如图 13 所示，自适应流水线总能选择最佳策略，与基线相比，在 $f=4$ 时可实现高达 39% 的提升，在 $f=8$ 时可实现高达 57% 的提升。

        (a) 自适应流水线的平均改进。

        (b) 自适应流水线相对于最差情况的改进。
        表 6. 自适应流水线改进。

        图 13. 自适应流水线根据容量因子 f 带来的改进。D = 4096, H = 4096, Eg = 2, and tokens/step = 4096。
5.2 单个 MoE 层扩展性
  • 实验内容：在扩展到 2,048 个 GPU 的过程中，评估单个 MoE 层的单步时间。以 Fairseq 为基线，逐一添加 TUTEL 的优化特性，分析各部分的增益来源。

  • 实验结果：如图 14 所示，TUTEL 的各项优化技术均带来了显著性能提升。

    • TUTEL 内核 (快速编码/解码)：在小规模（16 GPU）上带来 3.52 倍的大幅增益，大规模（2048 GPU）上增益变小（1.04 倍）。
    • 2DH All-to-All：在大规模（2048 GPU）上带来显著增益（4.25 倍）。
    • 灵活 All-to-All：从 256 GPU 开始在大规模上带来增益（2048 GPU 上为 1.24 倍）。
    • 自适应流水线深度：在 16 GPU 和 2048 GPU 上分别进一步带来 1.43 倍和 1.04 倍的提升。
    • 总性能：与基线相比，TUTEL 最终在 16、128 和 2,048 个 GPU 上分别实现了 4.96 倍、3.11 倍和 5.75 倍的加速。


    图 14. 单个 MoE 层改进分解。基线是 Fairseq / DeepSpeed MoE 层。

    图 15. TUTEL 与 Fairseq / DeepSpeed MoE 之间的内核计算分解比较。

5.3 真实世界问题应用：SwinV2-MoE

  • 模型与任务：引入 SwinV2-MoE 模型，它是 SOTA 视觉模型 Swin Transformer V2 的 MoE 版本。在 ImageNet-22K 数据集上进行预训练，并在 ImageNet-1K 分类、5-shot 线性评估和 COCO 目标检测等下游任务上进行评估。

  • 实验结果：

    • 速度对比 (表 7)：在 8 到 128 个 GPU 的所有规模下，TUTEL 在训练和推理方面均显著快于 Fairseq。训练和推理速度分别提升了 1.14x-1.55x 和 1.95x-2.11x。
    • 精度对比 (表 8)：SwinV2-MoE-B 模型在预训练和下游任务中的准确率均高于其对应的密集模型 SwinV2-B。例如，在 ImageNet-22K 预训练中 top-1 准确率高出 1.3%，在 COCO 目标检测中 box/mask AP 分别高出 0.4/0.4。这首次证明了稀疏 MoE 模型在重要的下游视觉任务（如 COCO 目标检测）上的优越性。


    表 7. SwinV2-MoE 使用 Fairseq 和 TUTEL 的训练和推理速度（图像/秒）对比。
    表 8. 稀疏 SwinV2-MoE-B 模型与其密集对应模型 SwinV2-B 在预训练和微调准确率上的对比。

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

  • 动机：小消息传输问题。大多数流行的深度学习框架使用 NCCL 的点对点（P2P）API 来实现线性 All-to-All 算法。在该算法中，当扩展到大规模（GPU 数量 $n$ 很大）时，任意两个 GPU 之间传输的数据块大小 $S/n$ 会变得很小，这难以充分利用 NVLink 和 HDR InfiniBand 等高速链路的带宽（如图 16 所示）。

    算法 1 使用点对点 API 的线性 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

  • 方法与挑战。为了实现高链路带宽，我们的方法是聚合从多个本地 GPU 发送到同一个远程 GPU 的多个数据块。这通过将小数据块合并为单个大数据块来避免通过网络发送多个小消息，从而显著提高链路带宽利用率。然而，在大规模上高效实现此方法具有挑战性，因为聚合小消息会产生开销。具体来说，在朴素的本地聚合 All-to-All 的第一阶段，节点内 GPU 需要交换非连续的数据块，这会在每个 GPU 上产生 $O(n)$ 的非连续内存访问，导致延迟随 $n$ 增加而增加。

  • 2DH All-to-All 算法。为了避免非连续内存访问导致的性能下降，2DH All-to-All 包含额外的阶段，通过高效的跨步内存复制（stride memory copies）将非连续的数据块对齐到连续的地址空间。如图 17 所示，2DH All-to-All 首先通过跨步内存复制对齐共享相同本地目标 GPU 的数据块（阶段1），然后进行节点内 All-to-All（阶段2）。接下来，再次通过跨步内存复制对齐共享相同远程目标 GPU 的数据块（阶段3），最后进行节点间 All-to-All（阶段4）。通过利用跨步内存复制，2DH All-to-All 在前三个阶段实现了高内存带宽利用率，并保持了与 $n$ 无关的恒定低延迟。


    图 17. 朴素本地聚合 All-to-All 和二维分层 (2DH) All-to-All 各阶段数据布局示例。此例中有两个节点，分别包含 GPU 0~3 和 GPU 4~7。

  • 使用 MSCCL 进行优化。使用 NCCL API 实现 2DH All-to-All 需要在不同阶段之间设置额外的同步屏障，可能导致吞吐量下降。为了获得更好性能，我们利用 MSCCL【索引5，MSCCLang: Microsoft collective communication language，2023，ASPLOS】，通过领域特定语言（DSL）描述 2DH 算法并由编译器进行优化。该定制编译器还利用 LL128 协议进行 All-to-All，这在小尺寸 All-to-All 等低延迟场景中能实现比默认基于 NCCL 的实现更高的效率。

  • 评估。如图 18 所示，与 NCCL 的线性算法相比，所提出的 2DH 算法具有更好的可扩展性，其延迟增长梯度更低。对于小尺寸消息（1 MiB），2DH 算法从较小规模开始就能实现更低的延迟。对于较大尺寸消息，虽然 2DH 算法因额外数据复制导致延迟更高，但随着 GPU 数量的扩展，其性能会变得更好。如图 19 所示，使用定制编译器优化的实现比使用 NCCL API 的实现效果更好，并且 LL128 协议在小尺寸消息上表现更优。因此，在这些算法和协议之间进行动态适配是必要的。

    算法 2 二维分层（2DH）All-to-All

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 高效的快速编码和解码

  • 优化动机与方法。TUTEL 对 MoE 层的编码（在 dispatch 阶段从 MoE 层输入生成 All-to-All 输入）和解码（在 combine 阶段从 All-to-All 输出生成 MoE 层输出）阶段实施了精密的优化。现有实现需要时间复杂度很高的 einsum 操作。如图 20a 所示，该实现包含大量不必要的零乘法和加法。TUTEL 通过稀疏实现来解决此问题（图 20b）。稀疏版本的时间复杂度仅为 $O(T \cdot k \cdot D)$，远低于密集版本的 $O(T \cdot E \cdot Cg \cdot D)$。


    图 20. 从 MoE 层输入 (moe_input) 和门控函数输出 (logits) 生成 All-to-All 分发输入 (dispatch_input) 的密集和稀疏实现比较。

  • GPU 内核实现。尽管稀疏计算无法高效利用矩阵乘法加速器（如 Tensor Cores），但 TUTEL 通过实现可微分的快速编码和解码操作符来解决此问题。这些操作符基于三个特殊设计的 GPU 内核：K0、K1 和 K2（如图 21 所示）。TUTEL 通过将维度 $T$ 的不同索引分配给不同的线程阵列（或 warp）来加速这些内核，确保单个令牌沿维度 $M$ 的计算是 SIMT 高效的。通过这种方法，稀疏计算可以利用 warp shuffling、Blelloch 扫描算法和元素向量化等仅适用于密集计算的优化。最终，TUTEL 极大地减少了编码和解码的延迟（如图 15 所示）并显著节省了 GPU 内存（如表 9 所示，节省 20%~90%）。


    图 21. 快速编码和快速解码操作符的前向和后向计算。括号内为张量形状。张量 X, Y, Z 的形状分别为 (T, D), (T, ), 和 (E, Cg, D)。idxs 和 locations 没有后向计算，因为它们不是可训练的输入。

    表 9. 单个 MoE 层的 GPU 内存成本。(静态设置: D = H = 4096, top-k = 2, Eg = 2)

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

  • COCO 目标检测微调策略。如表 10 所示，直接对 MoE 模型进行微调会导致性能下降。研究发现，在微调期间固定所有 MoE 层可以缓解这个问题，并通过此策略获得了 +0.4/+0.4 的 box/mask AP 提升。

    表 10. COCO 目标检测结果。“fixed” MoE 表示在微调中 MoE 层被固定。

  • 消融研究。

    • 专家数量：表 11 显示，32 和 64 个专家表现最佳，这与先前工作一致。
    • 路由算法和容量因子：图 22 比较了有无批量优先路由（BPR）的路由方法，显示 BPR 对视觉 MoE 模型至关重要，尤其是在容量因子较低时。表 12 对不同 $k$ 和容量因子 $f$ 的性能进行了消融，发现 top-1 路由具有更好的速度-准确性权衡。


    图 22. ImageNet-22K top-1 准确率与推理容量因子的关系。“w/ BPR”表示使用批量优先路由进行训练，而“w/o BPR”则没有。
    表 11. SwinV2-MoE 模型与密集对应模型 (Liu et al., 2022) 的比较。

    表 12. top-k 和容量因子 f 的消融研究。

  • TUTEL 中支持的新余弦路由器。TUTEL 提供了一种新的余弦路由器，旨在通过增加模型大小来提高数值稳定性。其公式为：


    初步实验（表 13）表明，在使用 32 个专家时，余弦路由器在图像分类任务上的准确性与常见的线性路由器相当。尽管目前在图像分类中没有优势，但其对输入的归一化效果可能在输入特征的幅值或维度变化时带来更稳定的路由，并且有同期工作表明余弦路由器在跨语言任务中更准确。

    表 13. 线性路由器和余弦路由器的比较 (E = 32, k = 1, f = 1.25)。