背景
在 MoE 模型的训练过程中，EP rank 之间的 A2A 通信在端到端时间中占据了相当大比重，对训练效率影响很大，特别是对于 Fine-grained MoE model，EP size 会比较大，跨机通信基本无法避免。那么要如何减少 EP A2A 对整体训练效率的影响呢？
- 首先是尽可能提升 EP A2A 通信本身的效率。DeepSeek 开源的 DeepEP 通信库通过将通信拆分为机内通信和机间通信，并将二者进行 overlap， 实现了很高的通信效率。这一工作已经集成到 NVIDIA Megatron-LM 中；
- 其次是让 EP A2A通信尽可能和计算 overlap，用计算来掩盖通信的时间。本文从这个主题出发，介绍我们如何基于 1F1B 的流水线机制来实现 EP A2A 与计算的 overlap。
用计算掩盖 EP A2A 通信的相关研究
关于如何用计算掩盖 EP A2A 通信，我们注意到目前有以下这些研究（不完全统计）：
1, micro batch 内 overlap
- EP A2A 和 expert mlp overlap (minimax-01 方案、faster moe 方案、 tutel 方案)；
- 将 Attention + MLP 整体分成多个 block，block 间的 EP A2A 和计算 overlap (AMPipe)。
2, micro batch 间 overlap
- DeepSeek V3 论文提出了一种新的流水线并行算法：DualPipe。实现了在两个 micro batch forward 和 backward 之间的 EP A2A 通信与计算 overlap。
以上方案有各自的不足之处：
- EP A2A 和 expert mlp overlap： 其效果比较好的前提是计算时间要大于通信时间，但对于跨机 A2A 的情况，A2A 的通信可能会大于 expert mlp 计算时间，DeepSeek V3 论文中提到训练过程中 EP A2A 的整体耗时占比大约是 50%，此时这一方案的 overlap 效果就大打折扣了；
- Attention + MLP 整体分成多个 block 方案：需要分块 attention, 支持起来比较麻烦，而且分了 block 可能会降低计算效率，另外这个方案可能和其他分布式并行策略有冲突，比如 CP；
- DeepSeek DualPipe：毋庸置疑是很优秀的工作，但是调度逻辑比较复杂，难以理解与 debug，另外需要对现有的分布式训练框架（如 Megatron-LM）进行大范围的重构才能接入。
我们分析认为， Deepseek DualPipe 调度核心亮点是实现了 batch 之间 EP A2A 通信和 attention 、mlp 计算 overlap 的精细化调度, 从而使得跨机 EP 成为可行的分布式策略。
Dualpipe 起到的主要作用是创造了 overlap 的机会， 从这个角度而言，基于 1F1B 的调度也能做到 batch 之间 forward 和 backward 之间的 EP A2A 通信与计算的 overlap。
基于 1F1B 的 MoE A2A 通信计算 overlap
Conventional 1F1B
可以很自然地想到，通过调整 1F1B 调度来实现 1F1B 调度中稳态阶段 f 和 b 的 EP A2A 通信与计算 overlap。
原始的 1F1B 调度，最后一个 pp stage 稳态阶段的 f 和 b 之间存在数据依赖，无法实现 f 和 b 之间的 EP A2A 通信与计算 overlap。 为了解决这个问题， 我们可以把最后一个 pp stage 的 warmup step 加 1。同时，考虑到 pp stage 间前反向的数据依赖，其他 pp stage 也需要调整 warmup step。
具体的调度方式如下：
图 1: 1F1B 调度下 A2A 基于传统 1F1B 方案实现相邻 micro batch 间计算和通信的 overlap。
该图片来源于小红书 Agi Infra 团队，若您有任何疑问或需要使用该图片，请联系小红书
我们得到了一个 1F1B 通信计算 overlap 的初步方案，但是这个方案有很大的缺陷：因为 pipeline parallel rank 之间存在数据依赖，warmup step 过大，activation 显存占用过大。
Interleaved 1F1B
我们转向考虑 interleaved 1F1B ，沿用上面调整 warmup step 的逻辑，惊喜地发现，通过将稳态的 1F1B stage 中第一个 micro batch 的 fprop 提前到 warmup stage，即 warmup step + 1, 就可以在 interleaved 1F1B 实现 1F1B 稳态阶段不同 micro batch 前反向之间的 EP A2A 与计算的 overlap。
图 2: 基于 interleaved 1F1B 方案实现相邻 micro batch 间计算和通信的 overlap。
该图片来源于小红书 Agi Infra 团队，若您有任何疑问或需要使用该图片，请联系小红书
相比于原始的 interleaved 1F1B， 我们的这个方案虽然在 warmup stage 增加了一个 micro batch 的前向计算，但是 1F1B 阶段的调度是前反向细粒度地交叠调度，F 在增加 activation 显存的同时，B 会释放 activation 显存，使得峰值显存基本没有变化；Bubble rate 和原始的 interleaved 1F1B 完全一致。
因此我们最终选择使用 interleaved 1F1B with A2A overlap 方案。
Fine-grained schedule of 1F1B overlap
具体来说，对于两个 micro batch 之间 forward 与 backward 的 EP A2A 通信/计算 overlap, 我们采用了类似 DualPipe 的方案， 拆分了 dw。我们认为这是很有必要的。如果计算与通信 1:1, 那就意味着前向的计算比通信少， 反向的计算比通信多，前向的计算不足以 overlap 反向的通信, 需要用 dw 来 overlap 部分反向通信。
我们通过下面一组流程图来对比说明 dw 拆分的必要性。Baseline 版本将 forward 和 backward 的计算部分放到了同一个 stream 中，通信部分放到了另一个 stream 中，通过将 forward 和 backward 的各个模块交错排布，消除了同一时刻计算流和通信流上 kernel 的数据依赖，从而实现了 batch 间的 EP A2A 和计算的 overlap。
但如果 EP A2A 的时间比较长，超过了计算时间，比如下图里的 F/attn 和 F/mlp，此时 EP A2A 通信还是会暴露出来。这时我们可以把 mlp 和 attention 的 dw 从 backward 计算中拆分出来， 那么 B/dispatch 和 B/combine 的执行就可以提前，因为 B/dispatch 和 B/combine 只依赖 dx 的结果。这样用来 overlap B/dispatch 的计算除了 F/mlp 之外还有 W/mlp 的一部分，从下面两种方案的对比我们也可以看出，拆分 dw 之后，计算流上 kernel 的排布更加紧密，整体的执行时间也要比 Baseline 快。
图 3: 拆分 dw 前后，计算和通信 overlap 示意图。
该图片来源于NVIDIA，若您有任何疑问或需要使用该图片，请联系NVIDIA
与 Deepseek DualPipe 对比
| DualPipe | 1F1B with A2A overlap | |
| PP bubble overhead | (PP/2-1) x (F&B + B -3W) | (PP – 1) x (F + B) / vpp |
| Memory: param | 2X | 1X |
| Memory: activation | PP+1 | PP+(PP-1)/vpp |
注：F – forward chunk time, B – full backward chunk time, W – “backward for weight” chunk, F&B – two mutually overlapped forward and backward chunks.
为了更好地对比两种方案的 bubble overhead，接下来我们对 F/B/W 的时间做一个合理的假设（参考图三）：B=2F=2W，F&B 的耗时比较难以给出较准确的估计，我们合理地认为 F&B=B，即前向传播的时间完全被反向传播掩盖住，这种估计的合理性在于：
- 在拆分 dw 之前，反向传播的计算和通信时间都大于与其 overlap 的前向计算的时间；拆分 dw 之后，dw 的计算时间可以进一步被 overlap；
- 计算和通信的 overlap 在相当程度上会对计算 kernel 产生 overhead；
| DualPipe | 1F1B with A2A overlap | |
| PP bubble | (PP/2-1) x F | 3 x (PP-1) / vpp x F |
据以上公式，我们可以估算出两种方案 bubble overhead 的比例：
可以看出，vpp=8 大致是两种方案的 bubble overhead 比值的分界点，我们从一些常见模型的最佳配置里抽离出 vpp 的设置如下：
| Model | PP | VPP |
| GPT-175B | 8 | 8 |
| llama3-70B | 4 | 5 |
| llama3-405B | 9 | 7 |
| Nemotron-22B* | 4 | 10 |
| Nemotron-340B* | 8 | 12 |
| Mixtral-8x7B | 4 | 8 |
| Mixtral-8x22B | 4 | 14 |
| QWen* | 4 | 7 |
注：* 该模型的最佳配置为内部测试结果，训练脚本尚未开源
因此，多数情况下两种方案的 bubble overhead 是比较接近的。另外大多数情况下 micro batch 的数量比较多，PP bubble 在端到端的耗时占比处于比较低的水平，此时 PP bubble 对整体的性能影响比较有限。
关于 PP bubble 带来的 Activation memory 开销，从上面常见模型的最优配置可以看出，PP size 通常小于 VPP，此时 1F1B with A2A overlap 方案的 Activation memory 开销要低于 DualPipe。
总结
针对 MoE 模型训练过程中 EP A2A 通信占比过高的问题，我们分析了现有的优化策略，认为 在 batch 间进行通信计算 overlap 是比较高效的实现方案。参考 DeepSeek 提出的 DualPipe，我们设计了一种基于现有的 interleaved 1F1B 策略实现 batch 间计算通信 overlap 的方案，该方案具有以下几个特点：
- 基于现有的 1F1B 方案实现了和 DualPipe 相似的 batch 间 EP A2A 与计算 overlap 的效果；
- 与 DualPipe 相比, Interleaved 1F1B 大多数情况下展现出较接近的 bubble overhead。另外大多数训练场景下，micro batch 的数量都比较多，此时这两种方案下 bubble rate 在端到端训练中的耗时占比都非常低；
- 不需要存储一份额外的模型参数；
- 能很好地兼容 NVIDIA Megatron-Core 现有的代码结构和分布式策略；
NVIDIA 技术团队正与小红书技术团队一起将该优化方案实现到 Megatron-Core 中。目前，前期的 PoC 验证和 Design 已完成，预计 3 月中下旬可提供 EA 版本供试用。
