| 名称 | nemo-mbridge-perf-expert-parallel-overlap |
| 描述 | 在Megatron-Bridge中验证和使用MoE专家并行通信重叠，包括overlap_moe_expert_parallel_comm、delay_wgrad_compute以及DeepEP和HybridEP等灵活调度器后端。 |
| 开源协议 | Apache-2.0 when_to_use: 启用EP重叠以隐藏调度/合并延迟，或将吞吐量回归追溯至EP重叠配置更改；overlap_moe_expert_parallel_comm、delay_wgrad_compute、flex dispatcher、DeepEP overlap、HybridEP overlap。 |
MoE专家并行(EP)重叠技能
参考
- 稳定文档: @docs/training/communication-overlap.md
- 结构化元数据: @skills/nemo-mbridge-perf-expert-parallel-overlap/card.yaml
是什么
专家并行(EP)重叠通过在专家FFN计算的同时间执行token分发/合并的全对全(all-to-all)通信来隐藏其成本。可选的延迟专家权重梯度计算（delay_wgrad_compute）通过将wgrad延迟到与下一层前向计算重叠，进一步提供重叠。
Bridge支持两种调度器路径：
| 调度器 | 后端 | 使用时机 |
|---|---|---|
| alltoall | 标准MoE all-to-all | 默认，兼容性最广 |
| flex | DeepEP 或 HybridEP | Ampere/Hopper/Blackwell上提供更高重叠 |
快速决策
使用EP重叠的条件：
- 模型是MoE且EP > 1
- 专家分发/合并通信占step时间的重要部分
- 有内存余量并正在调优吞吐量
优先选择：
- 首次部署使用alltoall调度器
- 在支持的GPU上争取额外收益时使用flex + DeepEP/HybridEP
避免情况：
- 启用全激活重计算
- 启用moe_shared_expert_overlap
- 运行还在做正确性验证
- PyTorch < 2.6.0
预期结果：
- 若all-to-all分发是明显瓶颈，重叠可带来适度到可观的加速
- 若运行体量小、通信少或受其它因素主导，收益可能可忽略
先正确性的alltoall基准
纯EP重叠隔离基准中关闭flex调度和延迟wgrad。测量形态为Qwen3 MoE 30B-A3B SFT，16块H100：EP=16、alltoall、BF16、全局batch size 1024、CUDA graphs禁用、moe_permute_fusion=false，在迭代3-8上测量。
使用以下覆盖项：
python –cuda_graph_impl none --moe_flex_dispatcher_backend None --moe_a2a_overlap false comm_overlap.overlap_moe_expert_parallel_comm=true comm_overlap.delay_wgrad_compute=false model.moe_shared_expert_overlap=false
不要在本隔离测试中使用 --moe_a2a_overlap true：性能辅助函数会同时启用overlap_moe_expert_parallel_comm和delay_wgrad_compute，因此无法隔离纯EP重叠。
稳态窗口计时：
| 场景 | 稳态均值 | 相对 |
|---|---|---|
| 无EP重叠 | 41.25s | 1.000x |
| EP重叠 | 31.31s | 1.317x |
| EP重叠加delay_wgrad_compute | 31.20s | 1.322x |
该证据表明在该跨节点all-to-all形状上应启用纯EP重叠；未显示延迟wgrad有独立收益，也未验证融合MoE置换。
HybridEP生产形状基准
2026-07-25受控的Qwen3 30B-A3B预训练对比验证了纯EP重叠在生产HybridEP路径上的效果：
硬件: 16xH100
精度: BF16
序列长度: 4096
并行: TP1/PP1/CP1/EP16
Batch: MBS1/GBS1024
路由: force balance
调度器: flex + HybridEP
CUDA graph: moe_router和moe_preprocess的Transformer Engine作用域
延迟wgrad: 禁用
| 场景 | 稳态窗口 | 步时间 | 模型TFLOPS/GPU |
|---|---|---|---|
| overlap关闭 | 迭代5-20 | 24.7138s | 244.039 |
| overlap开启（搜索运行） | 迭代5-20 | 21.0725s | 286.208 |
| overlap开启（独立验证） | 迭代41-50 | 20.9920s | 287.305 |
独立运行将步时间降低15.059%，吞吐量提高17.729%。损失有限，跳过和NaN迭代为零，rank-0峰值内存为62.166 GiB。Nsight对比显示启用overlap后通信与GEMM/attention并发时间从9.079ms（0.11%）增至3958.997ms（36.55%），GPU活跃区间并集从22.821s降至21.221s。
该证据证明机制有效，但不承诺普适加速。调度器、图形作用域、路由、并行、批形状和运行时均保持不变，仅改变纯EP重叠。
启用方法
alltoall调度器
python cfg.comm_overlap.overlap_moe_expert_parallel_comm = True cfg.comm_overlap.delay_wgrad_compute = False cfg.model.moe_shared_expert_overlap = False cfg.model.expert_model_parallel_size = 8 cfg.model.num_moe_experts = 64 cfg.model.moe_token_dispatcher_type = ‘alltoall’ cfg.model.bf16 = True cfg.model.fp16 = False
仅在纯重叠路径已验证且兼容性约束满足后，再启用delay_wgrad_compute=True。
flex调度器（DeepEP或HybridEP）
python from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend cfg.comm_overlap.overlap_moe_expert_parallel_comm = True cfg.comm_overlap.delay_wgrad_compute = False cfg.model.moe_shared_expert_overlap = False apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend=‘deepep’)
或：apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend=‘hybridep’)
先对纯EP重叠做基准，延迟wgrad仅在CUDA graph和TE兼容性约束满足后单独做A/B。
兼容性与约束
- expert_model_parallel_size > 1
- num_moe_experts > 1
- moe_token_dispatcher_type必须为’alltoall’或’flex’
- moe_shared_expert_overlap = False
- 基础精度为BF16或FP16
- PyTorch >= 2.6.0
- PP > 1时必须设置virtual_pipeline_model_parallel_size
- recompute_granularity != ‘full’，recompute_method = None，recompute_num_layers = None
- mtp_num_layers必须为None或1
- delay_wgrad_compute要求overlap_moe_expert_parallel_comm为前置
- delay_wgrad_compute与overlap_grad_reduce需TE >= 2.7.0
- delay_wgrad_compute与gradient_accumulation_fusion需TE >= 2.7.0
- CUDA graph的attn作用域 + delay_wgrad_compute需TE >= 2.12.0、gradient_accumulation_fusion=True且无attention bias
- DeepEP仅支持Ampere、Hopper、B200、B300 GPU
- HybridEP支持Ampere、Hopper、B200、B300以及带NVL72的GB200/GB300
最小可用配置
python cfg.comm_overlap.overlap_moe_expert_parallel_comm = True cfg.comm_overlap.delay_wgrad_compute = False cfg.model.expert_model_parallel_size = 4 cfg.model.num_moe_experts = 64 cfg.model.moe_token_dispatcher_type = ‘alltoall’ cfg.model.moe_shared_expert_overlap = False cfg.model.bf16 = True
这是正确性优先的起点。之后才添加延迟wgrad、flex调度和CUDA graph交互。
最小可运行命令
在Slurm分配内保持模型、并行、调度器和运行时固定，仅改变两个overlap覆盖项：
bash 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
分离纯EP重叠与延迟wgrad时不要用–moe_a2a_overlap true。单元测试验证：
bash uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py -k ‘moe’ tests/unit_tests/training/test_deepep.py -q
验证
单元测试
bash uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py tests/unit_tests/training/test_deepep.py -q
日志检查
EP overlap成功运行后：
- 确认CommOverlapConfig完成时无断言错误
- 确认overlap_moe_expert_parallel_comm在日志配置中为True
- 若使用flex调度器，则确认moe_token_dispatcher_type='flex’及正确后端
成功标准
- 所选调度器和overlap配置通过校验
- 训练无挂起或断言失败
- 吞吐量提升或至少不下降
- 损失轨迹与基线一致
Profile解读
使用未profile的稳态窗口判断吞吐量；用匹配的profile解释机理：
- 保持调度器、路由、graph作用域、批形状、并行布局和运行时固定
- 切换纯EP overlap时捕获相同rank和稳态迭代
- 构建通信与计算kernel的区间并集，取交集
- 不要将kernel持续时间和当作wall time
- 用分发/合并NVTX区间、最终step时间、损失有限性、跳过/NaN计数和峰值内存佐证
代码锚点
Bridge overlap验证（src/megatron/bridge/training/comm_overlap.py 470-505）
if self.user_comm_overlap_cfg.overlap_moe_expert_parallel_comm is True: assert model_cfg.expert_model_parallel_size > 1 assert model_cfg.num_moe_experts > 1 assert model_cfg.moe_token_dispatcher_type in [‘alltoall’, ‘flex’] assert model_cfg.bf16 or model_cfg.fp16 assert is_torch_min_version(‘2.6.0’) # PP+VPP检查、recompute检查、shared_expert_overlap检查…
延迟wgrad验证（src/megatron/bridge/training/comm_overlap.py 507-557）
if self.user_comm_overlap_cfg.delay_wgrad_compute is True: # TE版本检查，用于overlap_grad_reduce和gradient_accumulation_fusion # CUDA graph作用域校验 assert overlap_moe_expert_parallel_comm
flex调度器激活（src/megatron/bridge/training/flex_dispatcher_backend.py 27-72）
def apply_flex_dispatcher_backend(…): # GPU架构检查 model_config.moe_token_dispatcher_type = ‘flex’ model_config.moe_flex_dispatcher_backend = moe_flex_dispatcher_backend model_config.moe_shared_expert_overlap = False
性能辅助覆盖（scripts/performance/utils/overrides.py 149-156）
def _set_moe_a2a_overlap_overrides(recipe, moe_a2a_overlap=False): if moe_a2a_overlap: recipe.comm_overlap.overlap_moe_expert_parallel_comm = True recipe.comm_overlap.delay_wgrad_compute = True recipe.model.moe_shared_expert_overlap = False
测试覆盖
| 文件 | 覆盖范围 |
|---|---|
| tests/unit_tests/training/test_comm_overlap.py | EP overlap验证、延迟wgrad、CUDA graph+wgrad交互 |
| tests/unit_tests/training/test_deepep.py | DeepEP/HybridEP辅助激活和GPU门控 |
故障诊断
| 症状 | 可能原因 | 确认方式 | 修复 |
|---|---|---|---|
| 断言expert_model_parallel_size>1 | EP未配置 | 检查expert_model_parallel_size | 设置EP>1 |
| 断言moe_token_dispatcher_type | 调度器错误 | 检查调度器类型 | 使用alltoall或flex |
| BF16/FP16断言 | 精度错误 | 检查bf16和fp16 | 设置bf16=True |
| 训练挂起 | PyTorch<2.6 | 检查版本 | 升级至>=2.6.0 |
| 断言virtual_pipeline_model_parallel_size | PP>1但无VPP | 检查PP和VPP配置 | 设置VPP |
| 断言recompute_granularity | 启用了全重计算 | 检查recompute设置 | 关闭全重计算 |
| 断言overlap必须开启 | 延迟wgrad未搭配EP overlap | 检查delay_wgrad_compute | 先启用EP overlap |
| 断言gradient_accumulation_fusion | CUDA graph+延迟wgrad | 检查graph作用域和wgrad配置 | 启用gradient_accumulation_fusion |
| attention bias断言 | CUDA graph attn+wgrad+bias | 检查add_bias_linear/add_qkv_bias | 禁用attention bias |
| flex无收益 | 未调用apply_flex_dispatcher_backend | 日志中检查调度器类型 | 调用apply_flex_dispatcher_backend |
| DeepEP/HybridEP被跳过 | GPU不支持 | 查看警告 | 在支持GPU上运行 |
| 启用后kernel总时间增加 | 并发争用或回归 | 比较区间并集和真实step时间 | 根据暴露的wall time判断 |
已知限制
- 仅设置moe_flex_dispatcher_backend不会激活flex，必须调用apply_flex_dispatcher_backend(…)。
- 公开recipes通常保守，默认禁用MoE overlap。
- 受控端到端证据只针对一个Qwen3 30B-A3B/HybridEP/H100形状，推广前需重复A/B。
- MoE overlap与shared-expert overlap互斥。
- CUDA graph联合延迟wgrad是一条多约束路径，需谨慎验证TE版本和graph作用域。
上次签名刷新：2026-08-03。