FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training

  • 作者: Yunqi Gao1, Bing Hu1∗, Mahdi Boloursaz Mashhadi2, A-Long Jin3, Yanfeng Zhang4, Pei Xiao2, Rahim Tafazolli2, Mérouane Debbah5
  • 机构: 1浙江大学信息与电子工程学院, 2萨里大学5GIC & 6GIC通信系统研究所(ICS), 3西交利物浦大学未来技术学院, 4东北大学计算机科学与工程学院, 5哈利法大学KU 6G研究中心

A1 主要贡献

本文针对分布式混合专家（MoE）模型训练中的效率瓶颈问题，提出了一个名为 FlowMoE 的可扩展流水线调度框架。

核心问题: 现有用于分布式MoE训练的流水线方法（如ScheMoE, Tutel等）主要关注MoE层内部的任务调度，例如专家计算和all-to-all（A2A）通信，而忽略了其他关键操作，如多头注意力（MHA）计算、门控（gating）以及all-reduce通信。如表1所示，实验表明，MHA计算、门控和all-reduce通信占用了每次迭代时间的30%-40%，这构成了严重的性能瓶颈。

研究目标: 设计一个能够统一调度Transformer块中所有主要任务（包括MHA计算、门控、专家计算、A2A通信和all-reduce通信）的流水线方法，以最大化计算与通信的重叠，从而提高分布式MoE训练的扩展效率。

主要技术贡献:
1. 构建统一流水线：FlowMoE构建了一个统一的流水线，能够持续地调度MHA计算、门控、专家计算和A2A通信，从而将MHA层和门控的计算时间与A2A通信时间重叠。
2. 基于张量块的优先级调度机制：FlowMoE引入了一种基于all-reduce张量块的优先级调度机制，以进一步将all-reduce通信与所有计算任务重叠。该机制将A2A通信任务的优先级设置得高于all-reduce张量块通信任务，从而在不阻塞关键路径的前提下，利用通信间隙执行all-reduce操作。
3. 自适应参数调优：利用贝叶斯优化（BO）自动调整all-reduce张量块的分割大小（Sp），以在最优调度和系统开销之间找到平衡，使框架具有自适应性。
4. 实现与验证：在PyTorch之上实现了FlowMoE框架，并将其开源。通过在两个GPU集群上对675个典型的MoE层和四个真实的MoE模型进行广泛实验，证明了FlowMoE相比现有最先进的MoE训练框架，能够将训练时间减少13%-57%，能耗降低10%-39%，内存使用减少7%-32%。

Table 1: 在一个16-GPU (NVIDIA RTX3090) 集群（100Gbps带宽）上运行原生专家并行[19]训练四个MoE模型时，每次迭代中不同任务的时间。“MHA + Gating Time”表示MHA层和门控函数的计算时间。“All-Reduce Time”表示all-reduce通信的时间。“Ratio”表示“MHA + Gating Time”和“All-Reduce Time”之和占每次迭代总时间的比例。

A3 背景知识与挑战

2.1 带有MoE层的Transformer块

Transformer块结构。典型的带有MoE层的Transformer结构如图1a所示，其中Transformer块通常由一个MHA层和一个MoE层组成。对于输入张量$I \in R^{B \times N \times M}$，MHA层利用查询、键、值矩阵$W^Q, W^K, W^V \in R^{M \times M}$来计算每个令牌的注意力数据，并通过一个线性变换矩阵$W^O \in R^{M \times M}$获得输出张量$I' \in R^{B \times N \times M}$，其中B表示每个GPU每次迭代的样本数（或小批量大小），N表示每个样本的令牌数，M表示令牌的嵌入大小。MoE层包括一个门控函数和多个专家。门控函数G是一个小型的可学习神经网络，后跟一个softmax层，用于为输入令牌选择激活的专家。通常，只有top-k个专家被选择来处理一个令牌【索引9，Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity，2022，Journal of Machine Learning Research】。门控函数的输出张量$G(I') \in R^{E \times C \times M}$将被分派到相应的专家（每个专家接收一个形状为$C \times M$的张量），其中E表示每个MoE层的专家总数，C表示分配给一个专家的最大令牌数。C可以通过$f \times k \times B \times N/E$计算，其中f是容量因子，用于确定分配给每个专家的最大令牌数并调整C【索引5，GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding，2020，arXiv】。专家通常是一个具有两个结构对称的前馈层（第一层和第二层的大小分别为$M \times H$和$H \times M$）的小型神经网络，其中H表示前馈层的隐藏大小，每个专家被认为拥有自己的专业领域。专家计算后，所有专家的输出被组合成一个形状为$B \times N \times M$的张量，作为下一个Transformer块的输入。

(a) 带有MoE层的Transformer结构。(b) 专家并行示意图。

Figure 1: 一个带有MoE层的Transformer块和专家并行的示例。

2.2 专家并行

