Nemo Mbridge Perf Moe Optimization Workflow

MoE Training Optimization Workflow

Stable docs: @docs/training/moe-optimization.md Card: @skills/nemo-mbridge-perf-moe-optimization-workflow/card.yaml Source: Scalable Training of MoE Models with Megatron Core

Quick Reference

Start with the paper's Three Walls:

• memory wall
• communication wall
• compute-efficiency wall

For operational diagnosis, split the compute-efficiency wall into compute and host/launch bottlenecks. They need different evidence and different fixes. MoE tuning is iterative, so use this order:

freeze the measurement contract -> fit -> scale -> profile -> retune -> validate

First Answer Checklist

For MoE optimization workflow prompts, present the response in this order:

1. Freeze the measurement contract: record the exact model and task, hardware and topology, container and commits, data and routing semantics, precision, sequence and batch shape, parallelism, graph scopes, and the steady-state metric window. Label each candidate as training-equivalent or benchmark-only.
2. Fit: make the model memory-feasible first. Use the smallest model parallelism that fits, prefer selective recompute before full recompute, add offloading only after recompute and parallelism are insufficient, and use --fake-init-process-group to sanity-check large layouts.
3. Scale: maximize DP after the model fits, keep hot communication inside the fastest interconnect, use PP plus VPP for multi-node scaling, prefer EP over extra TP for expert layers, and add CP when long context makes attention memory dominant.
4. Profile: identify the dominant wall: memory, communication, host overhead, or compute.
5. Retune: change one variable at a time based on the profiled bottleneck. Dispatcher, overlap, lower precision, CUDA graphs, and recompute are candidates, not hardware defaults.
6. Validate: use short matched screens to reject candidates, then run the winner for at least 50 steps. Verify the requested backend or graph replay actually ran, time a declared post-warmup window, and report loss health, skipped/NaN iterations, memory, step time, and model TFLOPS/GPU.
7. Include the exact Parallel Folding meshes: Attention: TP x CP x DP x PP and MoE: ETP x EP x EDP x PP.
8. Use alltoall for safe bring-up, then A/B flex + deepep and flex + hybridep when their packages and target topology support them. Start from BF16 and eager execution; introduce lower precision or the narrowest useful CUDA-graph scope only after profiling justifies it.

Phase 0: Freeze The Measurement Contract

A comparison is valid only when the following stay fixed unless they are the single variable under test:

• Bridge, MCore, Transformer Engine, container, CUDA, and NCCL versions
• GPU count, SKU, node topology, and launcher/environment settings
• model, task, data path, sequence length, MBS, GBS, and optimizer settings
• TP, PP, VPP, CP, EP, ETP, and DP layout
• routing semantics, precision, recompute, dispatcher, overlap, and graph scope
• warmup and steady-state timing windows

Separate two acceptance classes:

• Training-equivalent changes preserve the intended routing, loss, data, optimizer, and checkpoint/resume behavior.
• Benchmark-only changes such as forced load balancing are useful for controlled kernel studies, but cannot establish production-training correctness or convergence.

Phase 1: Make The Run Memory-Feasible

Start with a configuration that fits reliably before chasing throughput.

Recommended order:

1. Use the smallest amount of model parallelism that still fits.
2. Turn on selective recompute before falling back to full recompute.
3. Add offloading only when recompute and parallelism are still insufficient.
4. Use --fake-init-process-group to sanity-check large parallel layouts on a single GPU before burning cluster time.

Recompute guidance

Prefer selective recompute for MoE runs:

• good first choices: layernorm, core_attn, moe_act, mlp, or model-specific modules (shared_experts, mla_up_proj)
• use full recompute only when the run still does not fit
• revisit recompute after enabling CUDA graphs, because some graph scopes and full recompute paths do not mix well

As a rule of thumb, fine-grained recompute often recovers most of the needed memory while keeping throughput much closer to the non-recompute baseline than full-layer recompute does.

Phase 2: Choose Parallelism For Scale

Priority order:

1. Maximize DP once the model fits.
2. Keep the hot communication path inside the fast interconnect when possible.
3. Use PP, plus VPP if needed, for multi-node scaling.
4. Prefer EP over extra TP for expert layers.
5. Add CP for long context once sequence length makes attention memory dominant.

Parallel Folding

Parallel Folding decouples attention and MoE parallelism so you do not have to pick a single compromise layout:

Attention: TP × CP × DP × PP
MoE:       ETP × EP × EDP × PP

Key knobs:

• --expert-model-parallel-size
• --expert-tensor-parallel-size

Use it when attention prefers some TP or CP, but expert layers benefit from a larger EP degree than the dense layers can tolerate.

Phase 3: Profile The Dominant Bottleneck

Profile overlap without misreading it

Use unprofiled steady iterations for the acceptance metric and a matched profile for causal explanation:

1. Change one overlap or dispatcher variable at a time; keep routing, graph scopes, parallelism, batch shape, and runtime fixed.
2. Build interval unions for communication kernels and compute kernels, then measure their intersection to quantify hidden communication.
3. Do not add kernel durations and call the result wall time. Concurrent kernels may run longer because of SM or bandwidth contention even while the exposed GPU-active union and end-to-end step time fall.
4. Corroborate the trace with dispatch/combine NVTX ranges, steady step time, model TFLOPS/GPU, loss finiteness, skipped/NaN counts, and peak memory.

On a controlled 16×H100 Qwen3 30B-A3B HybridEP run, plain EP overlap increased communication hidden by GEMM/attention from 0.11% to 36.55%. The unprofiled step fell from 24.7138s to 20.9920s and throughput rose from 244.039 to 287.305 model TFLOPS/GPU. delay_wgrad_compute remained disabled.

Phase 4: Retune From Evidence

Choose the smallest candidate that targets the profiled bottleneck and change one variable at a time.