DisagMoE：通过解耦 AF-Pipe 并行实现计算与通信重叠的 MoE 训练

# DisagMoE：通过解耦的 AF-Pipe 并行实现计算与通信重叠的 MoE 训练
来源：https://arxiv.org/html/2605.11005
1] ByteDance Seed 2] 华盛顿大学 3] 康奈尔大学

Chi-Chih Chang, Jiayi Wang, Zezhou Wang, Ningxin Zheng, Zheng Zhong, Cesar A. Stuardo, Dongyang Wang, Mohamed S. Abdelfattah, Haibin Lin, Banghua Zhu, Ang Li, Ziheng Jiang

(2026年5月10日)

###### 摘要

混合专家（Mixture-of-Experts, MoE）架构使得拥有万亿参数的稀疏激活大语言模型（LLMs）成为可能。专家并行（Expert Parallelism, EP）是一种广泛采用的 MoE 训练策略，但它面临着严重的全对全（all-to-all）通信瓶颈。随着模型规模的增大，需要将专家分布在不同的 GPU 节点上，而有限的节点间网络带宽进一步加剧了这一问题。以往的研究主要集中在将这些全对全通信与前馈网络（FFN）和自注意力（self-attention）计算进行重叠，但由于注意力层和 FFN 层在计算-通信比例上存在固有的不平衡，往往仍会留下受限于网络的停滞时间。我们提出了 **DisagMoE**，这是一个解耦的 MoE 训练系统，通过对模型放置和调度进行联合优化以实现最大效率。DisagMoE 将注意力层和 FFN 层分离到不同的 GPU 组中，引入了一种具有单向多对多通信的多级流水线，并采用计算-通信屋顶线模型（roofline model）来平衡注意力组和 FFN 组之间的 GPU 与网络带宽分配。DisagMoE 基于 Megatron-LM 实现，评估表明，DisagMoE 在多种 MoE 模型上提高了训练效率，在 16 节点 8×H800 集群上实现了高达 1.8 倍的加速。

## 1 引言

混合专家（MoE）架构被广泛用于扩展基于 Transformer 的大语言模型（LLMs）[24, 9, 3, 14]。通过将前馈网络（FFN）层划分为多个专家，并为每个 token 仅激活一小部分专家，MoE 架构将激活参数数量与总参数数量解耦，在扩展模型容量的同时实现了每 token 计算量的次线性增长 [16, 6, 20, 31, 26]。例如，DeepSeek-V4-Pro 的总参数规模达到 1.6T，但每个 token 仅激活 49B 参数 [8, 4]。

尽管具备计算效率，大型 MoE 层仍因总参数数量巨大而面临 GPU 内存容量挑战。专家并行（EP）通过在 GPU 设备间切分专家来解决这一挑战，其中注意力层等密集组件通过数据并行（DP）在同一 GPU 组内进行复制。这种 EP+DP 策略产生了两次全对全通信：token 必须从注意力层的微批次排名（ranking）洗牌到 FFN 层的专家排名（dispatch），然后再洗牌回来（combine）。由于动态专家路由的原因，这些通信传输了几乎整个微批次的数据，并且它们位于注意力层和 MoE 层之间的关键路径上。此外，随着模型规模继续扩大，FFN 层无法再放入单个多 GPU 节点中，迫使 EP 域跨越多个节点，较低的内部节点带宽进一步加剧了通信瓶颈。正如我们在 64×H800 GPU 上的性能剖析所示，通信占据了训练步骤高达 50% 的时间（图 1）。

**图 1：** MoE 执行分析。使用 Megatron-LM 在 8 个节点的 H800 GPU 上执行的 MoE 模型的时间分解。

优化节点间的 dispatch/combine 是大规模 MoE 训练的首要优化目标，该领域已有丰富的文献 [23, 22, 12, 13, 32]。一类工作将通信与计算重叠。Tutel [13] 和 Comet [32] 提出了算子级的分块（chunking），将 FFN 计算划分为细粒度的块，并通过块级流水线部分地将 dispatch/combine 与 MoE 层重叠。然而，如图 4 所示，重叠窗口受限于 FFN 计算量，导致部分通信暴露在外。Lancet [15] 和 DualPipe [7] 通过流水线处理多个微批次，用另一个微批次的计算来隐藏当前微批次的通信，从而扩大重叠窗口。尽管如此，它们仍然面临一个根本性的不平衡问题：如图 4 所示，注意力计算时间随序列长度呈二次方增长，而 FFN 计算和全对全通信大致呈线性增长。这种不平衡使得在广泛的工作负载范围内难以实现完美重叠，如图 6(a) 所示。

