/nemo-mbridge-perf-expert-parallel-overlap

name: nemo-mbridge-perf-expert-parallel-overlap
description: Validate and use MoE expert-parallel communication overlap in Megatron-Bridge, including overlap_moe_expert_parallel_comm, delay_wgrad_compute, and flex dispatcher backends such as DeepEP and HybridEP.
license: Apache-2.0
when_to_use: Enabling EP overlap to hide dispatch/combine latency, or tracing a throughput regression to an EP overlap config change; 'overlap_moe_expert_parallel_comm', 'delay_wgrad_compute', 'flex dispatcher', 'DeepEP overlap', 'HybridEP overlap'.

MoE Expert-Parallel Overlap Skill

References

• Stable docs: @docs/training/communication-overlap.md
• Structured metadata: @skills/nemo-mbridge-perf-expert-parallel-overlap/card.yaml

What It Is

Expert-parallel (EP) overlap hides the cost of token dispatch/combine all-to-all communication by running it concurrently with expert FFN compute. Optionally, delayed expert weight-gradient computation (`delay_wgrad_compute`) provides additional overlap by deferring wgrad to overlap with the next layer's forward.

Bridge supports two dispatcher paths:

| Dispatcher | Backend | When to use | |---|---|---| | `alltoall` | Standard MoE all-to-all | Default, broadest compatibility | | `flex` | DeepEP or HybridEP | Higher overlap on Ampere/Hopper/Blackwell |

Quick Decision

Use EP overlap when:

• the model is MoE with `EP > 1`
• expert dispatch/combine communication is a meaningful part of step time
• you have memory headroom and are tuning for throughput

Prefer:

• `alltoall` dispatcher for the first rollout (broader compatibility)
• `flex` + DeepEP/HybridEP when running on supported GPUs and seeking

additional gains

Avoid EP overlap when:

• full activation recompute is enabled
• `moe_shared_expert_overlap` is enabled
• the run is still being brought up for correctness
• PyTorch < 2.6.0

Expected outcome:

• if all-to-all dispatch is a clear profile bottleneck, overlap can produce a

modest to meaningful speedup

• if the run is tiny, communication-light, or dominated by another wall, the

gain may be negligible

Correctness-First alltoall Benchmark

For the plain EP-overlap isolation benchmark, keep flex dispatch and delayed wgrad disabled. The measured shape was Qwen3 MoE 30B-A3B SFT on 16 H100 GPUs: `EP=16`, `alltoall`, BF16, global batch size 1024, CUDA graphs disabled, `moe_permute_fusion=false`, measured over iterations 3-8.

Use these overrides for the plain-overlap case:

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

Do not use `--moe_a2a_overlap true` for this isolation test: the performance harness helper enables both `overlap_moe_expert_parallel_comm` and `delay_wgrad_compute`, so it does not isolate plain EP overlap.

Steady-window timing from that benchmark:

| Case | Steady mean | Relative | |---|---:|---:| | no EP overlap | 41.25s | 1.000x | | EP overlap | 31.31s | 1.317x | | EP overlap plus `delay_wgrad_compute` | 31.20s | 1.322x |

This is evidence for enabling plain EP overlap on this inter-node all-to-all shape. It does not show a meaningful independent win from delayed wgrad, and it does not validate fused MoE permutation because that path was disabled for the runtime stack.

Enablement

alltoall dispatcher

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

Enable `delay_wgrad_compute=True` only after the plain overlap path is known to work and its extra compatibility constraints have been checked.

flex dispatcher (DeepEP or 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")

Compatibility And Constraints

• `expert_model_parallel_size > 1`
• `num_moe_experts > 1`
• `moe_token_dispatcher_type` must be `"alltoall"` or `"flex"`
• `moe_shared_expert_overlap = False`
• Base precision is BF16 or FP16
• PyTorch `>= 2.6.0`
• If `PP > 1`, `virtual_pipeline_model_parallel_size` must be set
• `recompute_granularity != "full"`, `recompute_method = None`,

`recompute_num_layers = None`

• `mtp_num_layers` must be `None` or `1`
• `delay_wgrad_compute` requires `overlap_moe_expert_parallel_comm` as a

prerequisite

• `delay_wgrad_compute` with `overlap_grad_reduce` requires TE >= 2.7.0
• `delay_wgrad_compute` with `gradient_accumulation_fusion` requires TE >= 2.7.0
• CUDA graph `attn` scope + `delay_wgrad_compute` requires TE >= 2.12.0,

`gradient_accumulation_fusion = True`, and no attention bias

• DeepEP: Ampere, Hopper, B200, B300 GPUs only
• HybridEP: Ampere, Hopper, B200, B300, GB200/GB300 with NVL72

Minimal Working Config

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

Use this as the correctness-first starting point. Add delayed wgrad, flex dispatch, and CUDA-graph interactions only after the plain overlap path is known to work.

Minimal Runnable Command

Performance harness example inside a Slurm allocation. Keep the model, parallelism, dispatcher, and runtime fixed, and vary only the two overlap overrides:

uv run python