Zhongzhu's Blog

Keep

0%

DisagMoE：用解耦 Attention 和 FFN 打通 MoE 训练的 all-to-all 瓶颈

Posted on Edited on In ,
一篇关于 DisagMoE 的中文阅读笔记：把 attention 和 FFN 分别放到独立 GPU 池，用 AF-Pipe 调度和 M2N 通讯原语把两侧拼起来，从而把 MoE 训练里的 all-to-all 瓶颈藏进计算之下。

笔记日期： 2026-05-14
笔记作者： Zhongzhu Zhou
所读论文： DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism
论文作者： Zhichen Zeng, 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（字节跳动 Seed；华盛顿大学；康奈尔大学）
arXiv： 2605.11005v1，2026-05-10
收录状态： arXiv preprint

简短结论

这篇论文回答了一个我以前没认真想过的"结构性"问题：如果 attention 层和 FFN 层在每个训练 step 里的 计算 / 通讯 比例是天生不一样的，那为什么我们至今还是把它们放在同一组 GPU 上、用同一份网络带宽预算来训？

DisagMoE 的回答是：别这么做了，把它们分开放。具体来说：

  • 一组 GPU 只跑 attention 层（组内用 data parallel）。
  • 另一组 GPU 只跑 FFN/MoE 层（组内用 expert parallel）。
  • 两组之间用一种新的 many-to-many 通讯原语（论文里叫 M2N / N2M）来传 hidden state 和反向梯度。
  • 整条流水线被设计成一个叫 AF-Pipe 的 schedule：两组工人不停地在不同 micro-batch 上做 forward 和 backward，A-side 和 F-side 永远互为彼此的下一个生产者/消费者。
  • 怎么分 GPU、怎么分 NIC，由一个 roofline 模型 + 一个很小的 MILP（混合整数线性规划）+ 一段局部搜索来一次性解出来。

在 8~16 节点 × 8 张 H800 的集群上，对 DeepSeek-MoE、GPT-OSS-120B、Qwen3-235B 这三个真实大模型、4K~32K 序列长度做测试，DisagMoE 相比 Megatron-LM 1F1B 交错流水线取得 最高 1.81× 吞吐提升；相比目前最强的 MoE overlap 系统（Tutel、Comet、DualPipe），最高 1.34×。在 iteration 时间分解里，DisagMoE 把 Tutel 留下的"没盖住的通讯尾巴"砍掉 88%，Comet 的砍掉 75%，DualPipe 的砍掉 45%。

它解决的关键观察是这样的：在一个混合 EP+DP 的 EP-group 里，attention 的 FLOPs 是序列长度 s 的二次函数（α₁s² + α₂s），而 FFN 的 FLOPs 和 all-to-all 通讯量都只是 s 的线性函数。所以随着上下文变长，attention 的算术强度（FLOPs / byte）一直往上爬，FFN 的算术强度几乎不动。结果是：attention 已经撞上 compute roof 了，FFN 还死死卡在通讯带宽这条斜线下面。如果你只用一种网络带宽配置来"喂"这两种性格完全不同的层，FFN 永远吃不饱通讯，attention 永远多了带宽没地方花。

我自己看完最深的一个体会是：DisagMoE 走的是和 DistServe（prefill / decode 解耦）、MegaScale-Infer / StepFun（推理时 attention / FFN 解耦）一脉相承的路线，只是把它从"推理"搬到了"训练"。而训练的难点是它在时间上是对称的——既有 forward 也有 backward，每个 worker 都要存自己负责那部分参数和 optimizer state——所以解耦必须同时在"空间上分组"和"时间上排流水线"两个维度都做对，AF-Pipe 就是那条让"空间解耦"在"时间维度"也能跑得动的 schedule。这部分是我觉得最值得对着论文图细看的一段。

1. 前置知识

这一节是为还没有把 Megatron-LM / DeepSpeed 这套训练框架完整啃过一遍、但已经训过普通 transformer 的同学准备的。我会把 MoE、expert parallel、all-to-all、pipeline parallel、roofline 这五件事都用一段话讲清楚，再回到论文本身。

1.1 一句话讲清 MoE

