Scalable Training of Mixture-of-Experts Models with Megatron Core

TL;DR

NVIDIA’s systems report on Megatron-Core MoE — the production training stack used to train trillion-parameter-class MoE models on thousands of GPUs. The central framing is that MoE sparsity breaks dense-model parallelism assumptions in three coupled ways — the Memory Wall, Communication Wall, and Compute Efficiency Wall — which Megatron-Core addresses through Expert Parallelism + Parallel Folding (decoupling attention and MoE layer parallelism), DeepEP/HybridEP dispatchers, Grouped GEMM, FP8/FP4 selective precision, and CUDA-Graph-friendly sync-free execution. Reported headline throughput is 1233/1048 TFLOPS/GPU on GB300/GB200 for DeepSeek-V3 (685B total / 37B active) and 974/919 TFLOPS/GPU for Qwen3-235B. Less of a research paper than a comprehensive design document for MoE-at-scale, with a DeepSeek-V3 case study showing how the optimizations compose.

Key claims

MoE sparsity creates a fundamental parameter-compute mismatch: DeepSeek-V3 has 18× more total parameters (685B) than active per-token (37B), so memory scales with total parameters but per-token FLOPs scale only with active parameters, leaving naive parallelism communication-bound [§1.2, Fig. 3].
The three coupled bottlenecks — Memory Wall, Communication Wall, Compute Efficiency Wall — cannot be optimized independently because relaxing one shifts pressure to another (e.g. larger batch helps GEMM utilization but inflates activation memory and all-to-all volume) [§1.2].
Expert Parallelism (EP) is necessary because the alternative — sharding expert matrices via Tensor Parallelism — fragments already-small fine-grained-expert GEMMs, making them less efficient; EP keeps full-size expert GEMMs and introduces all-to-all dispatch instead [§3.2].
The dense-sparse mismatch between attention (which prefers TP/CP) and MoE layers (which prefer EP with ETP=1) is resolved by MoE Parallel Folding, which decouples the two layer types’ parallelism configurations rather than forcing a single global TP/DP/PP grid [§3.3, §2.2.1].
Unoptimized all-to-all can consume up to 60% of training time in DeepSeek-V3 because experts span multiple nodes and the communication volume drops from intra-node NVLink to ~10× narrower inter-node interconnect [§1.2].
MoE breaks GEMM dominance: in Llama-3 405B (dense), GEMMs are ~70% of execution time; in DeepSeek-V3 (MoE), GEMMs are <50%, with routing/permutation adding ~9% layer-time and the remainder going to operations that scale with tensor count rather than FLOPs [§1.2].
Three dispatcher backends are exposed — AllGather (simple, memory-heavy, small EP), all-to-all (NCCL, scales well), and Flex (DeepEP / HybridEP, NVL72-aware bandwidth-optimized kernels with overlap) [§2.1.3].
The ChainedOptimizer separates dense and expert parameter optimization with three design points: allreduce=False marking for expert params, separate reduction groups (dp_cp vs expt_dp), and gradient scaling by edp_size / dp_size to compensate for the routing-dependent batch experts see [§2.2.2].
ZeRO-style optimizer-state sharding composes with MoE: expert optimizer states shard across the EP group while dense states follow standard DP sharding, all under a single ChainedOptimizer abstraction [§2.2.2].
Reduced-precision training (FP8 and NVFP4) is treated as a cross-cutting optimization because it simultaneously reduces activation memory, halves communication volume, and accelerates Tensor Core GEMMs — but requires selective precision strategies for convergence [§5].
The full stack reports 1233 / 1048 TFLOPS/GPU for DeepSeek-V3-685B on GB300 / GB200 and 974 / 919 TFLOPS/GPU for Qwen3-235B, scaling to thousands of GPUs [Abstract, §8].
The systems primitives also support RL post-training workflows including packed sequences, dynamic context parallelism, online weight export to inference engines (Megatron-Bridge), and “router replay” for RL [§10].

Method

