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

Abstract: Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models. Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation, creating coupled constraints across memory, communication, and computation. Optimizing one dimension often shifts pressure to another, demanding co-design across the full system stack. We address these challenges for MoE training through integrated optimizations spanning memory (fine-grained recomputation, offloading, etc.), communication (optimized dispatchers, overlapping, etc.), and computation (Grouped GEMM, fusions, CUDA Graphs, etc.). The framework also provides Parallel Folding for flexible multi-dimensional parallelism, low-precision training support for FP8 and NVFP4, and efficient long-context training. On NVIDIA GB300 and GB200, it achieves 1,233/1,048 TFLOPS/GPU for DeepSeek-V3-685B and 974/919 TFLOPS/GPU for Qwen3-235B. As a performant, scalable, and production-ready open-source solution, it has been used across academia and industry for training MoE models ranging from billions to trillions of parameters on clusters scaling up to thousands of GPUs. This report explains how these techniques work, their trade-offs, and their interactions at the systems level, providing practical guidance for scaling MoE models with Megatron Core.

• The paper introduces a production-ready MoE training stack that tackles memory, communication, and compute bottlenecks in distributed GPU systems.
• It employs Parallel Folding to decouple dense and MoE parallelism, enhancing throughput and resource utilization.
• The framework integrates reduced-precision training and dynamic routing to enable efficient scaling of trillion-parameter models.

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

Mixture-of-Experts (MoE) architectures facilitate efficient scaling of LLMs by conditionally activating only a small subset of sub-modules (experts) per input token. While this sparsity enables impressive reductions in per-token computation relative to parameter count, it introduces unique systems challenges that make efficient distributed training fundamentally distinct from dense transformer models. "Scalable Training of Mixture-of-Experts Models with Megatron Core" (2603.07685) presents a production-ready, open-source MoE training stack within the Megatron-Core framework that systematically addresses these challenges, enabling training at trillion-parameter scale with strong empirical throughput.

---

Mixture of Experts: Model and Systems Challenges

MoE models replace dense feed-forward submodules (FFN) within transformers with collections of independent experts, each a small MLP, and a dynamic router that assigns tokens to experts on a per-input basis. While only $K$ of $E$ experts are activated per token ($K \ll E$), all expert parameters reside in GPU memory, creating a severe parameter-compute mismatch: the number of total parameters far exceeds those involved in any specific computation. As a result:

• Memory capacity becomes a dominant constraint, with activation, gradient, and optimizer memory easily exceeding single-device capacity.
• Communication overhead is dominated by all-to-all collectives for expert parallelism (EP), often crossing node boundaries where bandwidth is limited.
• Compute efficiency is compromised by fragmented, small GEMMs and non-trivial CPU overhead from token permutation, routing, and dynamic kernel launches.

These tightly-coupled bottlenecks, termed the Three Walls—Memory, Communication, and Compute Efficiency—necessitate coordinated multi-dimensional optimization and parallelism strategies.

---

MoE System Architecture and Parallelism

The Megatron-Core MoE stack is architected for both modularity and systems co-design.

MoE Layer Design

Every MoE layer consists of modular components: (1) a router implementing top-$K$ expert assignment, (2) a dispatcher for cross-GPU communication of tokens, and (3) highly parallelized grouped expert computation using optimized GEMM libraries. The four-stage forward pass—route, dispatch, compute, combine—allows for flexible backends for dispatch (AllGather, All-to-All, or token-wise approaches like DeepEP/HybridEP) and efficient management of token-expert assignments.

Figure 1: Data flow through an MoE layer: Route, Dispatch, Compute, and Combine stages.

Figure 2: Router architecture: linear projection, score function, top-k selection, and load balancing.

Multi-Dimensional Parallelism: Parallel Folding

The report introduces Parallel Folding, which decouples parallelism strategies for attention and MoE blocks. Traditional approaches tie EP (expert parallelism) to DP (data parallelism) and enforce restrictive mappings (e.g., $\text{EP} \leq \text{DP}$), leading to suboptimal resource utilization and unnecessary scaling of GPU count. Parallel Folding eliminates this constraint by maintaining independent process groups for each type of layer, allowing:

• High TP (tensor parallelism) for dense attention layers (improving QKV GEMM utilization).
• High EP for MoE layers (maximizing per-expert GEMM size and communication efficiency).
• Arbitrary mappings that keep both EP and CP/TP communication within high-bandwidth NVLink domains, reducing cross-node traffic.

  Figure 3: Parallelism mappings: traditional constraints vs. MoE Parallel Folding decoupling.

  

  Figure 4: Parallel Folding: decoupled attention and MoE parallelism mappings.

---

Crossing the Memory, Communication, and Compute Efficiency Walls

Memory Optimization

The majority of MoE memory burden arises from activations, not parameters. The report details layered memory optimizations:

• Fine-grained recomputation avoids checkpointing bulky activations, selectively recomputing only those portions that are expensive to store but cheap to recalculate.
• Memory-efficient permutation eliminates unnecessary buffer allocation by algebraic rearrangement of expert output accumulation.
• Reduced-precision training (FP8/FP4) halves or quarters activation and parameter storage cost, with careful recipe selection for expert GEMMs to preserve convergence.
• Fine-grained offloading, using activation and optimizer state transfers to host memory, leverages CUDA streams for overlap, ensuring minimal training step overhead.

  Figure 5: Memory-Efficient Permutation.

  

  Figure 6: Selective Recomputation.

  

  Figure 7: Fine-grained activation offloading: stream overlap for forward and backward passes.

  

  Figure 8: Fine-grained offloading and recomputation: complementary memory optimization strategies.

  

Figure 9: FSDP2 per-parameter uniform sharding.