最近的工作 [29, 28, 34] 将注意力层和 FFN 层（AFD）解耦到不同的设备组上，以重叠组间通信。在推理服务方面，诸如 MegaScale-Infer [34] 和 StepFun [28] 等系统观察到，在解码阶段，注意力层相对更受内存限制，而 FFN 层更受计算限制，因此将注意力层和 FFN 层放置在异构硬件上。虽然这种设计对于延迟驱动的服务非常有效 [28, 34]，但很难应用于训练：每个组必须持有其组件在所有层中的完整参数，加上仅用于训练的状态（例如，优化器状态）。HeterMoE [29] 将 AFD 引入训练，并在使用旧款 GPU 的 FFN 组和采用新款 GPU 的注意力组之间提出了斑马并行（Zebra Parallelism, ZP），两者通过全对全原语进行通信。然而，这种两级流水线对于大规模 MoE 训练来说是不够的，因为每个设备组必须保留其对应组件在所有 Transformer 层中的参数，导致大规模下出现内存溢出（OOM）问题（图 14）。HeterMoE 仅展示了约 43 亿参数的规模，如何将通信重叠的 MoE 训练扩展到千亿参数规模仍是一个开放性问题。此外，现有方法维持固定的计算-通信比例，限制了在流水线各阶段平衡通信和计算的重叠设计空间。

为了解决这种不平衡并实现可扩展的重叠训练，我们提出了 **DisagMoE**，这是一个大规模 MoE 训练系统，它通过自适应的工作者分配在设备组间解耦注意力层和 FFN 层，联合优化模型放置和调度以缓解全对全瓶颈。我们的主要贡献如下：

*   **解耦的组件放置：** DisagMoE 按组件类型划分模型，将注意力层和 FFN 层分配给不同的工作组。注意力组复制密集组件并分发微批次训练数据，而 FFN 组切分专家并根据专家选择处理 token。每个组可以拥有来自多个 Transformer 层的同类组件，从而在 GPU 内存约束下最大化设备利用率。
*   **重叠的多级流水线：** 我们提出了 **AF-Pipe**，这是一种多级调度方案，用组间多对多（M2N）token 交换替代传统的全对全通信，并将其视为与注意力和 FFN 计算并列的一级阶段，通过对齐各组间的阶段边界，系统地将通信与计算重叠。
*   **自适应工作者分配：** DisagMoE 采用自适应分配器，联合调整注意力组和 FFN 组的规模，并重新分配它们之间的网卡（NIC）带宽，以平衡不同序列长度、top-k 选择标准和 EP 大小下的通信和计算时间。在轻量级计算-通信屋顶线模型的指导下，这种重新调整扩大了有效的重叠窗口，减少了关键路径上的通信，并缓解了限制先前重叠策略的注意力-FFN 不平衡问题。

我们在拥有数百亿参数的 MoE 模型上进行了广泛的预训练实验，使用了 128×H800 GPU 集群，序列长度范围为 4K 到 32K。DisagMoE 始终优于强大的基线方法，比 Megatron-1F1B 交错重叠高出高达 1.81 倍的吞吐量，比最先进的 MoE 重叠系统高出高达 1.34 倍。

## 2 背景与动机

### 2.1 混合专家（MoE）结构

混合专家（MoE）架构是一种在提高性能的同时高效扩展 LLM 的关键技术 [7]。MoE 架构用一组多个 FFN 替换标准的前馈网络（FFN），这些 FFN 被称为**专家（experts）**。

