nemo-mbridge-perf-expert-parallel-overlap

MoE 专家并行重叠技能

参考资料

• 稳定文档：@docs/training/communication-overlap.md
• 结构化元数据：@skills/nemo-mbridge-perf-expert-parallel-overlap/card.yaml

这是什么

专家并行（EP）重叠通过让 token 分发/合并的全互换通信与专家 FFN 计算并发执行，隐藏其开销。可选地，延迟专家权重梯度计算（delay_wgrad_compute）可通过推迟 wgrad，使其与下一层的前向计算重叠，从而实现额外的重叠。

Bridge 支持两种分发器路径：

分发器	后端	适用场景
alltoall	标准 MoE 全互换	默认选项，兼容性最广
flex	DeepEP 或 HybridEP	在 Ampere/Hopper/Blackwell 上实现更高的重叠度

快速决策

在以下情况下使用 EP 重叠：

• 模型为 MoE，且 EP > 1
• 专家分发/合并通信占用了显著的步时
• 有足够的显存余量，并且正在针对吞吐量进行调优

优先选择：

• 首次部署时使用 alltoall 分发器（兼容性更广）
• 在支持的 GPU 上运行并寻求额外收益时，使用 flex + DeepEP/HybridEP

在以下情况下避免使用 EP 重叠：

• 已启用完整激活重计算
• 已启用 moe_shared_expert_overlap
• 运行仍在进行正确性验证
• PyTorch < 2.6.0

预期结果：

• 如果全互换分发是明显的性能分析瓶颈，重叠可带来一定到显著的加速
• 如果运行规模很小、通信量很少，或受其他因素严重限制，收益可能微乎其微

以正确性为先的 alltoall 基准测试

对于纯 EP 重叠隔离基准测试，请保持 flex 分发和延迟 wgrad 处于禁用状态。测试使用的模型配置为：在 16 张 H100 GPU 上对 Qwen3 MoE 30B-A3B 进行 SFT：EP=16、alltoall、BF16、全局批次大小为 1024、禁用 CUDA graphs、moe_permute_fusion=false，测量迭代 3-8。

纯重叠场景使用以下覆盖参数：

--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

这为在此节点间全互换场景中启用纯 EP 重叠提供了依据。但这并未显示延迟 wgrad 有显著的独立收益，也未验证融合 MoE 排列路径，因为该路径在此运行时栈中处于禁用状态。

HybridEP 生产场景基准测试

2026-07-25 对 Qwen3 30B-A3B 预训练进行的受控对比，使用生产环境的 HybridEP 路径验证了纯 EP 重叠：

Hardware: 16×H100
Precision: BF16
Sequence: 4096
Parallelism: TP1 / PP1 / CP1 / EP16
Batch: MBS1 / GBS1024
Routing: force balance
Dispatcher: flex + HybridEP
CUDA graph: Transformer Engine scopes moe_router + moe_preprocess
Delayed wgrad: disabled

场景	稳定窗口	步时	每 GPU 模型 TFLOPS
关闭重叠	迭代 5-20	24.7138s	244.039
开启重叠，搜索运行	迭代 5-20	21.0725s	286.208
开启重叠，独立验证	迭代 41-50	20.9920s	287.305

与复现的基线相比，独立运行将步时缩短了 15.059%，吞吐量提高了 17.729%。损失值为有限值，跳过的迭代和 NaN 迭代均为零，rank-0 峰值已分配显存为 62.166 GiB。

匹配的 Nsight Systems 对比捕获了每种场景下相同的 463,348 个 rank-0 内核。启用重叠后，与 GEMM 和注意力计算并发的通信量从 9.079ms（占通信时间的 0.11%）增加到 3,958.997ms（36.55%）。GPU 活跃区间的并集从 22.821s 降至 21.221s。

应将此结果作为该机制有效的依据，而非普遍适用的加速承诺。分发器、图范围、路由、并行配置、批次形状和运行时均保持不变，只有纯 EP 重叠发生了变化。

启用

alltoall 分发器

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）

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")
# or: apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="hybridep")

先对普通 EP 重叠进行基准测试。只有在满足 CUDA 图和 TE 兼容性约束后，才将 delay_wgrad_compute=True 作为单独的后续 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 图 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

最小可用配置

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 图交互。

最小可运行命令

在 Slurm 分配中使用的性能测试框架示例。保持模型、并行策略、分发器和运行时不变，仅调整两个重叠覆盖参数：

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：性能测试框架的辅助函数会同时启用 overlap_moe_expert_parallel_comm 和 delay_wgrad_compute。

单元测试验证：

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py -k "moe" \
  tests/unit_tests/training/test_deepep.py -q

验证

单元测试

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py \
  tests/unit_tests/training/test_deepep.py -q

日志检查

