nemo-mbridge-perf-moe-comm-overlap

MoE 通信重叠

更高层级的概述请参阅：

• @docs/training/communication-overlap.md
• @skills/nemo-mbridge-perf-moe-comm-overlap/card.yaml

快速决策

在以下情况下使用 MoE 通信重叠：

• EP > 1
• 配置文件中可见 token 分发或合并耗时
• 运行已正确无误，现在要调优吞吐量

不要在初期启动阶段就启用。等分发器、路由模式和重计算方案都已稳定后，再进行验证会更容易。

启用

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True

# Optional: delayed wgrad for additional overlap
cfg.comm_overlap.delay_wgrad_compute = True

# IMPORTANT: disable shared expert overlap when using dispatch overlap
cfg.model.moe_shared_expert_overlap = False

先决条件

• expert_model_parallel_size > 1
• num_moe_experts > 1
• moe_token_dispatcher_type 必须为 "alltoall" 或 "flex"
• 精度：BF16 或 FP16
• 如果使用 PP，则必须设置 VPP（virtual_pipeline_model_parallel_size）（不能为 None）

激活 Flex 分发器

仅设置 moe_flex_dispatcher_backend 不会激活 Flex 分发。 还必须设置 moe_token_dispatcher_type = "flex"。

与重计算和 CUDA Graph 的交互

• 完全重计算与重叠路径并不适配。
• 如果 CUDA Graph 作用域包含注意力或 MoE 路由器工作，delay_wgrad_compute 会带来更多限制。
• 实际使用中，启用重叠时搭配选择性重计算更稳妥。

测量证据

HybridEP 生产规模验证

2026-07-25 进行的一次受控 Qwen3 30B-A3B 预训练对比使用了 16 张 H100 GPU、BF16、序列长度 4096、TP=1、PP=1、CP=1、EP=16、MBS=1、GBS=1024、强制均衡路由、HybridEP，以及 Transformer Engine CUDA Graph 作用域 moe_router 和 moe_preprocess。唯一的性能改动是启用普通 EP 重叠；延迟 wgrad 保持禁用。

案例	稳态窗口	步骤耗时	模型 TFLOPS/GPU
关闭 EP 重叠	迭代 5-20	24.7138s	244.039
启用 EP 重叠，搜索运行	迭代 5-20	21.0725s	286.208
启用 EP 重叠，独立验证	迭代 41-50	20.9920s	287.305

独立验证结果相较于复现的基线，步骤耗时降低了 15.059%，吞吐量提高了 17.729%。损失值始终有限，没有跳过任何迭代，也没有出现 NaN；rank-0 峰值已分配内存为 62.166 GiB。

采用相同方法进行的 rank-0 Nsight Systems 对比，两种情况下各采集了 463,348 个内核：

配置文件指标	关闭重叠	启用重叠
与 GEMM/注意力并发的通信	9.079ms	3,958.997ms
被计算隐藏的通信耗时	0.11%	36.55%
GPU 活跃区间并集	22.821s	21.221s
HybridEP dispatch-with-permute NVTX	4.253s	1.767s
HybridEP metadata-preprocess NVTX	3.109s	0.670s

这是直接证据，表明性能提升来自隐藏原本暴露的 HybridEP 分发/合并工作，而非更改分发器、路由、图作用域、批次规模或并行布局。

优先保证正确性的 alltoall 冒烟测试

2026-05-18，在当前主分支上进行的一次 H100 x16 Qwen3 30B-A3B 模拟预训练冒烟测试使用了 EP=16、alltoall、全局批次大小 1024、禁用 CUDA Graph，并设置 moe_permute_fusion=false，因为在之前的启动调试中，PyTorch 25.11 / TE / Triton 技术栈在 Transformer Engine 融合置换中发生故障。

结果仅供参考，不代表发布级别的验证：

• 不启用 EP 重叠：迭代 3-8 的稳态平均耗时为 41.25s
• 启用 EP 重叠：迭代 3-8 的稳态平均耗时为 31.31s
• 启用 EP 重叠并设置 delay_wgrad_compute：迭代 3-8 的稳态平均耗时为 31.20s

这说明，当通信暴露在外时，EP 重叠可能有助于跨节点 alltoall MoE 配置。这并不能证明延迟 wgrad 本身带来了额外收益，也不能验证融合置换路径。2026-05-16 在相同配置上进行的较早一次短时冒烟测试也显示了相同趋势。

代码锚点

• 重叠验证：src/megatron/bridge/training/comm_overlap.py
• Flex 分发器后端：src/megatron/bridge/training/flex_dispatcher_backend.py
• 配置：src/megatron/bridge/training/config.py
• 单元测试：tests/unit_tests/training/test_comm_overlap.py
• DeepEP 测试：tests/unit_tests/training/test_deepep.py

注意事项

1. 共享专家重叠冲突：moe_shared_expert_overlap 和 overlap_moe_expert_parallel_comm 可能冲突。使用分发重叠路径时，请禁用共享专家 重叠。

2. 不使用 VPP 的 PP：启用流水线并行时，MoE 重叠需要 VPP。否则，重叠调度 无法正确交错执行。

3. Flex 不等于后端标志：如果 moe_token_dispatcher_type 仍为 "alltoall"， 单独设置 moe_flex_dispatcher_backend="deepep" 不会产生任何效果。

4. 保守的配方默认值：大多数公开配方默认禁用 MoE 重叠。你需要通过覆盖项 明确启用它。

5. 性能提升取决于工作负载：当分发通信已占据明显的步时比例时，重叠最有帮助。 但这并不保证能提升所有小规模或低负载的 EP 运行性能。

6. 累加的内核时间不等于实际耗时：并发内核可能因争用 SM 或带宽而运行更久， 因此，重叠可能增加每条流的累加内核时长，同时缩短暴露的时间区间并集以及端到端步时。

验证

初始化期间，请留意与重叠相关的日志消息。如果未满足前置条件，comm_overlap.py 中的 通信重叠验证会报错，因此顺利启动即可确认该功能已启用。

运行短时间的性能测试时，请明确指定命令格式，每次只调整一个重叠参数：

uv run python scripts/performance/run_script.py \
  -m qwen \
  -mr qwen3_30b_a3b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 16 \
  -gn 8 \
  --max_steps 8 \
  --cuda_graph_impl none \
  --moe_flex_dispatcher_backend None \
  --moe_a2a_overlap false \
  --tokenizer_type NullTokenizer \
  comm_overlap.overlap_moe_expert_parallel_comm=true \
  comm_overlap.delay_wgrad_compute=false \
  model.moe_shared_expert_overlap=false

如果启动时融合 MoE 置换失败，请添加 model.moe_permute_fusion=false，将重叠计时与运行时堆栈验证分开，然后使用匹配的生产容器重新测试。

验证性能时，请使用未经分析的稳定运行窗口作为验收指标。使用匹配的 Nsight A/B 测试确定因果关系：

1. 固定分发器、路由、CUDA 图、批次形状、并行方式和运行时。
2. 仅切换 overlap_moe_expert_parallel_comm；首次单独验证时，保持 delay_wgrad_compute=false。
3. 比较通信与计算时间区间的并集及其交集，而不只是累加的内核时长。
4. 报告稳定步时、每 GPU 的模型 TFLOPS、损失是否有限、跳过/NaN 迭代次数以及峰值已分配内存。

上次签名刷新：2026-08-03。