Mixture-of-Experts（MoE）就是：把 transformer 的某些 FFN 层从"1 个大 FFN"换成"E 个并行的小 FFN（experts）+ 1 个小的 gate 网络"。对每个 token，gate 选出概率最高的 top-k 个 expert（典型 k 在 2 到 8 之间），只把这个 token 喂进这 k 个 expert，最后把输出加权求和。这样做的好处是：模型总参数随 E 涨，但每个 token 实际跑的参数只随 k 涨。论文里反复用 DeepSeek-V4-Pro 当例子：总共 1.6T 参数，但每个 token 只激活 49B（从 384 个 expert 里选 6 个）。MoE 就是怎么把"万亿参数"塞进"我们还真能训"的工程答案。

1.2 Expert Parallel（EP）：tokens 要去找 experts

experts 多了之后，一张 GPU 装不下所有 expert。Expert Parallelism 就是把不同 expert 分别放到不同 GPU 上。一个 EP-group 是一组 GPU，它们合起来正好凑齐一层 MoE 里的全部 expert。当 gate 决定了"这个 token 要去 expert 42 和 expert 17"，这个 token 就得从当前所在的 GPU 跳到那两个 expert 所在的 GPU。这件事本质上是一个 all-to-all：每张 GPU 都要把自己手里的 token 切碎，发给所有需要它的其他 GPU。

一层 MoE 在 forward 里有 两次 all-to-all：

  1. Dispatch：按 expert 重新排 token，发到对应 GPU。
  2. Combine：experts 算完，把结果发回原来 micro-batch 的位置。

backward 是镜像，所以一个训练 step 里每层 MoE 一共会有 四次 all-to-all。每次的通讯量基本和整个 micro-batch 的激活量一样大。这就是后面所有问题的根源。

1.3 为什么 all-to-all 是瓶颈

同节点 GPU 之间走 NVLink，约 900 GB/s。跨节点 GPU 之间走 InfiniBand（IB）或 RoCE，约 100 GB/s，慢了快一个数量级。一旦模型大到一层 MoE 的 experts 一张 8 卡节点装不下，EP-group 就必须跨节点，all-to-all 就被迫走慢的 IB。论文在 64 张 H800 上的 profiling：单节点时 all-to-all 占 forward step 时间约 22%，扩到 8 节点直接爬到约 78%。也就是说在大规模时，计算单元有 78% 的时间在等通讯，再加 GPU 也没用。

1.4 Pipeline Parallel、1F1B、流水线 bubble

一张 GPU 装不下整个模型的时候，Pipeline Parallelism (PP) 把模型按层切成 stage：stage 0 拥有第 0~k 层，stage 1 拥有第 k+1~2k 层，依此类推。一个 micro-batch 顺着 stage 走完 forward，再倒回去走 backward。最经典的调度是 1F1B（one-forward-one-backward）：每个 worker 不停在"forward 第 i+1 个 micro-batch + backward 第 i 个 micro-batch"之间切换，这样每个 worker 的显存峰值大致稳定。Megatron 还加了一种 interleaved 变种：每个 worker 拥有不连续的多个 virtual stage，方便用 P2P 的 inter-stage 传输来盖过 compute。

Bubble 就是流水线里"这个 stage 在干等下一个输入到达"的空档。warm-up bubble 出现在 step 开头（最后一个 stage 还没收到第一个 micro-batch），cool-down bubble 出现在结尾。PP 研究的一大目标就是把 bubble 压小。

1.5 一分钟讲清 roofline

Roofline 图 的横轴是 算术强度（arithmetic intensity）= FLOPs / 单位字节通讯量，纵轴是实际拿到的算力 FLOPs/s。图上有两条线：

  • 一条从原点出发的"带宽斜线"：实际算力上限 ≤ 带宽 × 算术强度。
  • 一条水平的"compute roof"：实际算力上限 ≤ 峰值 FLOPs。

两条线交点处的算术强度叫 Î = 峰值 FLOPs / 峰值带宽。算术强度高于 Î 的负载是 compute-bound，低于 Î 是 bandwidth-bound。论文里把"带宽"换成了"网络带宽（IB）"，问题就变成"每层 MoE 的计算能不能盖过它的 all-to-all 通讯"。关键观察是：随着 s 变大，attention 和 FFN 会落在 Î 的不同侧。

