Agent skill · NVIDIA
nemo-mbridge-perf-moe-optimization-workflow
Evidence-gated workflow for MoE performance optimization in Megatron Bridge. Covers measurement contracts, the Three Walls framework, parallel folding, profiling, matched A/B tuning, and final validation.
What it needs
About 7k tokens when loaded.
What this skill does
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: 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. …
How to use it
Reference it in AdaL, Claude Code, Cursor or any coding agent — nothing to install:
@skills NVIDIA/nemo-mbridge-perf-moe-optimization-workflow