专家并行机制。我们将P定义为集群中的工作节点（或GPU）数量，L为MoE模型中的Transformer块数量。图1b展示了一个通过专家并行在P个工作节点上训练MoE模型的示例，其中每个工作节点拥有不同的专家参数。然后，$AT_r^{(l)}, E_r^{(l)}, D_r^{(l)}$ 和 $C_r^{(l)}$ 分别表示第l个Transformer块的MHA层（包括门控函数）、专家计算、分发A2A通信和组合A2A通信的第r个子任务。$AR^{(l)}$表示第l个Transformer块的all-reduce通信任务。图2a和2b展示了在原生专家并行【索引19，FastMoE: A Fast Mixture-of-Expert Training System，2021，arXiv】中，一个Transformer块在前向计算和反向传播过程中的多个计算和通信任务的时间线。值得注意的是，反向传播的时间线与前向计算相反，但MHA层和门控函数的参数需要通过额外的all-reduce通信操作进行同步。

2.3 分布式MoE训练中的性能瓶颈

现有流水线方法的局限性。大多数现有的流水线调度工作（例如，ScheMoE 【索引10，ScheMoE: An Extensible Mixture-of-Experts Distributed Training System with Tasks Scheduling，2024，EuroSys】，Tutel 【索引12，Tutel: Adaptive Mixture-of-Experts at Scale，2023，MLSys】，PipeMoE 【索引21，PipeMoE: Accelerating Mixture-of-Experts through Adaptive Pipelining，2023，INFOCOM】）根据流水线程度R，在数据维度上划分MoE层的输入令牌张量，其中图2c和2d展示了一个在Transformer块的MoE层中R=2的流水线示例。此外，FasterMoE 【索引11，FasterMoE: Modeling and Optimizing Training of Large-Scale Dynamic Pre-Trained Models，2022，PPoPP】根据工作节点数量拆分MoE层的输入张量，实现了工作节点之间的点对点通信。尽管这些工作通过流水线化专家计算任务和A2A通信任务有效减少了训练时间，但它们忽略了MHA层计算任务、门控任务和all-reduce通信任务。此外，一些MoE框架（例如，FSMoE 【索引24，FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models，2025，ASPLOS】和Lina 【索引20，Accelerating Distributed MoE Training and Inference with Lina，2023，USENIX ATC】）也探索了all-reduce通信任务的流水线。然而，FSMoE更侧重于MoE层内的节点间和节点内通信重叠，而Lina仅优化了MoE层特定的通信瓶颈，而不是整个Transformer块。因此，我们的目标是流水线化整个Transformer块中的所有主要计算和通信任务，以最小化每次迭代的时间，这涉及三个主要挑战：
1. 多类型任务间的复杂依赖关系。在分布式MoE训练中，计算和通信任务之间存在高度复杂的依赖关系（见图2a和2b）【索引10，ScheMoE: An Extensible Mixture-of-Experts Distributed Training System with Tasks Scheduling，2024，EuroSys】。特别是，典型的并行化方案（例如，流水线并行【索引25, 26】，张量并行【索引27, 28, 29】）难以直接叠加在专家并行之上。
2. A2A通信与all-reduce通信的共存。尽管已有大量研究【索引30, 16, 31, 15, 32, 33, 34】为数据并行训练的all-reduce通信任务提供了高效的调度算法，但它们主要关注传统的基于卷积的深度神经网络或没有MoE层的LLM。这些研究不能直接应用于分布式MoE训练，因为额外的A2A通信任务不等同于all-reduce通信任务（见图2b）。
3. 设计自适应和通用的流水线调度框架。目前，大多数调度框架的性能依赖于一些超参数的设置【索引35, 36, 37】，这在训练前引入了额外的调优工作量。理想情况下，一个自适应框架应该能自动调整超参数，并且只需修改模型定义或数据集接口即可直接部署。同时，通用框架也应设计为与不同的优化框架和通信栈兼容。

本文的主要目标是通过提出一个名为FlowMoE的可扩展流水线调度框架来解决以上三个挑战，以应对分布式MoE训练中的多类型任务。


Figure 2: 一个由两个带有MoE层的Transformer块组成的模型的计算和通信任务执行时间线示例。

A2 方法细节

3.1 概述

FlowMoE 工作流程。图3展示了FlowMoE的工作流程。FlowMoE利用给定的MoE结构和数据集进行高效的流水线训练。首先，多类型计算和通信任务根据流水线程度被分解，等待流水线调度。其次，FlowMoE根据它们的依赖关系和定义的优先级，在GPU集群上为分布式训练调度所有计算和通信任务（详见3.2和3.3节）。第三，FlowMoE在训练过程中，根据前几次迭代通过BO分析自动调整all-reduce张量的分区大小（详见4.1节）。


Figure 3: FlowMoE的工作流程。

3.2 MHA和MoE层的流水线

统一的流水线设计。基于MoE层的流水线，FlowMoE首先对MHA层计算和门控进行流水线化。值得注意的是，对于第l个Transformer块，门控任务仅依赖于MHA层的计算任务，在后续表达式中，我们将两者的计算任务视为一个整体，并用$AT^{(l)}$表示（$1 \le l \le L$）。FlowMoE将每个Transformer块的输入张量划分为R个大小相等的部分，MHA层和MoE层中的每个计算或通信任务（all-reduce通信任务除外）将被划分为R个独立的子任务。特别是，相同类型的任务具有相同的执行时间。划分后的任务可以表示为集合：