Megatron-Core MoE rebuilds the MoE layer as three swappable modules (router / dispatcher / experts) connected by a four-stage forward pass (Route → Dispatch → Compute → Combine). The router is a learned linear projection + softmax or sigmoid + top-k selection, optionally in FP32. The dispatcher permutes tokens so same-expert tokens are contiguous (required for dense Grouped GEMM), then runs one of three backends: AllGather (filter-on-receive), all-to-all (point-to-point NCCL), or Flex (DeepEP for NVLink overlap, HybridEP for NVL72 topologies). Experts use TEGroupedMLP (Transformer Engine grouped FP8/FP4 GEMM) so all local experts run in a single GEMM call, with optional shared-expert compute overlapped with the dispatch-combine pipeline.

The systems contribution is multi-dimensional parallelism with MoE Parallel Folding: rather than forcing attention and MoE to share TP/DP/PP groups, the runtime exposes separate tp, cp, ep, expt_tp, expt_dp, and tp_ep process groups, with the router using tp/cp/tp_cp (weights duplicated across EP), the dispatcher using ep/tp_ep, the experts using ep/expt_tp/expt_dp, and shared experts using plain tp. Memory work-around recipes (fine-grained activation recomputation, memory-efficient permutation, activation offloading, FSDP, distributed optimizer) and compute work-around recipes (Grouped GEMM, kernel fusion, CUDA Graphs, sync-free execution for dropless MoE) round out the design. Selective FP8/FP4 lives across expert GEMMs, activations, and communication, with a stability strategy keeping the router in higher precision.

Results

Throughput (Abstract / §8): 1233 TFLOPS/GPU on GB300 and 1048 on GB200 for DeepSeek-V3-685B; 974 / 919 TFLOPS/GPU on GB300/GB200 for Qwen3-235B.
Sparsity asymmetry quantified: DeepSeek-V3 has 685B total / 37B active = 18:1 ratio vs. dense Llama-70B’s 1:1, directly explaining the memory-vs-compute imbalance [§3.2.1, Table].
GEMM share drops: 70% of execution time → <50% in MoE; routing/permutation adds 9% even after optimization [§1.2].
Memory accounting for context: Llama-405B BF16 + Adam needs ~12.9 TB across params (810 GB) + optimizer states (4860 GB) + gradients (1620 GB) + activations 8K seq (5575 GB) — far beyond any single GPU [§3.1.1].
Communication-wall framing: dispatch volume ≈ L · k · d per GPU with full cycle 2×; saturates with EP growth as it spills off NVLink onto narrower inter-node links [§1.2].
The report does not present headline FID/perplexity numbers — it’s a systems paper; quality is implicitly inherited from DeepSeek-V3 and Qwen3-235B as configurations.
Production features include distributed checkpointing with flexible resharding, upcycling from dense checkpoints, MTP integration, capacity-controlled token dropping, and integration with the open-source volcengine/Megatron-LM lineage.

Why it’s interesting

This is the open NVIDIA write-up of the engineering substrate behind the recent wave of large open MoE models — DeepSeek-V3, Qwen3-235B, Nemotron-3 — and pairs naturally with the algorithmic side of the wiki’s MoE coverage. The “Three Walls” framing is genuinely useful: it crisply names the asymmetry that makes MoE-at-scale qualitatively different from dense, and Parallel Folding is the cleanest argument on file for why attention and MoE layers should be allowed different TP/DP configurations rather than the usual shared-grid orthodoxy.

Direct counterpart to veScale-FSDP: Flexible and High-Performance FSDP at Scale — ByteDance’s veScale-FSDP solves the same MoE-padding-inflation pain point (the report quantifies that FSDP2-style per-parameter even-sharding doubles the AllGather buffer for 128-expert models on 256 devices) but at the sharding-format layer (RaggedShard) rather than at the process-group layer (Parallel Folding). The two papers are largely complementary: veScale-FSDP fixes how parameters are laid out on devices; Megatron-Core MoE fixes how communication is routed between them. Both effectively argue that stock open-source FSDP is not MoE-native.

Also the systems dual of Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method (Quantile Balancing — algorithmic fix for router load imbalance) and DynaMoE: Dynamic Token-Level Expert Activation with Layer-Wise Adaptive Capacity for Mixture-of-Experts Neural Networks (architectural fix via dynamic expert counts). All three are in the “make MoE less fragile” cluster but at different loci (systems vs. routing-time correction vs. layer-architecture sweep). Worth reading alongside Mixture of Experts (MoE), Visually Explained as the conceptual primer.