1.6 什么叫"disaggregation（解耦）"

在 2024–2026 的 LLM 系统圈子里，"解耦"指的是 把模型不同组件放到不同硬件池里，让每个池子按它自己工作模式被独立尺寸化、组网、排程。代表作：

  • DistServe（OSDI 2024）：把 prefill（compute-bound、可大批量）和 decode（memory-bound、对 latency 敏感）放到不同 GPU 池。
  • MegaScale-Infer / StepFun / StepMesh：把 attention 层（decode 时 memory-bound）和 FFN/MoE 层（decode 时 compute-bound）放到不同硬件。
  • HeterMoE：把 AFD（attention/FFN disaggregation）思路引入训练，但只在单层粒度做，模型一大就 OOM。

DisagMoE 是这条线的下一步：训练版本的 AFD，靠"深度交错的分组 + 流水线 + roofline 自动定尺寸"把它做到可以训 235B+ 的规模。

2. DisagMoE 想攻击的核心问题

2.1 attention 和 FFN 的计算-通讯不对称

论文第 3 节把这个问题写得非常干净。在一个 EP-group 内，给定序列长度 s、hidden H、FFN-hidden De、GQA group size g、batch b、top-k k：

  • attention 计算量：约 α₁s² + α₂s（s² 来自 QK^T 和 softmax·V，α₂s 是 QKV 投影和 output 投影）。
  • FFN 计算量：约 βs（两次 GroupGEMM：up 投影和 down 投影，k 个 expert）。
  • all-to-all 通讯量：约 γs（与 top-k × hidden 成正比）。

换算成 算术强度（FLOPs / byte）：

  • I_attn = (H(2 + 2/g) + 4s) / (2k) — 随 s 线性增长。
  • I_ffn = 2 De — 跟 s 无关。

论文 Figure 6a 给的 Qwen-3 跨 8 节点数据：

序列长度 attention 占 step 算力 I_attn I_ffn
4K 28.4% 1.08 0.73
32K 50.3% 2.78 0.77

H800 在 IB 网络下的"turning point" Î 正好落在两者之间。结果是：用同一份带宽喂这两类层，attention 已经在 compute roof 上跑得很满，FFN 还卡在通讯带宽这条斜线下面。整体系统的实际算力 = 这两点的加权平均，所以会被 FFN 那一侧死死拽下来。给系统加带宽对 attention 是浪费，加算力对 FFN 也是浪费。

2.2 为什么 operator-level overlap 不够用

Tutel / Comet / FasterMoE 这一脉是把 FFN 的 GroupGEMM 切成小块，让每块的 compute 和切碎的 dispatch/combine 在 tile 级别上 overlap。问题是这条 overlap 窗口被 FFN 的 compute 时间卡住——而 FFN 的 compute 时间在跨节点场景下小于 all-to-all 时间，所以永远会留下一截盖不住的通讯尾巴。论文 Figure 5 里那一条橙色的"non-overlapped communication"条带就是这个尾巴。

Lancet / DualPipe 把 overlap 拉到 micro-batch 级别（用 batch B 的 compute 盖 batch A 的通讯），但仍然解决不了"attention 是 s²、FFN 和 a2a 是 s"这种 结构性的不对称。所以在长序列场景里，DualPipe 在 warm-up 和 cool-down 也有大块的暴露通讯。

2.3 为什么"推理用的 AFD"不能直接搬到训练

MegaScale-Infer 之流的 AFD 把 attention 丢到老 GPU，把 FFN 丢到新 GPU，因为推理 decode 时 attention 是 memory-bound、FFN 是 compute-bound。但训练有两条要命的差别：

  1. 每个 worker 必须保留自己负责那部分 layer 的全部参数 + optimizer state + 梯度。Qwen3-235B 有 94 层 attention，如果一个 worker 要单独存 94 层 attention 的优化器状态，根本装不进 80 GB H800。
  2. HeterMoE 证明这条路在小模型（约 4.3B）上能跑，但到 100B+ 就 OOM。

DisagMoE 要解决的就是 怎么让 AFD-for-training 不 OOM 地扩展到上百 B。

3. DisagMoE 的设计

