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: keep dense TP as low as capacity allows while using a large legal EP for expert weights. 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.