如图 2(a) 所示，MoE 模型使用可训练的网关网络来选择概率得分最高的 top-k 个专家。这些激活专家的输出通过加权和进行缩减，结果作为输入传递给下一个注意力层。现代大规模 MoE 模型表现出非常稀疏的激活，其条件计算原则使得模型参数扩展到数千亿的同时，浮点运算（FLOPs）仅增加次线性。例如，拥有 1.6T 总参数的 DeepSeek-V4-Pro 模型 [8]，在 384 个总专家中参与 6 个专家，为每个 token 激活 49B 参数。

**图 2：** MoE 架构、EP 以及注意力-FFN 解耦架构概述。

### 2.2 分布式 LLM 训练

分布式训练通过互补的并行化策略，使 LLM 超越单设备限制。

*   **数据并行（DP）：** DP 在设备间复制模型状态并分割批次以进行独立计算，通过 All-Reduce 同步梯度。DeepSpeed-ZeRO [21] 和 FSDP [33] 等内存高效变体对参数、梯度和优化器状态进行切分。
*   **流水线并行（PP）：** PP 将层划分为由激活传输连接的顺序阶段。Megatron-LM [25, 17] 采用 1F1B 调度以减少激活内存，仅留下较小的流水线气泡。
*   **专家并行（EP）：** MoE 使用 EP 在 GPU 间切分专家，token 动态路由到其分配的专家（图 2(b)）。每次前向/反向传递包括两次全对全操作，即 **dispatch** 和 **combine**，用于专家通信。

表 1：DisagMoE 符号描述。

## 3 挑战与动机

在本节中，我们首先阐述大规模 MoE 训练下的主要通信开销，并解释现有方法重叠不足的原因。然后我们介绍这些系统级问题如何促使 AF 解耦。

### 3.1 主导性的 MoE 通信开销

在大规模 EP 训练 [7, 30] 中，专家被切分到多个 GPU 和节点上。使用表 2 中的符号，由于每个 GPU 与每个非本地 GPU 交换 token，每个 GPU 的预期全对全通信量为 $V = \frac{EP-1}{EP}Nk$。传输在节点内使用 NVLink，在节点间使用 InfiniBand (IB) 或以太网 [10]。端到端通信延迟由较慢的互连主导，通常是接近 100GB/s 的 IB，远低于接近 900GB/s 的 NVLink。随着 EP 规模跨越更多节点，更大比例的同行通过 IB 到达，从而增加了通信延迟开销。

图 3(a) 显示，随着 EP 程度跨越更多节点，全对全通信在总训练时间中的占比从单节点上的 22% 上升到 8 个节点上的近 78%。同时，由于通信量 $V$ 几乎与所选 top-k token 的数量呈线性增长，如图 3(b) 所示，全对全通信的百分比大幅增加。因此，沿任一维度扩展都会在训练过程中产生大量的通信开销。随着通信开始占据主导地位，计算资源在交换完成前处于空闲状态，将系统推向通信受限模式。

**图 3：** (a) 跨节点的每层 MoE 全对全占比百分比。(b) 不同 top-k 选择的影响。两者均在 DeepSeek-MoE 的 8K 序列长度下。

### 3.2 未解锁的算子级重叠空间

鉴于主导性的全对全开销以及围绕注意力和 FFN 的两个对称全对全阶段（dispatch 和 combine），先前的系统 [13, 32, 15, 12] 试图通过将计算切片为算子级块，并将 dispatch 和 combine 与跨块的计算进行流水线化，从而隐藏通信，实现通信与计算的重叠。

由于 FFN 中的 Group GEMM 允许自然分块，token 可以分块传输和接收，计算可以与分块通信重叠。因此，大多数先前系统 [13, 32, 12] 将重叠限制在全对全阶段和 FFN 计算之间，如图 4 所示。一种由编译器驱动的方法 [15] 通过微批次级划分和流水线化来缩小注意力周围的通信气泡，从而扩展重叠。然而，带有 FlashAttention [5] 的自注意力仍然难以完全分块，因此当全对全通信占主导时，仍然存在未重叠的通信尾部，如图 5 所示。

**图 4：** (a) 原始 MoE 前向传播。(b) 全对全与 FFN 重叠，Comet 和 Tutel。(c) 微批次级流水线以减少气泡。当全对全占主导时，仍存在未重叠的尾部。

我们剖析了算子级重叠，发现狭隘的焦点……