$$\mathbb{T} = \left\{ AT_r^{(l)}, D_r^{(l)}, E_r^{(l)}, C_r^{(l)}, AR^{(l)} | 1 \le r \le R \right\},$$

其中$AT_r^{(l)}$和$E_r^{(l)}$是计算任务，而$D_r^{(l)}$、$C_r^{(l)}$和$AR^{(l)}$是通信任务。

前向计算的调度顺序。对于$1 \le l < L$，在前向计算期间，计算任务的调度顺序可以表示为：

$$AT_{1}^{(l)} \rightarrow AT_{2}^{(l)} \rightarrow \ldots \rightarrow AT_{R}^{(l)} \rightarrow E_{1}^{(l)} \rightarrow E_{2}^{(l)} \rightarrow \ldots \rightarrow E_{R}^{(l)} \rightarrow AT_{1}^{(l+1)} \rightarrow \ldots \rightarrow E_{R}^{(l+1)}$$

A2A通信任务的调度顺序可以表示为：

$$D_{1}^{(l)}->D_{2}^{(l)}->...->D_{R}^{(l)}->C_{1}^{(l)}->C_{2}^{(l)}->...->C_{R}^{(l)}->D_{1}^{(l+1)}->...->C_{R}^{(l+1)}$$

反向传播的调度顺序。在反向传播期间，计算任务的调度顺序可以表示为：

$$E_{R}^{(l+1)}->...->A T_{1}^{(l+1)}->E_{R}^{(l)}->E_{R-1}^{(l)}->...->E_{1}^{(l)}->A T_{R}^{(l)}->A T_{R-1}^{(l)}->...->A T_{1}^{(l)}$$

A2A通信任务的调度顺序可以表示为：

$$C_{R}^{(l+1)} \rightarrow \ldots \rightarrow D_{1}^{(l+1)} \rightarrow C_{R}^{(l)} \rightarrow C_{R-1}^{(l)} \rightarrow \ldots \rightarrow C_{1}^{(l)} \rightarrow D_{R}^{(l)} \rightarrow D_{R-1}^{(l)} \rightarrow \ldots \rightarrow D_{1}^{(l)} .$$

流水线效果。图2e和2f展示了在MHA层和MoE层中使用R=2进行流水线化时，计算和通信任务的执行时间线示例。与仅对MoE层进行流水线化相比，MHA层计算任务和门控任务可以与A2A通信任务重叠，从而缩短了每次迭代的时间。

3.3 All-reduce通信的流水线

All-reduce 通信的调度挑战。现有的最先进调度框架在每次迭代的反向传播结束时集中执行all-reduce通信任务（我们称之为all-reduce通信任务的集中式调度）。为了减少all-reduce通信时间，我们进一步考虑了all-reduce通信任务的流水线。直观上，在反向传播中，Transformer块l的all-reduce通信任务可以与Transformer块l-1的计算任务重叠，因为它们都依赖于Transformer块l的MHA层计算任务的完成。然而，在实际训练中，Transformer块l的all-reduce通信任务会与Transformer块l-1的A2A通信任务发生冲突。因此，我们首先对一次迭代中反向传播的时间线进行数学建模，以找到最小化训练时间的两种通信任务的最优调度。

数学建模。考虑到真实训练环境中的资源竞争，我们假设GPU上只能同时执行计算和通信任务，而不能同时运行多个计算或多个通信任务。此外，任务之间没有抢占：一旦任务开始执行，它必须不间断地运行至完成。我们将$\tau_b(\cdot)$定义为反向传播期间任务的开始执行时间戳，将$t_b(\cdot)$定义为反向传播期间任务的耗时。根据图2b，目标函数和任务间的依赖关系可以表示如下（$1 \le r \le R$）：

$$\min \quad T_b = \tau_b(AR^{(1)}) + t_b(AR^{(1)}) - \tau_b(C_R^{(L)})$$ $$\text{s.t. } \tau_{b}(C_{r}^{(l-1)}) \geq \tau_{b}(AT_{r}^{(l)}) + t_{b}(AT_{r}^{(l)}), 1 < l \leq L,$$ $$\tau_{b}(E_{r}^{(l)}) \geq \tau_{b}(C_{r}^{(l)})+t_{b}(C_{r}^{(l)}), 1 \leq l \leq L,$$ $$\tau_{b}(D_{r}^{(l)}) \geq \tau_{b}(E_{r}^{(l)})+t_{b}(E_{r}^{(l)}), 1 \leq l \leq L,$$ $$\tau_{b}(AT_{r}^{(l)}) \geq \tau_{b}(D_{r}^{(l)})+t_{b}(D_{r}^{(l)}), 1 \leq l \leq L,$$ $$\tau_{b}(A R^{(l)}) \geq \tau_{b}(A T_{r}^{(l)})+t_{b}(A T_{r}^{(l)}), 1 \leq l \leq L.$$