可以把 DisagMoE 拆成三层来读，每层回答一个问题：

  1. 空间上，哪些 GPU 装哪些层？ — 3.1 节，AF disaggregated placement。
  2. 时间上，工人之间按什么顺序工作、怎么 overlap？ — 3.2 节，AF-Pipe schedule。
  3. 尺寸上，A-pool 给多少 GPU、多少 NIC？F-pool 又给多少？ — 3.3 节，adaptive worker allocation by roofline MILP。

3.1 AF 解耦的空间分组（带深度交错）

设模型 L 层。选一个共享的"流水线深度" p。attention 一侧被切成 p 个组 A₀, A₁, …, A_{p-1}，FFN 一侧也切成 p 个组 F₀, F₁, …, F_{p-1}。层分配按深度交错：

组 g 拥有的层是 {g + kp | k = 0, 1, …, v-1}，其中 v = L/p 是每个组的 virtual stage 数。

举个例子：L = 8、p = 2。attention 这边：

  • A₀ 拥有第 0、2、4、6 层 attention；
  • A₁ 拥有第 1、3、5、7 层 attention。

FFN 这边同样的模式。最后输出 embedding 跟着拥有最后一层 FFN 的那个 F-group 走。

为什么不直接"每组一层"？ 因为那样每个 group 只装 1 层，需要的 group 数 = 总层数 L，optimizer state 在跨多组复制后会爆显存。HeterMoE 卡在这一点上。深度交错让每组装 v 层，但 group 数只有 p，这是关键的一笔显存账。

组内：

  • A-Worker（attention 组）：组内 GPU 走 data parallel，每张 GPU 保留这个组负责的全部 attention 层的完整参数，处理 micro-batch 的不同切片，组内 all-reduce 梯度。
  • F-Worker（FFN 组）：组内 GPU 走 expert parallel，experts 在组内 GPU 间分片，token 进入组后再路由到对应 GPU。

A-group 的 GPU 数 M 和 F-group 的 GPU 数 N 不要求相等——这个不对称就是 3.3 节那个 roofline 模型要去调的旋钮。

3.2 AF-Pipe：调度

把 attention 和 FFN 物理分开之后，A-Worker 和 F-Worker 形成一个双向的生产者-消费者关系：

  • forward：A 算出 hidden state h → F 消费它，吐出经过 expert mix 的 h′ → A 再消费 h′，算下一层的 h，循环往复。
  • backward：梯度反向流。

如果一次只跑一个 micro-batch，一半的工人会闲着。经典解法：塞很多 micro-batch 进流水线，让两侧永远有事做。AF-Pipe 就是这个流水线，比 Megatron 1F1B 多了两个关键改动：

1. MoE 的 all-to-all 升格成一个一等公民的 pipeline stage。 Megatron 1F1B 里 dispatch / combine 是层内的事，只能靠 chunked GEMM 来 overlap。AF-Pipe 把它们抽出来变成 显式的 M2N / N2M 阶段，bubble 分析也因此变得清爽。

论文给出明确的 bubble 公式：设 PP 流水线总 stage 数为 PP，每 stage 含 L 层、每组 v 个 virtual stage，attention / FFN / a2a 每层耗时 T_a, T_f, T_a2a，P2P 耗时 T_p2p：

  • Megatron 1F1B 的 warm-up bubble：B_base = ((PP-1)/v) · [L(T_a + T_f + 2 T_a2a) + T_p2p]
  • AF-Pipe 的 warm-up bubble：B_AF = ((PP-1)/(2v)) · [max(T_a, T_f) + T_M2N]

当 L = 1 且 T_M2N ≈ T_a2a，AF-Pipe 的 bubble 大约只有 baseline 的 1/4。

2. M2N 把"P2P + all-to-all"两段合并成一段全双工传输。 Megatron 的 interleaved 流水线里：inter-stage 是 P2P，intra-stage 是 all-to-all（dispatch/combine）。DisagMoE 把这两件事合并成一个 many-to-many 原语——A-Worker 的每张 GPU 直接发到 F-Worker 的每张 GPU，反之亦然，全双工让发送和接收方向各自跑满 NIC。论文说这样大约把整体通讯成本降低了 1/k。

