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）可通过将权重梯度计算推迟到下一层的前向计算期间进行重叠，带来额外的重叠收益。

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 分发和延迟权重梯度。测试的模型配置为：在 16 张 H100 GPU 上进行 Qwen3 MoE 30B-A3B SFT：EP=16、alltoall、BF16、全局批次大小 1024、禁用 CUDA 图、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 重叠是有益的。结果并未显示延迟权重梯度计算带来显著的独立收益，也未验证融合 MoE 排列路径，因为运行时堆栈中禁用了该路径。

启用

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 = True
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")

兼容性与限制

• 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 graph 交互。

最小可运行命令

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

分离普通 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"，且后端正确

成功标准

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

代码位置

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 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	检查 PyTorch 版本	升级到 >= 2.6.0
断言 virtual_pipeline_model_parallel_size	PP > 1，但未配置 VPP	检查 PP 和 VPP 配置	PP > 1 时设置 VPP
断言 recompute_granularity	启用了完整重计算	检查重计算设置	禁用完整重计算
断言 overlap_moe_expert_parallel_comm required	启用了延迟 wgrad，但未启用 EP 重叠	检查是否在未启用重叠时设置了 delay_wgrad_compute	先启用 EP 重叠
断言 gradient_accumulation_fusion	CUDA graph 与延迟 wgrad 同时启用	检查 graph 作用范围和 wgrad 设置	启用 gradient_accumulation_fusion
attention bias 相关断言	CUDA graph attention、延迟 wgrad 和 bias 同时启用	检查 add_bias_linear / add_qkv_bias	禁用 attention 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 重叠。
• 尚未在受控的 Bridge 实验中测量每个模型系列的端到端吞吐量提升。与普遍适用的性能结论相比，代码验证更为可靠。
• MoE 重叠和共享专家重叠不能同时使用。
• CUDA graph 与延迟 wgrad 同时使用是一条受多项约束的路径，需要仔细验证 TE 版本和作用范围。