理论分析。我们使用符号来表示all-reduce通信任务集中调度的时间线。现在，令$T_b$表示当一个Transformer块的all-reduce通信任务插入到任意两个A2A通信任务之间时的反向传播时间。并用$T_b^*$表示all-reduce通信任务集中调度下的反向传播时间。基于此，我们有以下定理：
定理 1*。如果调度顺序满足公式4和5，则有$T_b \le T_b^*$。
证明在附录B中提供。

基于 All-reduce 张量块的优先级调度机制。根据定理1，我们发现将all-reduce通信任务划分并插入到任何A2A通信任务之间的间隙中，可以增加all-reduce通信任务与计算任务之间的重叠。因此，我们在FlowMoE中设计了一种基于all-reduce张量块的通信任务优先级调度机制。具体来说，在反向传播期间，FlowMoE首先将每层的all-reduce通信任务张量切片成大小为$S_p$的张量块。其次，FlowMoE维护一个通信任务池，该池包含所有的all-reduce张量块和A2A通信任务，并将all-reduce张量块的调度优先级设置为低于A2A通信任务。换句话说，当没有A2A通信任务时，all-reduce张量块的通信任务将立即执行。图2f展示了在MHA层和MoE层中使用2级流水线以及基于all-reduce张量块的通信任务优先级调度机制时的执行时间线示例。可以观察到，all-reduce张量块充分占用了A2A通信任务之间的间隙，从而最大化了计算任务与通信任务之间的重叠。

4.1 通过贝叶斯优化自动调整分区大小

理论最优与实际开销。理想情况下，如果all-reduce张量块的通信不引入额外的启动开销，我们有以下定理：
定理 2。当调度顺序满足公式4和5，并使用基于all-reduce张量块的通信任务优先级调度机制时，如果$S_p \rightarrow 0$且通信all-reduce张量块没有启动开销，则每次迭代的时间将最小化。
证明在附录C中提供。

在实际训练中，all-reduce张量块的通信任务会引入额外的启动开销【索引15，US-Byte: An Efficient Communication Framework for Scheduling Unequal-Sized Tensor Blocks in Distributed Deep Learning，2023，IEEE Transactions on Parallel and Distributed Systems】。因此，all-reduce张量块的分区大小$S_p$成为决定最优调度与系统开销之间权衡的关键。由于将迭代时间明确建模为$S_p$的函数非常困难，为了保证FlowMoE的自适应性，我们采用贝叶斯优化（BO）在训练期间自动调整$S_p$，它试图在尽可能少的试验次数内找到一个未知目标函数的良好参数。BO在FlowMoE中的目标是最小化每次迭代的时间。具体来说，BO通过采样不同的($S_p$，每次迭代时间)对来拟合一个目标函数，并不断建议下一个$S_p$以获得新的目标函数值，其中所有对中不同$S_p$对应的每次迭代时间是通过多次迭代（例如10次）的平均值来测量的。换句话说，BO可以基于足够的样本准确预测一个接近最优的$S_p$。图4展示了使用BO调整$S_p$的一个示例。在实际训练中，仅用8个样本，BO就能以很高的置信度返回一个接近最优的值2.5MB。


Figure 4: BO示例：在16-GPU集群上为训练BERT-Large-MoE调整Sp。

4.2 算法设计

算法流程。算法1正式给出了使用FlowMoE进行R级计算和通信任务的流水线调度训练过程。此外，算法2提供了通信池管理，包括all-reduce张量的拆分和通信任务的优先级调度。两个算法的详细描述见附录E。

Algorithm 1: FlowMoE Pipeline Scheduling
Input: Dataset D = {(x1, y1), ..., (xn, yn)}, L.
Output: Model Weights W = {W1, ..., WL}.
1: Initialize DataQueue for data transferred between tasks;
2: Initialize A2AQueue for A2A communication tasks;
3: Initialize ARQueue for all-reduce communication tasks;
4: Split D for d1, d2, ..., dR into DataQueue;
5: Communication_pool_management.start();
6: for l = 1 → L do // Feed-forward Computing
7:   for r = 1 → R do
8:     dr = AT.FFcomp(DataQueue.get());
9:     A2AQueue.put(dr);
10:  for r = 1 → R do
11:    dr = E.FFcomp(DataQueue.get());
12:    A2AQueue.put(dr);
13: for l = L → 1 do // Backward Propagation
14:   for r = 1 → R do
15:     A2AQueue.put(DataQueue.get());
16:  for r = 1 → R do
17:    dr = E.BPcomp(DataQueue.get());
18:    A2AQueue.put(dr);
19:  for r = 1 → R do
20:    dr = AT.BPcomp(DataQueue.get());
21:    DataQueue.put(dr);

Algorithm 2: Communication Pool Management
1: /* Partition an all-reduce tensor and enqueue tensor chunks. */
2: Get Sp from BO;
3: procedure PARTITION(ARTensor)
4:   ARChunks = ARTensor.partition(Sp);
5:   ARQueue.put(ARChunks);
6: /* Communication task priority scheduling. */
7: procedure COMMPOOLMANAGER
8:   while True do
9:     if A2AQueue is not empty then
10:      d = A2A.Comm(A2AQueue.get());
11:      DataQueue.put(d);
12:    else if ARQueue is not empty then
13:      AR.Comm(ARQueue.get());