3. 多流异步执行。 每个 worker 内部并行跑三条 CUDA stream：forward 计算流、backward 计算流、通讯流。论文 Figure 8(c) 给了一个具体快照：F-worker 在同一时刻并行做三件事——往前一个 A-worker 发 F⁶₁ 的反向梯度（通讯流），在 compute 流里做 F⁶₆ 的前向，从下一个 A-worker 收 A²₇ 的反向输出。如果 3.3 节的分配器把尺寸算对了，这三件事的耗时差不多，流水线就跑在峰值。

实现：基于 Megatron-LM 和 PyTorch v2.6，约 6K 行 Python + 2K 行 C++。M-to-N / N-to-M 原语跑在 GPUDirect + GPUCopy 上，思路沿用 MegaScale-Infer 的 StepMesh。

3.3 自适应资源分配 = roofline MILP + 局部搜索

只把两边分开是不够的，还得把 GPU 和 NIC 分对：A-pool 给太多 GPU，FFN 这边就饿；反之亦然。论文把这个尺寸化问题写成一个很小的 MILP：

设总 GPU 数 W、总 NIC 数 M_tot。A-pool 拿 M 张 GPU、M_a 张 NIC；F-pool 拿剩下的 N = W - M、M_f = M_tot - M_a。两边的 stage 耗时是 计算时间 和 通讯时间 的取大：

  • T_a = max(C_a / (M · P), V / (M_a · B_IB))
  • T_f = max(C_f / (N · P), V / (M_f · B_IB))

AF-Pipe 把两边耦合成生产者-消费者，所以 step 时间 = max(T_a, T_f)。不平衡就是 bubble。优化分两步：

  1. 第一步：最小化瓶颈 stage — T* = min max(T_a, T_f)，扫所有可行的 (M, M_a)。
  2. 第二步——同样最优 T* 时用 MFU 来 tie-break：在所有达到 T* 的配置里，挑能最大化 MFU = (C_a + C_f) / (T* · W · P) 的，这在 T* 固定时等价于最大化 I_attn + I_ffn，也就是把"被 communication-bound 的那一侧的 compute roof 抬一抬"。

MILP 给一个种子 (M₀, M_{a,0})，然后 Algorithm 1 在它附近做半径 r 的局部搜索（采样 K 次），每次 profile 真实耗时来确认是否更好。整个过程一次性离线解出——pretraining 的 shape 是不变的，根本不需要在线重分配。

论文 Figure 9 那张 roofline 图把意图画得很清楚：aggregated baseline 里，attention 和 FFN 共享同一根带宽斜线，FFN 被卡在通讯下面；AFD 之后，A-pool 拿到的 NIC 少了一些（它本来也用不到那么多带宽），F-pool 拿到的 NIC 多了一些（斜线变陡），两边都更靠近自己的 compute roof。整个系统的总性能高于任何对称分法。

4. 实验和它们到底证明了什么

4.1 硬件 + 模型 + 基线

  • 硬件：16 节点，每节点 8 张 H800（80 GB）+ 168 CPU + 8× 400 GbE ConnectX-7 NIC。节点内 NVLink 400 GB/s。
  • 模型：DeepSeek-MoE（28 层、64 expert、top-4、hidden 2048、MoE-hidden 1408）；GPT-OSS-120B（36 层、128 expert、top-4、2880/2880）；Qwen3-235B-A22B（94 层、128 expert、top-8、4096/1536）。Qwen3 是最难的，又深又宽。
  • 基线：Megatron-LM 1F1B-interleaved（3D 并行）、Tutel（operator-level overlap）、Comet（细粒度 overlap 库）、DualPipe（DeepSeek-V4 的双向流水线 overlap）。
  • 扫描：序列长度 4K / 8K / 16K / 32K。所有基线 EP=16 跨 2 节点，PP / virtual-stage 各自调到不 OOM。

per-device batch size 固定，所以总 batch 随 GPU 数线性扩展。

4.2 端到端吞吐

