mcpField
N
nemo-mbridge-perf-moe-comm-overlap

nemo-mbridge-perf-moe-comm-overlap

github.com/promptingcompany开发039 次安装17 周前更新
代码仓库

关于 nemo-mbridge-perf-moe-comm-overlap

介绍 Megatron Bridge 中 MoE 通信重叠的训练调优技能，面向已验证训练正确性并希望优化吞吐量的开发者。说明启用前提、配置冲突、验证方式及性能结果的适用范围。

主要功能

  1. 01说明 MoE 专家并行通信重叠
  2. 02列出启用前提与配置要求
  3. 03说明重计算和 CUDA graph 交互
  4. 04标注共享专家重叠冲突
  5. 05提供性能测试与验证建议

适用场景

  1. 01优化通信耗时明显的 MoE 训练
  2. 02配置专家并行分发重叠
  3. 03排查流水线并行与重叠配置问题

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/promptingcompany/nv-skills --skill nemo-mbridge-perf-moe-comm-overlap

安装后在 Claude Code 中输入/nemo-mbridge-perf-moe-comm-overlap确认已加载。

包含的文件

除 SKILL.md 外，安装时一起下载的参考资料和脚本。

BENCHMARK.md
card.yaml
skill-card.md
skill.oms.sig
evals/
evals.json

详细介绍

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 会带来更多限制。
  • 实际使用中，启用重叠时选择性重计算是更稳妥的搭配。

短时测试的测量注意事项

2026-05-18，在当前 main 分支上使用 H100 x16 对 Qwen3 30B-A3B 进行模拟预训练冒烟测试时，使用了 EP=16、alltoall、全局批次大小 1024，并禁用了 CUDA Graph。由于此前启动时 PyTorch 25.11 / TE / Triton 技术栈在 Transformer Engine 融合置换中出错，因此设置了 moe_permute_fusion=false。

结果仅供参考，未达到发布级别：

  • 无 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. PP 未设置 VPP：启用流水线并行时，MoE 重叠需要 VPP。否则，重叠调度无法正确交错执行。

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

  4. 配方默认值较保守：大多数公开配方都默认禁用 MoE 重叠。你需要通过 overrides 显式启用。

  5. 性能提升取决于工作负载：当分发通信已占据明显的单步耗时比例时，重叠最有帮助。它不保证能改善每种小型或低负载 EP 运行的性能。

验证

初始化期间，请留意与重叠相关的日志消息。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 \
  --config_variant v1 \
  --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，将重叠计时与运行时技术栈验证分开，然后使用匹配的生产容器重新测试。

评论

支持 Markdown