开销分析。与原生专家并行【索引19，FastMoE: A Fast Mixture-of-Expert Training System，2021，arXiv】相比，FlowMoE中流水线调度的开销主要来自两部分：all-reduce张量的划分和BO期间的目标函数拟合。首先，当门控函数结构为典型的$M \times E$线性层时，MHA层和门控函数的总参数量为$4M^2 + M \times E$。根据算法2，FlowMoE调度算法在每次迭代中的计算复杂度为$O(L \times \frac{4M^2+M \times E}{S_p})$，这是可以接受的，因为它在实际实验中不到一次迭代时间的1%，并且all-reduce张量划分的时间可以与某些计算/通信任务的时间进一步重叠（我们在附录E中讨论其可扩展性）。其次，BO利用训练过程的前几次迭代来自动将$S_p$调整到接近最优的值。在我们的实验中，BO采样8个$S_p$值，并通过平均10次迭代记录每个值对应的迭代时间。BO引入的计算开销与总训练时间相比可以忽略不计，具体量化见附录D。

4.3 系统实现

部署与架构。我们在PyTorch中利用其API部署FlowMoE，它通常支持灵活的Python语言类继承，并具有很高的通用性，可以与不同的优化框架和通信栈兼容。特别是，我们在Tutel【索引12，Tutel: Adaptive Mixture-of-Experts at Scale，2023，MLSys】之上实现FlowMoE，这是一个高度优化的MoE加速库，深度集成到PyTorch中，并支持通信和计算任务的异步执行。Tutel也已被DeepSpeed【索引38，DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale，2022，ICML】用作默认的MoE训练模块。图5展示了FlowMoE架构的概览（左侧更接近用户层），其详细描述和技术实现可在附录F中找到。


Figure 5: FlowMoE架构概览（深灰色部分是新增的）。

A4 实验环境与结果

实验环境

  • 硬件配置:

    • 集群1: 2个节点，通过100Gb/s带宽连接。每个节点配备8块NVIDIA RTX3090 GPU（每块24GB显存），通过PCIe3.0x16连接。CPU为Intel Xeon(R) Gold 6248R。
    • 集群2: 4个节点，通过10Gb/s带宽连接。每个节点配备2块NVIDIA RTX2080Ti GPU（每块12GB显存），通过PCIe交换机连接。CPU为Intel Xeon(R) Gold 5118。

  • 模型与数据集:

    • 自定义MoE层: 覆盖了多种典型配置，参数组合包括$B \in \{2, 4, 8\}, f \in \{1.0, 1.1, 1.2\}, N \in \{512, 1024, 2048\}, M \in \{512, 1024, 2048, 4096, 8192\}, H \in \{512, 1024, 2048, 4096, 8192\}$。专家数量E等于GPU数量P，k=2。
    • 真实MoE模型: 选择了四个流行模型，包括用于语言建模任务的GPT2-Tiny-MoE和DeepSeek-V2（在OpenWebText数据集上），以及用于文本生成任务的BERT-Large-MoE和LLaMA2-MoE（在wikitext-103数据集上）。这些模型是将原始模型（GPT2-Tiny, BERT-Large, LLaMA2）中的所有前馈层替换为MoE层构建的。详细配置见下表2。所有模型参数和梯度均使用32位单精度浮点数存储。

  • 软件配置:

    • 实现: 基于PyTorch，并建立在Tutel库之上。
    • 基线: PyTorch-based vanilla expert parallelism (vanillaEP), ScheMoE, FSMoE, Tutel, FasterMoE。除非特别说明，流水线程度R=2。

实验结果

  • 真实MoE模型的端到端时间: 如表3所示，在不同规模的集群和模型上，FlowMoE均获得了最佳的可扩展性和最短的单次迭代时间。具体而言，FlowMoE分别比ScheMoE快14%-31%，比FSMoE快13%-25%，比Tutel快29%-42%，比FasterMoE快26%-57%，比vanillaEP快43%-82%。在带宽受限的情况下，随着GPU数量增加，流水线调度的优势减弱。

Table 3: 平均每次迭代时间的比较（毫秒）。S1、S2、S3、S4和S5分别是FlowMoE相对于ScheMoE、FSMoE、Tutel、FasterMoE和vanillaEP的加速比。

  • 流水线程度与训练速度: 表4展示了在16-GPU的集群1上，不同流水线程度R下的平均迭代时间。FlowMoE在所有情况下都优于ScheMoE和Tutel。

Table 4: 在DeepSeek-V2-S上不同流水线程度下的平均每次迭代时间。S1和S2分别是FlowMoE相对于Tutel和ScheMoE的加速比。

  • 自定义MoE层的加速效果: 在集群1和集群2上分别测试了490个和393个有效的自定义MoE层配置。如图6所示，FlowMoE在所有情况下都比ScheMoE快，平均实现了26%的性能提升。