Figure 10 画了三个模型、四个序列长度下的 speedup（相对 Megatron）：

  • 对 Megatron-1F1B：1.59× ~ 1.81×。序列越长，attention 和 FFN 的不对称越严重，AFD 收益越大。
  • 对 Tutel / Comet：1.2× ~ 1.5×。operator-level overlap 进不到 attention 层内部，全部通讯尾巴都摆在外面。
  • 对 DualPipe：1.05× ~ 1.13×。DualPipe 已经用 micro-batch 流水线把大块通讯盖了，多出来的这点 win 主要来自 AFD 把 GPU/NIC 重分配和 M2N stage 融合。论文也说：DualPipe 在 warm-up / cool-down 仍然有大块暴露通讯，steady state 也仍然受 attention/FFN 不对称的折磨。

4.3 iteration 时间分解

Figure 11 把一次 step 拆成 三段颜色：绿色 = overlap 的计算；紫色 = overlap 的通讯；橙色 = 没盖住 的通讯。看橙色那一段就够了：

  • Megatron-1F1B 的橙色非常长（约占 step 的 1/3）。
  • Tutel 的橙色，DisagMoE 砍掉约 88%。
  • Comet 的橙色，砍掉约 75%。
  • DualPipe 的橙色，砍掉约 45%。

这是论文里最有说服力的一张图，因为它隔离出"解耦"具体到底在哪里赚到了时间：绿色（compute）几乎没变，橙色（暴露通讯）大幅缩小。这正好对应它的设计意图。

4.4 资源分配 ablation

Figure 12 固定 NIC 16:16，只改 GPU 的 A:F 比例。三个 pattern：

  • 4K 序列下，最优 A:F 约 16:16（attention/FFN 计算量接近）。
  • 16K 序列下，最优 A:F 是 16:10——attention 算得多，应该拿更多 GPU；FFN 算得少，少一些 GPU 也跑得动。
  • 强行使用 16:8 或 8:16，又把不对称放回去，speedup 跌回 1× 附近。

所以那个 MILP 不是可有可无的。序列长度对最优分配有真实影响，靠手调猜不准。

4.5 top-k 和 EP-size 的 ablation

Figure 13 固定 64 expert、扫 top-k ∈ {2, 4, 8, 16}。top-k 大了，通讯量 V 和 FFN 算力 C_f 都线性涨，整体 step 时间涨。DisagMoE 在整条扫描线上保持 1.08× ~ 1.92× 的 speedup，最大值出现在 top-k = 16 + 大 EP 这种"对基线最压力"的配置上——也就是 all-to-all 最痛苦的时候。

4.6 virtual stage 的 ablation

Figure 14 固定 p，扫 virtual stage v。throughput 随 v 接近线性上升，直到 v = 16 开始 OOM——一个组里要存 16 层的参数和激活，显存撑不住。这正是 HeterMoE 在单层粒度上的失败模式。DisagMoE 的 p ≥ 1 这一深度交错的设计，是让 v 维度能往上爬的原因。

5. 局限性与我会进一步追问的地方

论文自己很诚实，列了两个明确的 limitation，我再补两条工程视角的疑问。

5.1 只适用于 pretraining 的固定 shape

MILP+局部搜索假设 sequence length、micro-batch、global batch 都是固定的。pretraining 满足（你启动一次跑几周，shape 永远不变）。但是：

  • RL 后训练（GRPO、DAPO、PPO 这类 rollout）每条样本的 prompt 和 response 长度都不一样。
  • SFT 用 packed sequence 训 的话，每个 batch 的有效 token 数也不一样。

要把 DisagMoE 搬到这些场景，需要 在线重分配 (M, M_a)。论文明确把这件事留给 future work。

5.2 attention 和 FFN 共享同一个 p

论文目前让 p 在 attention 和 FFN 之间共享，公式比较干净。但是显然的下一步：不对称的 (p_A, p_F)。如果 FFN 是瓶颈，那应该给 FFN 切更多的组（每组更少层），让瓶颈侧拥有更多并发 stage。论文说这是一个 promising extension，但调度复杂度会增加，所以留给后续工作。

5.3 我想看但论文没给的几张图

  • loss 曲线对比。所有报告都是吞吐数字。论文没给"DisagMoE vs Megatron 训出来的最终 loss 曲线对比"，从数学上说梯度应该完全相同（同样的模型、同样的样本），但生产化前我还是想看一张并列的最终 loss 表。
  • straggler 行为。同时跑 p 个 attention 组 + p 个 FFN 组，任何一张慢卡都会把整条流水线拖住。在百 B 规模上 straggler 是日常，论文没注入 straggler 做对比。
  • 网络拓扑敏感性。评测集群是均匀的 8× 400 GbE，真实生产集群常常 rail-aligned 且在 spine 上有 2:1 的 oversubscription。M2N 原语对 rail alignment 应该是敏感的，论文里没有刻画这部分。