Figure 10: Persistent double-buffer design: two pre-allocated buffers are cycled across FSDP collectives, eliminating allocation overhead and enabling NCCL User Buffer Registration.

Communication Optimization

Optimized dispatchers (DeepEP, HybridEP) implement token-wise collective communication, reducing both latency and memory usage compared to classic all-to-all schemes. Combining:

• All-to-all overlap with computation: FWD-BWD merging (1F1B scheduling) and weight/data gradient splitting interleave communication with expert computation, hiding up to 93% of all-to-all latency under forward and backward passes.
• Efficient pipeline parallelism: Flexible VPP supports arbitrary per-stage layer placement, balancing workload even with highly heterogeneous expert and shared layers.

  Figure 11: The dispatch kernel design of HybridEP.

  

  Figure 12: The combine kernel design of HybridEP.

  

  Figure 13: Merged FWD-FWD Timeline with all-to-all Overlapping.

  

  Figure 14: Merged FWD-BWD Timeline with all-to-all Overlapping.

  

  Figure 15: EP all-to-all communication overlap strategies: baseline vs. 1F1B with W/D split.

  

  Figure 16: Interleaved PP Timeline with all-to-all Overlapping.

Compute Efficiency & Host Overhead

To address underutilization from small, fragmented expert GEMMs and excessive kernel launches:

• Grouped GEMM and kernel fusions batch multiple small GEMMs and route/permute kernels in one launch.
• Partial and Full CUDA Graph execution captures as much of the critical path as possible, drastically reducing host launch overhead, with paged stashing and device-initiated kernels ensuring shape dynamism does not block graphing.
• Sync-free dispatch and ECHO (Elastic Cloning for Hot Experts) balance expert loads and reduce straggler effects, with minimal extra synchronization or memory cost.

  Figure 17: The pipeline for permute fusion in the training process.

  

  Figure 18: The workflow of the router fusion.

  

  Figure 19: Traditional execution (top) versus CUDA Graph execution (bottom).

  

  Figure 20: Full versus layer-wise CUDA Graphs in one training iteration (three layers, two microbatches).

  

  Figure 21: Partial CUDA Graphs capture static components (attention, shared experts, router, preprocessing) while leaving dynamic expert computation outside the graph.

---

Reduced-Precision Training

The framework systematically exploits FP8/FP4 in expert and attention computation, parameter storage, and communications, but does so selectively to safeguard router, embedding, and optimizer semantics. The resulting full-stack compression achieves:

• $50$–$75\%$ reduction in activation memory per-GPU.
• Halved communication volume on FP8/FP4-capable devices.
• Sustained per-GPU throughput of over 1,200 TFLOPS (DeepSeek-V3-685B on GB300) and over 900 TFLOPS on Qwen3-235B (Table 1 in paper), with no regression in training convergence.

---

Long-Context and RL Training Features

For sequences of $16$–$128$k+ tokens, attention's quadratic complexity becomes dominant. The stack coordinates context- and tensor-parallelism, selective recomputation, and optimizer offloading to maintain high utilization at extreme sequence lengths. Packed-sequence and dynamic context-parallelism approaches reduce padding waste and improve hardware balance for RL and fine-tuning workloads, supporting efficient scaling to variable-length, multi-modal, and on-policy RL data.

---

Practical and Production Considerations

The stack incorporates comprehensive operational features:

• Load balancing and token dropping schemes for router stability.
• Distributed checkpointing with automatic parallelism reconfiguration.
• Upcycling procedures for dense-to-MoE conversion, enabling re-use of valuable dense model weights.
• Flexible modular design supporting diverse model architectures, including LatentMoE, shared experts, and multi-token prediction modules.
• Seamless integration with leading RL frameworks via Megatron-Bridge, supporting online weight export and router replay.

---

Empirical Results

Under the full integration of these optimizations, performance benchmarks demonstrate:

• DeepSeek-V3 (685B) achieves up to 1,233 TFLOPS/GPU on GB300, 1,048 TFLOPS/GPU on GB200, and $>350$ TFLOPS/GPU on H100 at 1024-GPU scale.
• Qwen3-235B achieves 974/919 TFLOPS/GPU on GB300/GB200, and 320 on H100.
• Long-context runs (Qwen3-235B, 128k tokens) sustain 1,150 TFLOPS/GPU.
• Scaling laws and best practices are codified in iterative workflows, balancing memory, communication, and compute bottlenecks with platform-specific optimization selection.

---

Implications and Prospects

This work provides the technical basis for efficient, production-scale training of trillion-parameter MoE models, validating strong empirical throughput on contemporary GPU clusters. The introduction of systematic parallelism decoupling (Parallel Folding), integrated memory/communication/compute optimizations, and flexible hardware-aware scaling will inform best practices for both industry deployments and research prototyping.

On the research frontier, the detailed treatment of reduced-precision MoE training opens the door for new work on further algorithmic sparsity, more aggressive quantization, and refined communication primitives. Looking ahead, seamless scaling to future topologies (multi-node NVLink, advanced NVSwitch), support for heterogeneous and hybrid expert architectures, and co-design with hardware/software advances (e.g., memory-centric accelerators, fine-grained data movement) will be integral to pushing MoE capacities even further.

---

Conclusion

"Scalable Training of Mixture-of-Experts Models with Megatron Core" delivers a highly optimized, modular MoE training ecosystem that bridges the gap between sparse, high-capacity model architectures and the realities of distributed GPU systems. By addressing the parameter-compute and dense-sparse mismatches head-on with Parallel Folding and deeply integrated memory, communication, and compute optimizations, it sets a new empirical and methodological standard for scalable MoE research and deployment (2603.07685).