Figure 6: 相对于ScheMoE的加速比统计。

  • 消融实验: 表5展示了在自定义MoE层上各组件的贡献。仅流水线化MHA层和门控（FlowMoE-AT）相比Tutel提升了10.3%。进一步加入all-reduce流水线（FlowMoE-AR）和BO自动调优（FlowMoE-AR(BO)），性能分别再提升了14.3%和8.3%。最终，FlowMoE相比Tutel和vanillaEP分别快1.4倍和2.05倍。

Table 5: 在不同组件下MoE层的每次迭代时间。报告的加速比值以vanillaEP为基线。“w/ Pipe-MoE”表示对专家计算和A2A通信进行流水线处理。“w/ Pipe-AT”表示对MHA层计算和门控进行流水线处理。“w/ Pipe-AR”表示对all-reduce通信进行流水线处理。

  • 能耗与内存使用: 如表6所示，FlowMoE通过最大化计算与通信重叠，提高了资源利用率，从而降低了能耗。与ScheMoE、Tutel、FasterMoE和vanillaEP相比，FlowMoE分别节省了10%-16%、22%-27%、33%-39%和33%-41%的能耗。在内存使用方面，FlowMoE通过及时执行all-reduce通信减少了梯度缓存，内存使用量最低，分别比ScheMoE、Tutel、FasterMoE和vanillaEP节省高达7%、9%、32%和11%的内存。

Table 6: 一次迭代中平均每个工作节点的能耗和内存使用情况。

A5 结论

本文提出了一个名为FlowMoE的可扩展、经过数学证明的流水线调度框架，用于加速MoE模型的训练。FlowMoE解决了所有主要MoE相关任务的统一调度问题，并实现了异构通信任务（A2A和all-reduce）在MoE训练中的最优共存，这极大地推进了分布式解决方案，超越了传统流水线方法的简单扩展。我们在PyTorch上实现了FlowMoE框架，并在两个GPU集群上使用675个典型MoE层和四个真实世界的NLP模型进行了广泛的实验。实验结果表明，FlowMoE在训练时间、能耗和内存使用方面均优于包括ScheMoE、FSMoE、Tutel和FasterMoE在内的最先进MoE训练框架，性能提升分别为13%-57%、10%-39%和7%-32%。此外，FlowMoE可以与许多正交的优化工作相结合，将分布式流水线优化推向一个新的阶段。

A6 附录

A FlowMoE与关键文献的主要区别

对比分析。表A.2总结了FlowMoE与分布式MoE训练领域关键流水线调度文献的主要区别。FlowMoE是唯一一个同时对A2A通信、专家计算、MHA/门控计算和All-reduce通信进行流水线化，并提供自动调优和对动态硬件环境具有鲁棒性的框架。

Table A.2: FlowMoE与分布式MoE训练流水线调度的关键文献之间的主要区别，其中BO代表贝叶斯优化。

B 定理1的证明

证明思路。由于每个Transformer块的任务调度时间线相同，且all-reduce通信时间也相同，我们只需证明将$AR^{(l+1)}$插入到Transformer块l的任意两个A2A通信任务之间时的反向传播时间$T_b$将小于或等于集中式调度下的$T_b^*$。通过分析四种可能的插入位置（在$C^{(l)}$任务之间、在$C^{(l)}$和$D^{(l)}$之间、在$D^{(l)}$任务之间、在$D^{(l)}$和$C^{(l-1)}$之间），可以推导出，插入操作通过利用计算和A2A通信之间的空闲时间来执行all-reduce任务，从而提前了后续任务的开始时间，最终使得总的反向传播时间$T_b$不长于$T_b^*$。图A.1以图形化方式展示了这一过程，说明了插入调度如何优化时间线。


Figure A.1: 一个示例的演示，其中(a)、(c)和(e)代表使用集中式调度all-reduce通信任务时的三个时间线（虚线表示前后任务没有直接依赖关系），(b)、(d)和(f)代表使用调度顺序1时的三个时间线。比较(a)与(b)，(c)与(d)以及(e)与(f)，可以观察到$T_b \le T_b^*$。

C 定理2的证明

证明思路。在反向传播期间，如果不考虑all-reduce通信任务，通信资源的空闲时间（即A2A通信任务之间的间隙）是固定的。最小化一次迭代时间的方法是最大化all-reduce张量块的通信任务对这些空闲时间的占用。如果all-reduce张量块的通信没有额外的启动开销，那么对于不同的$S_p$，all-reduce通信任务的总时间将保持不变。当$S_p$较大时，未完成的all-reduce块通信任务可能会导致A2A通信任务无法及时执行，这将根据公式6b和6d使计算任务的时间线向后推移，并且部分all-reduce块的通信任务没有占用通信资源的空闲时间。相反，当$S_p \to 0$时，由于A2A通信任务的优先级更高，all-reduce块的通信任务不会影响A2A通信任务的时间线，因此也不会改变计算任务的时间线。同时，这将最大化地利用通信资源的空闲时间。因此，证明完成。

D BO的设计与性能评估

D.1 参数设置与优势分析

