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 (delaywgradcompute) provides additional overlap by deferring wgrad to overlap with the next layer's forward.

Bridge supports two dispatcher paths:

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
• moesharedexpert_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, moepermutefusion=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 --moea2aoverlap true for this isolation test: the performance harness helper enables both overlapmoeexpertparallelcomm and delaywgradcompute, so it does not isolate plain EP overlap.

Steady-window timing from that benchmark:

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.

HybridEP Production-Shape Benchmark

A 2026-07-25 controlled Qwen3 30B-A3B pretraining comparison validated plain EP overlap with the production HybridEP path:

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

The independent run reduced step time by 15.059% and raised throughput by 17.729% over the reproduced baseline. Loss was finite, skipped and NaN iterations remained zero, and rank-0 peak allocated memory was 62.166 GiB.

A matched Nsight Systems comparison captured the same 463,348 rank-0 kernels per case. Enabling overlap increased communication concurrent with GEMM and attention from 9.079ms (0.11% of communication time) to 3,958.997ms (36.55%). GPU-active interval union fell from 22.821s to 21.221s.

Use this as evidence for the mechanism, not as a universal speedup promise. The dispatcher, graph scopes, routing, parallelism, batch shape, and runtime were held fixed while only plain EP overlap changed.

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

Benchmark plain EP overlap first. Enable delaywgradcompute=True only as a separate follow-up A/B after its CUDA-graph and TE compatibility constraints are satisfied.

Compatibility And Constraints

• expertmodelparallel_size > 1
• nummoeexperts > 1
• moetokendispatcher_type must be "alltoall" or "flex"
• moesharedexpert_overlap = False
• Base precision is BF16 or FP16
• PyTorch >= 2.6.0
• If PP > 1, virtualpipelinemodelparallelsize must be set
• recomputegranularity != "full", recomputemethod = None,

recomputenumlayers = None

• mtpnumlayers must be None or 1
• delaywgradcompute requires overlapmoeexpertparallelcomm as a

prerequisite

• delaywgradcompute with overlapgradreduce requires TE >= 2.7.0
• delaywgradcompute with gradientaccumulationfusion requires TE >= 2.7.0
• CUDA graph attn scope + delaywgradcompute requires TE >= 2.12.0,

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

Do not use --moea2aoverlap true when separating plain EP overlap from delayed wgrad: the performance harness helper enables both overlapmoeexpertparallelcomm and delaywgradcompute.

Unit test verification:

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

Verification

Unit tests

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

Log checks

After a successful run with EP overlap:

1. Confirm no assertion errors during CommOverlapConfig finalization
2. Confirm overlapmoeexpertparallelcomm appears as True in the logged

config

1. If using flex dispatcher, confirm moetokendispatcher_type = "flex" and

the correct backend in logs

Success criteria

• Config validation passes for the selected dispatcher and overlap settings
• Training runs complete without hangs or assertion failures
• Throughput improves or at least does not regress for the target workload
• Loss trajectory matches baseline (overlap should not affect convergence)

Profile interpretation

Use an unprofiled steady window for the throughput acceptance result. Use a matched profile to explain the mechanism:

1. Keep the dispatcher, routing, graph scopes, batch shape, parallel layout,

and runtime fixed.

1. Capture the same rank and steady iteration while toggling only plain EP

overlap.

1. Build interval unions for communication and compute kernels, then measure

their intersection.

1. Do not use summed kernel duration as wall time. Concurrent kernels can run

longer under SM or bandwidth contention even when exposed time decreases.

1. Corroborate interval results with dispatch/combine NVTX ranges, final step

time, loss finiteness, skipped/NaN counts, and peak memory.

Code Anchors

Bridge overlap validation

```470:505:src/megatron/bridge/training/commoverlap.py if self.usercommoverlapcfg.overlapmoeexpertparallelcomm is True: assert modelcfg.expertmodelparallelsize > 1, ... assert modelcfg.nummoeexperts > 1, ... assert modelcfg.moetokendispatchertype in ["alltoall", "flex"], ... assert modelcfg.bf16 or modelcfg.fp16, ... assert istorchminversion("2.6.0"), ... # ... PP + VPP check, recompute checks, sharedexpertoverlap check ...

### Delayed wgrad validation

```507:557:src/megatron/bridge/training/comm_overlap.py
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-dispatcher activation

```27:72:src/megatron/bridge/training/flexdispatcherbackend.py def applyflexdispatcherbackend(...): # GPU architecture check for DeepEP / HybridEP modelconfig.moetokendispatchertype = "flex" modelconfig.moeflexdispatcherbackend = moeflexdispatcherbackend modelconfig.moesharedexpertoverlap = False

### Perf harness override

```149:156:scripts/performance/utils/overrides.py
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

Failure Diagnosis

Known Limitations

• Setting moeflexdispatcher_backend alone does not activate flex dispatch —

you must call applyflexdispatcher_backend(...).

• Public recipes are often conservative and leave MoE overlap disabled by

default.

• Controlled end-to-end and profile evidence exists for one Qwen3 30B-A3B

HybridEP H100 shape; repeat the matched A/B before generalizing it to another model, dispatcher, topology, precision, or batch shape.

• MoE overlap and shared-expert overlap are mutually exclusive.
• CUDA graph plus delayed wgrad is a multi-constraint path that requires

careful TE version and scope validation.

Last signature refresh: 2026-08-03.