5.4 哪些模型适合用 DisagMoE？

DisagMoE 收益大小其实由两个比率决定：

  1. I_attn / I_ffn — 你的模型每层算术强度的不对称程度。长序列放大它；MoE + 小 top-k 进一步放大。
  2. T_a2a / T_compute — all-to-all 占总时间的比例。跨节点 EP 放大它；单节点训练里 operator-level overlap 就够用，不需要 DisagMoE。

所以 DisagMoE 的甜区是 大 MoE + 跨节点 EP + 长序列，这正好是 2026 年生产 MoE 训练的常态。

6. 这篇论文在系谱里的位置

放回 2026 年这个时间点，"LLM 系统的解耦"已经是一个清楚的家谱：

  • 推理解耦：DistServe（prefill/decode）、MegaScale-Infer / StepFun / StepMesh（attention/FFN）、Mooncake（KVCache 解耦）。
  • 训练 overlap（不解耦）：Tutel、FasterMoE、Comet、Lancet、DualPipe。
  • 训练解耦的前辈：HeterMoE（单层粒度 AFD-for-training，在 4.3B 以上 OOM）。
  • DisagMoE（本文）：深度交错的 AFD-for-training，扩展到 235B+ 在 8~16 张 H800 节点上。

DisagMoE 的核心 intellectual move 是：把推理时的解耦认真地搬到训练上，并解决训练特有的障碍（optimizer state 爆显存、梯度反向流、warm-up + cool-down 都有 bubble），靠的是 schedule（AF-Pipe） 和 尺寸化模型（roofline MILP），而不是靠新硬件。

7. 伪代码 + 一个数字 worked example

想自己实现 DisagMoE 的同学可以把训练 step 想成三层嵌套——下面的伪代码是概念级别的，生产版本还要处理 worker 内部的 tensor parallel、梯度累积、optimizer state sharding 等，这些 Megatron-LM 都已经搭好了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def disagmoe_step(model, batch, allocator):
    # 1. 静态尺寸化（启动时跑一次）
    M_attn_gpus, M_ffn_gpus, M_attn_nics, M_ffn_nics = allocator.solve(
        model_cfg, hardware, seq_len)

    # 2. AF 解耦的空间分组（启动时建好）
    a_groups = make_attention_groups(
        model, p=allocator.p, gpus=M_attn_gpus, nics=M_attn_nics)
    f_groups = make_ffn_groups(
        model, p=allocator.p, gpus=M_ffn_gpus, nics=M_ffn_nics)

    # 3. AF-Pipe steady-state 调度
    micro_batches = split(batch, n=num_micro_batches)
    for tick in pipeline_schedule(a_groups, f_groups, micro_batches):
        # 每个 worker 内部独立推进三条 stream
        worker = tick.worker
        stage = tick.stage   # FWD_A / FWD_F / BWD_A / BWD_F / M2N / N2M
        run_async(worker.compute_stream, stage)
        run_async(worker.comm_stream, m2n_transfer_for(stage))
    optimizer.step()

pipeline_schedule 内部的 AF-Pipe warm-up 阶段先在 forward 方向发出 p × v 个 micro-batch，之后进入 steady state——每个 worker 每个 tick 都同时推进 1 个 forward compute + 1 个 backward compute + 1 个 communication。M2N 在 NIC 上全双工跑。

数字 worked example：speedup 到底从哪里来

拿 Qwen3-235B、s = 16K、8 节点（64 张 H800）。从论文 profiling 倒推（把 Megatron baseline 的 step 归一化为 1.0）：

  • Megatron-1F1B：

    • attention compute: 0.42
    • FFN compute: 0.18
    • 没盖住的 all-to-all + P2P: 0.40
    • 总 step ≈ 1.00。

  • DisagMoE（同硬件、A:F = 16:10）：

    • T_a 和 T_f 平衡到 max(0.45, 0.43) = 0.45。
    • M2N 完全被 compute 盖住，残余暴露通讯 ≈ 0.07。
    • 总 step ≈ 1.00 / 1.81 ≈ 0.55。