BO选择原因。我们选择BO来获取接近最优的$S_p$得益于以下三点：
1. 不依赖目标函数表达式：BO仅依赖于采样值，我们使用高斯过程回归与Matern核作为代理模型。
2. 搜索开销低：BO通过最大化采集函数（本文采用预期提升EI）来选择下一个配置$S_p$，通常只需少量试验即可找到高质量解。
3. 避免局部最优：通过调整超参数EI（本文设为0.1以偏向探索），BO可以避免陷入局部最优。$S_p$的搜索空间设置为(0MB, 每层Transformer块中待通信张量的最大值]，初始样本随机生成，能有效覆盖不同硬件和模型下的最优$S_p$。

D.2 BO的重要性分析

BO的必要性。BO自动调优是FlowMoE不可或缺的一部分。过大的all-reduce张量块会影响A2A通信任务的优先级，使其无法及时启动。而过小的张量块会导致过多的启动开销。因此，必须存在一个唯一的优解来最大化训练速度。BO优化器通过从实际迭代中采样和学习来平衡这两个相互竞争的效应，高效地找到一个既能最大化重叠又不产生过多开销的all-reduce张量块大小。

D.3 BO的性能评估

性能对比。如表A.3所示，与网格搜索和随机数生成相比，使用BO调整$S_p$获得了最短的单次迭代时间。网格搜索受限于采样点数量，难以覆盖最优解。如表A.4所示，不同固定分区大小对训练效率影响巨大，BO自动调优对于最大化FlowMoE性能至关重要。

Table A.3: 使用不同方法调整分区大小Sp时，平均每次迭代时间的比较（毫秒）。

Table A.4: 使用带BO自动调优的FlowMoE或不同固定分区大小时的平均每次迭代时间比较（毫秒）。

超参数敏感性。对于单峰、平滑的目标函数，BO寻找最优解的过程对采集函数和代理模型这两个超参数不敏感。如表A.5所示，不同的BO配置导致了相似的迭代时间，证明了其鲁棒性。

Table A.5: 在集群1上训练BERT-Large-MoE时，不同BO参数配置下的平均每次迭代时间比较（毫秒），其中代理模型使用不同核函数的高斯过程回归（GPR）。

计算开销。如表A.6所示，BO的计算开销占前1000次迭代训练时间的比例非常小（0.16%至3.22%），与其带来的性能增益相比可以忽略不计，证明了BO的轻量级和实用性。

Table A.6: BO的计算开销占前1000次迭代训练时间的百分比。

E 算法描述与可扩展性分析

算法描述。算法1描述了训练过程中R级计算和通信任务的流水线。第1-4行初始化参数和队列。第6-12行根据公式2和3执行每次迭代的前向计算。第13-21行根据公式4和5执行每次迭代的反向传播。算法2展示了通信池的管理。PARTITION过程负责在反向传播过程中拆分MHA层和门控函数的张量，并将其放入all-reduce通信任务队列。COMMPOOLMANAGER过程根据定义的优先级执行两类通信任务，其中第9-11行优先执行A2A通信任务，而第12-13行在没有A2A任务时执行all-reduce张量块的通信任务。

可扩展性分析。一次迭代的时间复杂度为$O(L \times (\frac{BNM^2 + B^2N^2M}{P} + \frac{BNME}{P} + 2BNMH))$。FlowMoE调度算法的复杂度为$O(L \times \frac{4M^2+M \times E}{S_p})$。当模型复杂度（L, B, N, M, H, E）或GPU数量P增加时，FlowMoE调度算法开销的增长远小于单次迭代时间的增长。因此，调度算法开销占总迭代时间的比例会下降，表明FlowMoE的调度算法具有良好的可扩展性。

F 系统描述

系统实现。FlowMoE部署在PyTorch API层面，以保证通用性，不修改用户脚本和框架引擎。其四个模块如下：
* 任务分解管理器（Task Breakdown Manager）：基于Tutel改进，负责根据流水线程度R拆分数据集，并根据任务间传输的张量请求不同的通信/计算子任务。
* BO自动调优器（BO autotuner）：负责搜索接近最优的分区大小，并指导all-reduce张量的划分。
* 通信任务池（Communication Task Pool）：维护A2A通信任务和all-reduce块通信任务的队列。
* 流水线调度管理器（Pipeline Scheduling Manager）：根据定义的调度顺序和优先级，将请求的多个计算和通信任务提交给PyTorch引擎和通信库。

实现细节。为了在反向传播期间及时获取每层Transformer块的梯度张量，我们使用register_full_backward_hook来访问梯度并进行拆分。通过多线程实现计算和通信任务的重叠，主线程负责调度所有任务，一个子线程根据优先级调度通信池中的任务，并使用threading.Lock保证线程安全。

G 两个放大MoE模型的压力测试

压力测试结果。为了对FlowMoE的性能进行压力测试，我们在两个接近集群1内存上限的放大MoE模型（LLaMA2-MoE-L和DeepSeek-V2-M）上测量了不同训练框架的性能。如表A.7所示，即使在更大模型上，FlowMoE仍然在所有基线中取得了最佳的训练性能。

Table A.7: 训练两个放大版MoE模型时的平均每次迭代时间比较。S1、S2和S3分别是FlowMoE相对于ScheMoE、Tutel和vanillaEP的加速比。

H 收敛性分析与实验

