nemo-mbridge-perf-moe-optimization-workflow

Files (6)

• BENCHMARK.md
• SKILL.md
• card.yaml
• evals.json · evals/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

# 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](https://arxiv.org/abs/2603.07685)

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

```text
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:

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

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills