MoE Training Optimization#

This page covers the optimization framework, techniques, and best practices for training Mixture-of-Experts models with Megatron-Core. It is based on Scalable Training of MoE Models with Megatron Core.

For tuning knobs, hardware-specific configs, and benchmark-oriented guidance, see:

The Three Walls#

MoE training is constrained by three tightly coupled barriers:

Wall

Root Cause

Metric

Memory

All E experts stored but only K active per token

GB per GPU

Communication

EP all-to-all dispatches tokens across GPUs

% of step time

Compute Efficiency

Small expert GEMMs + host overhead

GPU SM utilization

These walls interact: solving one often exposes another. FP8 reduces memory and can improve GEMM throughput, but it can also shift more of the remaining cost to quantization kernels and host overhead. Communication overlap can hide latency, but it may add scheduling and buffer constraints. Effective optimization treats all three as a unified system.

The Optimization Workflow#

Use the full evidence-gated loop:

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

Phase 0: Freeze the Measurement Contract#

Before tuning, record and hold fixed the software/runtime versions, hardware and topology, model/task/data/routing semantics, precision, sequence and batch shape, parallelism, graph scopes, and steady-state timing window. Label forced routing, synthetic data, or disabled optimizer/checkpoint paths as benchmark-only; they cannot establish training-equivalent acceptance.

Phase 1: Establish Memory-Feasible Parallelism#

Memory is the hard gate — if it doesn’t fit, nothing runs.

Strategy

Activation

Weight

Optimizer

Comm

TP

1/d (with SP)

1/d

1/d

High

EP

~1 (load-dependent)

1/d (MoE only)

1/d

Medium

PP

1 (>1 with VPP)

1/d

1/d

Medium

CP

1/d

1

1/d†

Medium

DP

1

1

1/d†

Low

†Requires --use-distributed-optimizer.

Quick test: use --fake-init-process-group to emulate distributed training on a single GPU for rapid parallelism iteration before spending cluster time.

Phase 2: Select Optimal Parallelism#

Five guidelines, in order of priority:

  1. Minimize model parallelism, maximize DP. Model parallelism adds communication; use distributed optimizer to free memory for larger DP.

  2. Keep EP × TP within the fast interconnect domain. EP and TP are communication-intensive, so keeping the hot path on the fastest available links usually matters more than theoretical FLOPs.

  3. Use PP for multi-node scaling. PP’s point-to-point comms scale better across nodes than TP/EP. Enable VPP to reduce bubbles.

  4. Prefer EP over TP for expert layers. EP gives better GEMM efficiency, lower communication, and eliminates local permutation when EP = num_experts. Use Parallel Folding to decouple attention TP from expert EP.

  5. Enable CP once sequence length makes attention memory dominant. In practice that often starts around the 8K-class regime, but the exact point depends on model size and hardware. Use hierarchical CP (a2a+p2p) on NVL72-class systems when appropriate.

Phase 3: Profile the Bottleneck#

Profile the training run and operationally split the paper’s compute-efficiency wall into compute and host/launch bottlenecks:

Memory bottleneck — forced into full recompute or excessive parallelism:

Optimization

Overhead

Flag

FP8 training

Low

--fp8-format --fp8-recipe

Selective recompute

Low

--recompute-granularity --recompute-modules

Precision-aware optimizer

Low

--use-precision-aware-optimizer

Activation offloading

Medium

--fine-grained-activation-offloading

Optimizer offloading

Medium

--offload-optimizer-states

Communication bottleneck — profiling shows time in collectives:

Comm Type

Fix

DP grad/param

--overlap-grad-reduce --overlap-param-gather

TP

--tp-comm-overlap

EP dispatcher

--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend {deepep|hybridep}

EP all-to-all

--overlap-moe-expert-parallel-comm

PP send/recv

--pipeline-model-parallel-layout (flexible VPP)

CPU overhead — gaps between GPU kernels in Nsight traces:

Fix

Flag

Disable Python GC

--manual-gc --manual-gc-interval 10

CUDA Graphs

--cuda-graph-impl transformer_engine

Reduce kernel launches

Decrease TP or increase MBS

Compute inefficiency — low SM utilization despite no comm/CPU issues:

Fix

Flag

Grouped GEMM

--moe-grouped-gemm

Kernel fusions

--moe-router-fusion --moe-permute-fusion

FP8 precision

--fp8-format --fp8-recipe

Phase 4: Retune One Variable at a Time#

Use the profile to choose one candidate: dispatcher, one overlap mode, lower precision, CUDA graphs, recompute, or a parallelism adjustment. Keep the rest of the measurement contract fixed. Hardware support narrows the candidate set; it does not guarantee which backend or precision mode will improve end-to-end step time.

Phase 5: Validate the Complete Winner#

Use short post-warmup screens to reject weak candidates, then run the winner for at least 50 steps and report a declared steady window. Require runtime evidence that the requested dispatcher, lower-precision kernels, overlap path, and graph replay actually ran. Report finite loss, skipped/NaN iterations, peak memory, step time, model TFLOPS/GPU, and checkpoint/optimizer-state behavior when the production path requires them.

This process is iterative: solving one bottleneck often exposes another.

Parallel Folding#

Attention and MoE layers have conflicting optimal parallelisms. Parallel Folding decouples their configurations.

Attention layers: TP × CP × DP × PP
MoE layers:      ETP × EP × EDP × PP  (PP must match)

Key benefits:

  • Breaks EP ≤ DP constraint: EP can “fold” across TP×CP groups

  • Independent optimization: Attention uses high TP; MoE uses ETP=1

  • Fewer GPUs needed: CP=8 and EP=8 share the same 8 GPUs

  • NVLink-local comms: Both CP and EP stay in high-bandwidth domain

Example: 256 GPUs with attention TP=4, CP=2, DP=8, PP=4. Traditional: EP ≤ DP = 8. With folding: EP=64, ETP=1, EDP=1.

Memory Optimization Stack#

Ordered by overhead (lowest first):

  1. Memory-efficient permutation (zero overhead): Absorbs routing weights into activations before FC2, eliminating saved tensors for router backward.

    Standard: y = Σ p_i · W2_i · φ(W1_i · x) Memory-efficient: y = Σ W2_i · (p_i · φ(W1_i · x))

    Mathematically equivalent when experts have no bias. Eliminates saving each expert output for router backward — activation is recomputed from already-saved inputs.

  2. FP8/FP4 activations: Store linear-layer inputs in lower precision than BF16. This usually gives a modest but useful activation-memory reduction.

  3. Fine-grained recompute: Recompute only cheap operations such as LayerNorm, activation functions, or model-specific up-projection modules. This often recovers much of the needed memory while keeping overhead much lower than full-layer recompute.

  4. Fine-grained offloading: Module-level D2H/H2D with stream overlap. This can free a meaningful amount of memory at a small throughput cost and may allow a better parallelism layout that more than repays the offload overhead.

  5. Optimizer state offloading: Move optimizer states to CPU between steps. This is especially attractive on GB200-class systems, where the host-device path is strong enough to make the trade practical.

  6. FSDP for MoE: Dual DeviceMesh — primary mesh for attention, expert mesh for MoE. AllGather/ReduceScatter stay within small EDP groups. Zero-copy comms via NCCL User Buffer Registration.

FP8 Recipe Selection#

Recipe

Platform

Granularity

Role after BF16 is stable

Per-tensor FP8

Hopper/Blackwell

1 scale/tensor

Starting point

Blockwise FP8

Hopper

128×128 blocks

Candidate when the target stack supports the intended kernels

MXFP8

Blackwell

1×32 elements

High-priority candidate

NVFP4

Blackwell

16 elements, 2-level

Speed-first candidate after BF16/FP8 validation

Key rules:

  • Router stays in FP32 always

  • Embeddings, output layer, gradients, optimizer stay in original precision

  • Expert GEMMs are the primary quantization target

  • MXFP8 on Blackwell communicates params in BF16 (can’t save on AllGather)

  • NVFP4 requires Random Hadamard Transforms, 2D scaling, stochastic rounding

CUDA Graphs for MoE#

Two modes, different trade-offs:

Mode

What’s Captured

When to Use

Full CUDA Graphs

Entire fwd+bwd

Drop-and-pad MoE only

Partial (layer-wise)

router + moe_preprocess, optionally attn

Dropless-MoE candidate when host/launch overhead is visible

Partial CUDA graphs capture static components while leaving dynamic expert computation outside the graph. Start with the narrowest useful scope only after profiling host/launch overhead. Validate replay against eager on the same full stack: TE-scoped capture can be neutral or slightly slower when the workload is not launch-bound, and an earlier win can disappear after dispatcher, overlap, precision, or recompute changes.

For full CUDA Graphs on dropless MoE, three techniques are needed:

  • Device-initiated Grouped GEMM: Reads shapes from GPU memory. cuBLASLt (CUDA 13.1+) or cuteDSL with fused activation/quantization.

  • ECHO (Elastic Cloning for Hot Experts): Clones hot experts to underutilized ranks via bin-packing. Reduces load variance so worst-case buffer sizing is closer to actual.

  • Paged Stashing: Single worst-case tmp buffer shared across layers for computation; paged stashing buffer stores only actual tokens. Reduces memory from O(layers × worst_case) to O(worst_case + actual). 64 tokens per page, free list via circular buffer.

Flexible Asymmetric VPP#

PP layout string controls per-stage layer distribution:

--pipeline-model-parallel-layout "Et*3|(tt|)*29m|L"

  • E = embedding, t = transformer, m = MTP, L = loss, | = stage boundary

  • Balance workload: embedding + N dense layers ≈ fewer MoE layers

  • Place MTP and loss on dedicated stages for memory isolation

EP Communication Overlap#

For exact overlap constraints and verification guidance, see docs/training/communication-overlap.md and the related MoE overlap skills.

Two overlap patterns for 1F1B:

  1. Merged FWD-FWD / BWD-BWD: Same-type passes from two microbatches run in parallel. Costs 2× activation memory. Less overlap (fwd compute is half of bwd).

  2. Merged FWD-BWD (preferred): Forward of microbatch i+1 overlaps with backward of microbatch i. No extra memory. Matches DualPipe design. Limited: first FWD and last BWD can’t be hidden.

Key optimization: W/D split. Split backward MLP work into weight-gradient and data-gradient pieces so the weight-gradient portion can overlap with forward compute when forward MLP alone is too short to hide dispatch cost.

MoE Token Dispatchers#

MoE models route tokens to experts via all-to-all communication. The dispatcher backend controls how this communication is implemented:

Dispatcher

Backend

Mechanism

alltoall

Standard MoE

Torch-native all-to-all collectives

flex + DeepEP

DeepEP library

Low-latency SM-based dispatch with GPU-side routing

flex + HybridEP

HybridEP library

Fused intra-node NVLink + inter-node IB dispatch

Hardware-informed candidate order#

Hardware

Bring-up

Tuned candidates

H100 / B200 (NVL8)

alltoall

A/B DeepEP and HybridEP when supported

GB200 / GB300 (NVL72)

alltoall

HybridEP first, then compare other installed backends

Topology and EP degree affect the candidate order, but they do not determine the winner. The canonical 16×H100 Qwen3 30B recipe uses HybridEP, while the current 256×H100 Qwen3 235B recipe uses standard alltoall plus overlap.

Treat dispatcher package availability as part of the experiment setup, not as a given. alltoall is the correctness fallback and should be the first smoke test on a new container. DeepEP and HybridEP require their corresponding runtime packages to be installed; otherwise the config can select moe_token_dispatcher_type="flex" and still fail during model construction. Keep the container, CUDA graph scope, routing mode, and MoE kernel-fusion flags fixed when comparing dispatcher throughput.

Short-run H100 sanity check#

The current canonical 16×H100 Qwen3 30B-A3B BF16 performance recipe uses HybridEP with 32 SMs, 64-token combine chunks, plain EP overlap, delayed weight-gradient compute disabled, and TE graphs over moe_router and moe_preprocess. Its verified 50-step run averaged 20.14729 seconds and 299.352 model TFLOPS/GPU over steps 41–50. This is a workload-specific result, not a universal H100 dispatcher rule.

On 2026-05-17, a 16-GPU H100 smoke run of Qwen3 30B A3B BF16 with EP=16 and the recipe’s Transformer Engine CUDA graph scopes (moe_router, moe_preprocess) completed with alltoall after disabling moe_permute_fusion for a Triton JIT compatibility issue in the container. The five-step run had a 45.65 s mean step time after the first warmup step, 132.9 mean TFLOP/s/GPU after warmup, final loss 11.44050, and 61.351 GB peak max allocated memory. In the same container, DeepEP and HybridEP selected the requested flex backend in the dumped config but failed before iteration 1 because the DeepEP/HybridEP packages were not installed. Use this as an availability caveat, not as evidence that alltoall is faster than flex dispatchers on H100.

Long-Context MoE Training#

At long sequences (64K+), SDPA dominates FLOPs. Context parallelism (CP) is the primary mechanism for scaling sequence length.

CP sizing rules of thumb:

  1. Start with CP ≈ seq_len / 4096: then round to a practical layout.

  2. Keep DP ≥ 1: CP × EP × TP × PP must not exceed total GPUs.

  3. Prefer selective recompute over full: Recompute up_proj, norm, moe, mlp rather than full recompute for better throughput.

  4. TP can sometimes substitute for some CP on NVLink systems: on NVL72 systems, higher TP can be competitive with a more CP-heavy plan.

  5. Optimizer CPU offload is often critical at long context because activation pressure consumes so much of the memory budget.

Long-context recommendations:

  • Keep sub-sequence length ~4096–8192 per CP/TP shard

  • Don’t recompute SDPA at long context: SDPA recompute adds significant compute overhead while saving relatively little memory. Recompute non-SDPA modules instead.

  • TP preferred within node (fast comms, reduces param memory)

  • P2P CP preferred across nodes (natural overlap with attention)

  • a2a CP + TP within node when ring exchange is undesirable

Dynamic Context Parallelism#

For variable-length training (RL, SFT):

  • Per-microbatch CP sizing instead of static CP for all

  • Pre-constructs multiple CP groups during init (powers of 2)

  • Scheduler selects effective cp_size per microbatch

  • Works with packed sequences (THD format)

MoE VLM Training#

MoE vision-language models combine a vision encoder with a MoE language decoder. Training requires choosing between two strategies:

Approach

Mechanism

When to Use

FSDP

Shards params, grads, and optimizer across all GPUs

Simpler setup and a better first bring-up path

3D Parallel

TP + PP + EP + DP

Higher throughput ceiling once the multimodal path is already stable

Key principles:

  • Always benchmark with real vision data — image-free mock runs can significantly overestimate throughput.

  • Freezing vision encoder saves compute when fine-tuning only the decoder.

  • MBS is critical for 3D-parallel VLM — larger micro-batch sizes often matter more than they do for text-only MoE.

  • FSDP is simpler and often competitive for initial bring-up.

Production Features Summary#

Feature

Purpose

Force-balance routing

Even token distribution for disclosed benchmark-only comparisons; validate natural routing separately

Aux-loss-free balancing

Learnable expert bias; adapts over time

Shared expert overlap

Hides shared expert latency behind dispatch/combine

LatentMoE

Reduces comms and per-expert params by compression ratio α

Distributed checkpoint

Parallelism-agnostic save/load with automatic resharding

Upcycling

Convert dense checkpoint to MoE without retraining

MTP

Multi-token prediction with flexible VPP placement

Muon optimizer

Matrix-aware updates; fewer steps than AdamW