理论分析。我们将每个流水线中处理的样本称为一个微批次（microbatch），微批次数量等于流水线程度R。在反向传播期间，来自所有$AT_{lr}$和$E_{lr}$（$1 \le r \le R$）的梯度被累加求和。只有当来自$AT_{l1}$的梯度也被累加后，才启动第l个Transformer块中MHA和门控函数的All-Reduce块通信任务。类似地，当$E_{l1}$的梯度被累加后，才更新专家参数。这有效防止了参数被提前更新，避免了梯度陈旧。此外，为了确保流水线内外的累积梯度等价，我们将每个微批次计算的损失按R进行缩放，即$\frac{loss^{(r)}}{R}$。数学推导表明，通过这种方式，所有微批次的累积梯度在数值上与不使用流水线时使用整个小批次（mini-batch）计算的梯度是等价的。

实验验证。我们在集群1上使用16个GPU训练GPT2-Tiny-MoE和BERT-Large-MoE来实验验证FlowMoE的收敛性。如图A.2所示，FlowMoE达到了与基线相同的损失，但由于单次迭代时间更短，花费的时间要少得多。


Figure A.2: 训练GPT2-Tiny-MoE和BERT-Large-MoE时的损失随时间变化图。

I 性能下限分析

性能分析。我们分析了FlowMoE在三种情况下的性能下限：
1. 通信任务时间远长于计算任务时间：此时，仅A2A通信任务的时间就完全覆盖了计算任务时间。FlowMoE的性能将与ScheMoE、Tutel和FasterMoE相同，但优于vanillaEP。
2. 计算任务时间远长于通信任务时间：此时，所有通信任务时间都可以被计算任务时间覆盖。FlowMoE将优于ScheMoE、Tutel和FasterMoE，因为它隐藏了AR任务时间。
3. 通信任务时间与计算任务时间相当：此时，FlowMoE通过最大化多类型任务的重叠而优于所有基线。

总结来说，在所有情况下，FlowMoE的性能都大于或等于ScheMoE、Tutel和FasterMoE，并且总是优于vanillaEP。

J GPU SM利用率

SM利用率分析。
1. 微批次大小的影响：如表A.8所示，较小的微批次（即较大的R）可能会导致GPU SM利用率降低（例如GPT2-Tiny-MoE），但对于较大的MoE模型，由于参与GPU计算的实际张量仍然足够大，SM利用率几乎不受影响。
2. 批次大小的影响：如表A.9所示，较小的批次大小更容易导致SM利用率降低，尤其是在较小的模型上。对于参数量大的模型，即使批次大小较小，SM资源也能被高效利用。
3. 激活专家数量的影响：如表A.11所示，在专家数量较多的情况下，激活的专家数量越少（即f值越大，路由越不均衡），GPU上的计算负载越不平衡，导致GPU之间SM利用率的差异越大。

Table A.8: 不同微批次大小下的平均GPU SM利用率。

Table A.9: 使用FlowMoE时不同批次大小下的平均GPU SM利用率。

Table A.11: 在集群1的16个GPU上使用FlowMoE时，大量专家在不同激活专家数量下的最大和最小GPU SM利用率。

K 鲁棒性分析

异构集群的鲁棒性。
1. 不同计算能力的GPU：由于all-reduce和A2A任务是同步的，迭代时间由最慢的GPU决定。FlowMoE能最大化最慢GPU的计算/通信重叠，性能依然优于基线。通过在集群1上模拟异构GPU（给一半GPU增加延迟），实验结果（表A.12）验证了FlowMoE在异构计算能力下仍能取得最短的训练时间。
2. 非对称带宽：由于集合通信任务中所有GPU同时开始和结束，任务调度时间线对所有GPU都相同。FlowMoE凭借其高效的流水线方案仍能取得更优性能。

Table A.12: 当GPU具有不同计算能力时，平均迭代时间的比较（毫秒）。S1、S2、S3和S4分别是FlowMoE相对于vanillaEP、FasterMoE、Tutel和ScheMoE的加速比。

动态硬件环境的鲁棒性。为了应对网络带宽和GPU算力动态变化，我们设计了重新贝叶斯调优机制。当当前迭代时间与上一次BO预测的最优$S_p$对应的迭代时间变化超过预设阈值$\delta$时，将自动重新执行BO以寻找新的最优$S_p$。

$$\frac{|T-\hat{\mathcal{F}}(S_p^{best})|}{\hat{\mathcal{F}}(S_p^{best})} > \delta,$$

节点掉线的鲁棒性。为提高系统容错性，FlowMoE采用以下方法：
1. 参数备份：每个专家参数的副本存储在两个不同节点上。当一个节点发生故障时，其备份副本可以继续提供服务。每隔固定迭代步数（如1000步）同步一次副本间的参数。
2. 故障检测与恢复：FlowMoE定期调用torch.distributed.barrier()并设置超时。当操作超时抛出异常时，判断为节点故障。恢复过程包括：更新所有幸存节点上的门控函数路由表，将请求重新映射到备份节点；然后销毁当前通信组并重新初始化一个新的分布式通信组，使训练在剩余的幸存节点上继续进行。