也就是说，头条 1.81× 的来源主要是橙色（暴露通讯）那一段被压扁了，而 compute 这一段几乎不变。这是实现者最应该记住的一句话。

8. 符号表

符号 含义
L transformer 总层数
E 一层 MoE 的 expert 数
k 每个 token 路由到的 top-k expert 数
H hidden size
De FFN/expert 的 hidden size
s 序列长度
b micro-batch size
W 总 GPU 数
m, M A-pool 的 node / GPU 数
n, N F-pool 的 node / GPU 数
EP expert-parallel size
P per-GPU compute（FLOPs/s）
B_NV, B_IB NVLink / InfiniBand 带宽（B/s）
p 每个组件类型切的 pipeline 组数（深度）
v 每个组里的 virtual stage 数
I_attn, I_ffn 算术强度（FLOPs / byte）
Î 系统 turning point = 峰值 FLOPs / 峰值带宽
T_a, T_f attention / FFN 每个 stage 的耗时
T_a2a 一层 all-to-all（dispatch+combine）
T_M2N 组间 many-to-many 传输
T_p2p 组间 point-to-point 传输
B_base, B_AF Megatron-1F1B 和 AF-Pipe 的 warm-up bubble

9. 复现注意事项

  • 代码。论文没有给出开源地址。文本说"约 6K 行 Python + 2K 行 C++ 基于 Megatron-LM + PyTorch v2.6"，所以一个熟悉 Megatron 的工程师大概能在数周内重写出 placement 和 M2N 原语；最耗时的部分是 AF-Pipe 的精确调度和 MILP 分配器。
  • 硬件。论文报告的吞吐数字基于 ConnectX-7 8×400 GbE NIC + 8×H800 节点。把 NIC 换成 200 G 或单 NIC 节点，M2N 就会变成新瓶颈，speedup 会缩水。
  • 模型。三种评测模型（DeepSeek-MoE / GPT-OSS-120B / Qwen3-235B）都是公开权重，训练 recipe 不一定公开，但被测的是训练 step 耗时，任何同尺寸 MoE 模型都应该能复现 throughput 比率。
  • Roofline MILP。直接照 Algorithm 1 实现就行：任何开源 MILP 求解器（PuLP / OR-Tools）都能处理整数分配，局部搜索这一步只需要每个 stage 的实测耗时。

10. 我的最终判断

我觉得 DisagMoE 值得被认真读，理由跟我两年前认真读 DistServe 是一样的：它发现了一个 结构性 不对称（MoE 训练里 attention 和 FFN 的计算-通讯比例本来就不同），点出这个不对称在浪费硬件预算，然后用 尺寸化 + 调度 的方式解决它，而不是发明新架构。论文第 3 节里的 roofline 论证是核心智识贡献，AF-Pipe + M2N 原语是把这一论证落地的工程脚手架。

如果我是一个 2026 年在生产里训 MoE 的团队，我会把这篇论文当作 本季度做三件事的强提示：

  1. 拿你当前的训练负载（你真实的 s）做一次每层 算术强度 profile。如果 I_attn 和 I_ffn 落在你的网络 Î 的两侧，这个解耦的 case 就成立。
  2. 在动手做完整移植前，先做一个小规模 AF-Pipe 原型。哪怕只是把 A 和 F 跑成两组 PyTorch worker（不上 M2N 原语），都能告诉你 AF-Pipe 的调度对你的模型合不合身。
  3. 把"不对称的 p"和"动态 shape"这两个 future work 提前规划进路线图。因为 RL 后训练（2026 年的标配）很快会撞到 limitation 5.1。

我也不会把这篇当作"通用替代 Megatron"。它是一个 专门化 的训练系统，专门为大 MoE + 跨节点 EP + 长序列这个区间做。但是在这个区间里，依据现有实验，它确实是目前公开发表的最强方法。

本笔记基于 v1 预印本（2026-05-10）阅读所得。所有量化推断和图表解读都基于论文公开内容；如有理解偏差，责任在我，而非作者。