启用 EP 重叠并成功运行后：

1. 确认 CommOverlapConfig 最终化期间没有断言错误
2. 确认日志配置中 overlap_moe_expert_parallel_comm 显示为 True
3. 如果使用 flex 分发器，确认日志中的 moe_token_dispatcher_type = "flex" 且后端正确

成功标准

• 所选分发器和重叠设置通过配置验证
• 训练运行完成，且没有挂起或断言失败
• 目标工作负载的吞吐量有所提升，或至少没有下降
• 损失变化轨迹与基线一致（重叠不应影响收敛）

性能分析解读

使用未进行性能分析的稳定时间窗口作为吞吐量验收结果。使用匹配的性能分析来解释其机制：

1. 固定分发器、路由、图范围、批次形状、并行布局和运行时。
2. 保持 rank 和稳定迭代相同，仅切换普通 EP 重叠并进行采集。
3. 分别计算通信与计算内核的时间区间并集，然后测量两者的交集。
4. 不要将内核时长总和当作墙钟时间。即使暴露时间缩短，并发内核也可能因 SM 或带宽竞争而运行更久。
5. 使用 dispatch/combine NVTX 范围、最终步时长、损失是否有限、跳过/NaN 计数和峰值内存来佐证时间区间结果。

代码定位

Bridge 重叠验证

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 check, recompute checks, shared_expert_overlap check ...

延迟 wgrad 验证

if self.user_comm_overlap_cfg.delay_wgrad_compute is True:
    # TE version checks for overlap_grad_reduce and gradient_accumulation_fusion
    # CUDA graph scope validations for delayed wgrad
    assert overlap_moe_expert_parallel_comm, ...

启用 flex 分发器

def apply_flex_dispatcher_backend(...):
    # GPU architecture check for DeepEP / HybridEP
    model_config.moe_token_dispatcher_type = "flex"
    model_config.moe_flex_dispatcher_backend = moe_flex_dispatcher_backend
    model_config.moe_shared_expert_overlap = False

性能测试框架覆盖参数

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 重叠验证、延迟 wgrad、CUDA 图与 wgrad 交互
tests/unit_tests/training/test_deepep.py	DeepEP/HybridEP 辅助函数启用和 GPU 限制

故障诊断

症状	可能原因	如何确认	修复方法
assert expert_model_parallel_size > 1	未配置 EP	检查 expert_model_parallel_size	将 EP 设为 > 1
assert moe_token_dispatcher_type	分发器错误	检查分发器类型	使用 "alltoall" 或 "flex"
BF16/FP16 相关 assert	精度设置错误	检查 bf16 和 fp16	设置 bf16 = True
训练期间挂起	PyTorch < 2.6	检查 PyTorch 版本	升级到 >= 2.6.0
assert virtual_pipeline_model_parallel_size	PP > 1，但未启用 VPP	检查 PP 和 VPP 配置	PP > 1 时设置 VPP
assert recompute_granularity	已启用完全重计算	检查重计算设置	禁用完全重计算
assert overlap_moe_expert_parallel_comm required	延迟 wgrad，但未启用 EP 重叠	检查是否启用了 delay_wgrad_compute，但未启用重叠	先启用 EP 重叠
assert gradient_accumulation_fusion	CUDA graph + 延迟 wgrad	检查图作用域和 wgrad 设置	启用 gradient_accumulation_fusion
注意力偏置相关 assert	CUDA graph 注意力 + 延迟 wgrad + 偏置	检查 add_bias_linear / add_qkv_bias	禁用注意力偏置
使用 flex 分发器后吞吐量未提升	未调用 apply_flex_dispatcher_backend	检查日志中的 moe_token_dispatcher_type	调用 apply_flex_dispatcher_backend(...)
DeepEP/HybridEP 被静默跳过	GPU 不受支持	检查警告日志	在 Ampere/Hopper/Blackwell 上运行
启用重叠后内核总耗时增加	预期的并发争用或回归	比较区间并集、通信/计算交集以及未进行性能分析的步耗时	根据暴露的墙钟时间判断重叠效果，而非各流时长之和

已知限制

• 仅设置 moe_flex_dispatcher_backend 不会启用 flex 分发 — 还必须调用 apply_flex_dispatcher_backend(...)。
• 公开的配置方案通常较为保守，默认不启用 MoE 重叠。
• 目前仅有一个 Qwen3 30B-A3B HybridEP H100 配置的受控端到端和性能分析证据；推广到其他 模型、分发器、拓扑、精度或批次形状前，请先重复进行匹配的 A/B 测试。
• MoE 重叠和共享专家重叠不能同时使用。
• CUDA graph 与延迟 wgrad 组合涉及多项约束，需要仔细验证 TE 版本和作用域。

最后一次签名刷新：2026-08-03。