PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zijie Yan, Hongxiao Bai, Xin Yao, Dennis Liu, Tong Liu, Hongbin Liu, Pingtian Li, Evan Wu, Shiqing Fan, Li Tao, Robin Zhang, Yuzhong Wang, Shifang Xu, Jack Chang, Xuwen Chen, Kunlun Li, Yan Bai, Gao Deng, Nan Zheng, Vijay Anand Korthikanti, Abhinav Khattar, Ethan He, Soham Govande, Sangkug Lym, Zhongbo Zhu, Qi Zhang, Haochen Yuan, Xiaowei Ren, Deyu Fu, Tailai Ma, Shunkang Zhang, Jiang Shao, Ray Wang, Vasudevan Rengasamy, Rachit Garg, Santosh Bhavani, Xipeng Li, Chandler Zhou, David Wu, Yingcan Wei, Ashwath Aithal, Michael Andersch, Mohammad Shoeybi, Jiajie Yao, June Yang

arXiv:2603.07685v2cs.DCcs.CLcs.LG

TL;DR

MoE sparsity creates coupled memory, communication, and computation constraints because total parameters grow faster than per-token computation. Megatron-Core addresses these constraints with integrated system optimizations and flexible parallelism, reporting high throughput on modern NVIDIA platforms. The report also identifies overheads and quantization-related risks that bound these techniques.

  • Problem

    MoE sparsity creates coupled memory, communication, and computation constraints because total parameters grow faster than per-token computation.

  • Method

    Megatron-Core combines memory, communication, computation, and parallelism techniques to scale MoE training across diverse model and hardware configurations.

  • Results

    1,233/1,048 TFLOPS per GPU are reported for DeepSeek-V3 on 256 GB300/GB200 GPUs, while Qwen3-235B reaches 974/919 TFLOPS on GB300/GB200.

  • Takeaways & Limitations

    The integrated optimizations reduce all-to-all’s contribution to training time under 10% and reduce memory from a blocking barrier to a manageable constraint.

  • Takeaways & Limitations

    The reported performance is a point-in-time snapshot based on Megatron-Core v0.16, and aggressive quantization can destabilize routing or cause expert collapse if sensitive components are not protected.

Abstract

from arXiv · show

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.

NVIDIA1

The passage directs readers to the Contributions and Acknowledgments section for the complete author list and identifies the corresponding authors.

  • The complete list of authors appears in the Contributions and Acknowledgments section.
  • Corresponding authors are identified by the email addresses zijiey, juney, and jiajiey at nvidia.com.
  • The passage provides author-contact information rather than technical content.

1. Introduction

MoE increases model capacity while activating only a subset of parameters per token, but this sparsity creates coupled memory, communication, computation, and parallelism challenges. Megatron-Core addresses these challenges with integrated parallelism, memory, communication, computation, and production optimizations for large-scale MoE training.

  • MoE Motivation: MoE routes each token to a small subset of expert networks, allowing capacity to grow independently of per-token computation.The architecture replaces dense FFNs with multiple experts and selects a subset using learned routing weights.
  • MoE Motivation: MoE sparsity creates a parameter-compute mismatch because total parameters scale with E while active computation scales with K, where K≪E.DeepSeek-V3 is presented as an example with substantially more total than active parameters.
  • Systems Challenges: Expert Parallelism preserves full-size expert computations but introduces all-to-all communication for routing tokens across GPUs.Naively sharding expert matrices can fragment already-small computations, motivating expert placement across devices.
  • Systems Challenges: The three coupled scaling barriers are memory pressure from all experts’ states, communication overhead from dispatch and collection, and fragmented computation.The memory and communication walls arise because all experts must be stored while only a subset activates, and tokens must move to assigned experts.
  • Megatron-Core Approach: Megatron-Core combines multidimensional parallelism, Parallel Folding, memory optimization, communication overlap, and compute kernels to address these barriers.The stack includes recomputation, offloading, reduced precision, optimized dispatchers, Grouped GEMM, fusion, and CUDA Graphs.
  • Megatron-Core Approach: Production features include load balancing, capacity-controlled token dropping, distributed optimization, FSDP, flexible checkpoint resharding, and dense-checkpoint upcycling.The modular stack is designed for experimentation and training from research-scale systems to trillion-parameter models.

2. Megatron-Core MoE Architecture

Megatron-Core implements MoE layers as modular router, dispatcher, and expert components connected by route, dispatch, compute, and combine stages. Separate process groups, dispatch backends, Grouped GEMM, and specialized optimizer handling support distributed execution.

  • MoE Layer Architecture: An MoE layer replaces a dense FFN with router, token-dispatcher, and expert modules connected through a four-stage forward pass.The stages are Route, Dispatch, Compute, and Combine.
  • Route: Routing maps token hidden states to expert logits, converts them to probabilities, and selects the highest-scoring top-k experts.The router outputs routing weights and a token-expert assignment mask.
  • Dispatch: Dispatch permutes tokens by destination expert and moves them across GPUs using AllGather, all-to-all, or Flex backends.AllGather is simple but memory-intensive, all-to-all scales through targeted sends, and Flex supports optimized backends such as DeepEP and HybridEP.
  • Expert Computation: Grouped GEMM executes local experts together, while SequentialMLP provides a slower one-expert-at-a-time implementation useful for debugging.TEGroupedMLP supports FP8 and FP4 quantization.
  • Combine: Combine returns processed tokens to their original GPUs, restores sequence order, and can add a shared expert’s output.Shared expert computation may run in parallel with routed expert processing to hide latency.
  • Distributed Integration: Distinct process groups and optimizer handling let attention and MoE layers use different parallel configurations and reduction groups.Parallel Folding can use configurations such as TP=4 for attention and ETP=1 with higher EP for MoE, while Chained-Optimizer separates dense and expert parameters.

3. Scaling MoE: Parallel Folding and Multi-Dimensional Parallelism

MoE sparsity creates a parameter-compute mismatch and conflicting parallelism needs between attention and MoE layers. Megatron-Core addresses this with Expert Parallelism and Parallel Folding, which decouples their mappings while integrating multiple parallelism dimensions.

  • 3.2. The Challenge of MoE Parallelism: MoE models require more GPUs for memory while per-token computation remains low, exposing communication overhead as expert count grows.Only K of E experts activate per token, so total parameters scale with E while computation scales with K.
  • 3.2.3. Expert Parallelism: The Fifth Dimension: Expert Parallelism distributes experts across GPUs, grouping tokens to improve GEMM efficiency while keeping all-to-all volume constant as expert count increases.The number of GPUs changes with expert count, but the all-to-all communication volume remains constant.
  • 3.2.4. The Challenges of Combining EP with Traditional Parallelism: Attention and MoE layers have conflicting optimal parallelism configurations, making a shared configuration a structural dense-sparse mismatch rather than a tuning problem.High TP benefits attention but fragments expert shards, high CP benefits long-context attention but not MoE, and high EP benefits MoE but not attention.
  • 3.3. Parallel Folding: Parallel Folding decouples attention and MoE parallelism mappings so each layer type can use its optimal topology.This is Megatron-Core’s response to the dense-sparse mismatch created by MoE sparsity.
  • 3.3.4. Summary: Megatron-Core combines Parallel Folding with Distributed Optimizer and FSDP support to coordinate multi-dimensional parallelism for large-scale MoE training.The framework is designed to decouple layer mappings while further reducing memory footprint.
  • 3.3.3. Benefits of Parallel Folding: Parallel Folding eliminates the EP ≤ DP constraint by folding EP across attention parallelism groups.With attention configured as TP=4, CP=2, DP=8, and PP=4, folding enables EP=64 instead of the traditional maximum EP=8 while preserving the attention configuration.

4. Scaling MoE: Breaking the Memory, Communication, and Compute Efficiency Walls

Megatron-Core addresses MoE’s memory wall through complementary optimizations that reduce activation, parameter, and optimizer-state footprints while preserving throughput. Fine-grained recomputation, offloading, precision-aware storage, FSDP, and communication improvements jointly make large-scale configurations feasible.

  • Memory Wall: Activations dominate large-scale MoE memory consumption, making activation optimization the highest priority for larger batches and flexible parallelism.Activations can exceed the combined memory of weights, gradients, and optimizer states.
  • Memory Wall: Memory-Efficient Permutation eliminates redundant intermediate tensors through an algebraic rearrangement without computational overhead.For DeepSeek-V3, it saves approximately 26.3 GB of activation memory per GPU.
  • Memory Wall: Fine-grained recomputation targets memory-intensive, computationally cheap operations, avoiding the higher overhead of full-layer recomputation for MoE layers.Full-layer recomputation can add approximately 33% computational overhead and retrigger expert-parallel all-to-all communication.
  • Memory Wall: Fine-grained offloading reduces memory by 10–18% with only 1.6–2% throughput overhead, while enabling a 15.0% throughput improvement for Qwen3-235B through reduced tensor parallelism.Asynchronous transfers overlap with computation to hide PCIe latency.
  • Memory Wall: Precision-aware optimization reduces optimizer-state memory by approximately 50%, while state offloading saves 15–20 GB of GPU memory with only 0.1–0.2 seconds of iteration overhead.These methods can be combined because lower-precision optimizer states reduce the amount of data that must be offloaded.
  • Memory Wall: FSDP with expert parallelism scales memory and collective volume with the expert-data-parallel group, while avoiding several pipeline-parallel configuration difficulties.This supports more experts or larger batches under the same hardware budget and simplifies parallelism configuration.

5. Reduced-Precision Training in FP8/FP4 for MoE

Megatron-Core treats reduced precision as a cross-cutting MoE optimization spanning memory, communication, and computation. Its selective-precision strategy protects numerically sensitive components while quantizing expert computation, with platform-specific FP8/FP4 recipes and alignment-aware kernels.

  • Selective precision: Router quantization can destabilize expert selection, so MoE training keeps routing and other numerically sensitive components in higher precision.The router remains in FP32, while embeddings, output layers, gradients, master weights, and optimizer states retain their original precision.
  • Cross-cutting benefits: FP8/FP4 training simultaneously improves memory, communication, and computation efficiency, making reduced precision a unifying optimization for MoE.The approach reduces activation storage and parameter AllGather traffic while accelerating expert GEMMs.
  • Memory and communication: Parameter AllGather communication is reduced by 50% when primary weights use FP8/FP4, and native low-precision paths bypass BF16 intermediates.The direct FP32-to-FP8/FP4 casting path reduces memory footprint and accelerates parameter AllGather.
  • Platform-specific recipes: Blockwise FP8 is recommended for Hopper, while MXFP8 is the default FP8 recipe for Blackwell; Megatron-Core also provides NVFP4 for Blackwell.Blockwise scaling uses 1×128 activation and gradient tiles and 128×128 weight blocks, whereas MXFP8 uses 1×32 granularity.
  • Kernel and shape optimization: Grouped quantization and fused padding align dynamic expert tensors while reducing CPU overhead, memory traffic, and preprocessing before grouped NVFP4 GEMMs.These kernels are CUDA-Graphable and produce aligned, zero-padded expert activations for downstream computation.

6. Long-Context MoE Training

Long-context training shifts the dominant computational concern from MoE layers to attention while intensifying activation-memory pressure. Megatron-Core combines CP, TP, offloading, selective recomputation, packed sequences, and Dynamic-CP to support long and variable-length workloads.

  • Long-context bottlenecks: At 64K tokens, SDPA consumes 69% of FLOPs, compared with 10–15% in short-sequence scenarios, shifting optimization emphasis toward memory and communication.SDPA scales as O(s^2), whereas MoE and other attention operations scale as O(s).
  • Memory management: Context Parallelism and Tensor Parallelism distribute activation memory across devices while keeping per-device sub-sequence lengths near 4096 or 8192 tokens.Scaling CP × TP with sequence length keeps per-device memory near baseline levels.
  • Memory management: Optimizer CPU offloading and selective recomputation provide additional memory headroom, but offloading trades memory savings for transfer and host-side optimizer overhead.For DeepSeek-V3 on 256 H100 GPUs at sequence lengths of at least 16K, the reported worst-case offloading overhead is about 2%.
  • Variable-length training: Packed sequences avoid padding waste for variable-length samples, reducing memory usage by 40–60% and improving throughput by 1.5–2× in reported RL training workloads.The benefits become increasingly important beyond 32K tokens, where padding waste can dominate memory use.
  • Variable-length training: Dynamic-CP jointly selects packing plans and per-microbatch CP sizes to reduce data-parallel synchronization stalls and unnecessary CP communication.It addresses both DP imbalance and CP inefficiency caused by static CP sizing for packed sequences.
  • Long-context results: At 256K tokens on 256 Hopper GPUs, DeepSeek-V3 reaches 88% of short-context MFU and Qwen3-235B-A22B reaches 129% using different TP/CP and memory-optimization combinations.The reported configurations combine selective recomputation with optimizer offloading for DeepSeek-V3 and with CP for Qwen3.

7. Production Features

Megatron-Core MoE adds production features for balancing workloads, supporting flexible expert architectures, reconfigurable parallelism, and heterogeneous pipeline stages. These mechanisms target operational robustness and efficient deployment across varied model structures.

  • Load Balancing: Load balancing and token dropping address routing-induced workload imbalance, memory bottlenecks, and degraded hardware utilization.Supported strategies include auxiliary loss, expert choice, auxiliary-loss-free biasing, dropless dispatch, and capacity-limited dropping.
  • Load Balancing: Pad-to-max converts variable expert token counts into static shapes, enabling CUDA Graphs that require fixed tensor dimensions.
  • Shared and Latent Experts: Shared experts process all tokens, and overlap runs shared-expert computation in parallel with dispatch and combine communication to hide latency.
  • Shared and Latent Experts: LatentMoE compresses routed expert inputs and outputs to dimension ℓ<d, reducing all-to-all volume and expert weight size by the compression ratio α=d/ℓ.The recommended ℓ-MoEacc scales expert count and top-k to restore inference cost while improving accuracy at iso-cost; it outperforms standard MoE in accuracy per FLOP and parameter at scales up to 95B.
  • Checkpointing: Distributed checkpointing stores sharded tensors independently and automatically reshares them for changed parallelism configurations without offline conversion.A checkpoint saved with TP=2, EP=4 can be loaded with TP=4, EP=8.
  • Flexible Asymmetric VPP: Flexible Asymmetric VPP assigns different layer types and counts to virtual stages, balancing computational costs and supporting arbitrary layer compositions.The DeepSeek-V3 example places dense, MoE, embedding, MTP, and loss layers according to their relative workloads.

8. Performance Evaluation

Megatron-Core MoE is evaluated on DeepSeek-V3-685B and Qwen3-235B across GB300, GB200, and H100 systems with the full optimization stack. The results show high per-GPU efficiency, thousand-GPU scalability, and strong long-context throughput.

  • Evaluation Setup: The evaluation measures per-GPU TFLOPS and tokens per second for fine-grained MoE workloads on GB300, GB200, and H100.Benchmarks use force-balanced routing and the fully enabled optimization stack.
  • Performance Results: 368 TFLOPS per GPU on H100 for DeepSeek-V3 demonstrates effective blockwise FP8 training for large-scale MoE models.The reported recipe maintains numerical stability while addressing dynamic routing and varying computation patterns.
  • Performance Results: GB200 achieves over 1,048 TFLOPS per GPU and GB300 reaches 1,233 TFLOPS per GPU for DeepSeek-V3 training.GB200 and GB300 deliver approximately 3× higher token throughput than H100 for both evaluated models at comparable or smaller GPU counts.
  • Scalability: Megatron-Core MoE scales to a 1,024-GPU DeepSeek-V3 run, validating parallel folding, optimized dispatchers, and kernel optimizations.
  • Long-Context Performance: A 131,072-token Qwen3-235B run on GB300 sustains 1,150 TFLOPS per GPU under the full optimization stack.This stress test targets memory- and communication-intensive long-context training.
  • Scope: The reported benchmark configurations are empirically tuned best-found settings rather than guaranteed global optima, and the results are a point-in-time snapshot of Megatron-Core v0.16.

9. Performance Best Practices

The performance methodology begins with memory feasibility, then selects communication-aware parallelism and profiles the dominant bottleneck. Hardware topology changes which optimization stack is most effective, so tuning proceeds iteratively.

  • Methodology: Megatron-Core uses an iterative workflow in which solving one bottleneck exposes the next, requiring profiling and refinement across memory, communication, and computation.
  • Memory Feasibility: Memory feasibility is the first constraint before throughput optimization, because parallelism and activation choices must fit GPU memory.For a 685B-parameter fine-grained MoE model, BF16 activations alone can exceed 130 GB per GPU.
  • Parallelism Selection: Keep parallelism degrees as small as possible while avoiding OOM, and use distributed optimizers to shard optimizer states across data-parallel ranks.
  • Communication-Aware Mapping: Keep EP×TP within the NVLink domain when possible, using pipeline parallelism beyond that domain and overlap when EP communication exceeds hidden bandwidth.
  • Expert Parallelism: EP is preferred over TP for expert layers because it preserves larger GEMMs, reduces communication, and simplifies communication-computation overlap.For Mixtral-8×7B, EP8×TP1 outperforms EP4×TP2.
  • Bottleneck Diagnosis: After establishing parallelism, profiling identifies the dominant wall and directs targeted optimizations for memory, communication, or CPU overhead.
  • Hardware-Dependent Optimization: On NVL8, all-to-all communication may consume 30–50% of step time, whereas on NVL72 the bottleneck can shift to CPU overhead after FP8 accelerates GPU computation.The same model therefore requires different optimization strategies on different hardware.

10. Megatron-Core MoE in Reinforcement Learning

RL post-training imposes MoE-specific requirements through variable-length sequences, interleaved inference and training, and routing differences. Megatron-Core addresses these demands with packing, dynamic parallelism, offloading, precision options, and router replay.

  • RL Requirements: RL frameworks often interleave training with inference and require rapid memory offloading, while routing can differ between inference and training engines.
  • RL Requirements: RL workloads produce highly variable sequences, with maxima reaching 128K or 1M tokens while mini-batch means are often one-half to one-quarter of the maximum.This distribution complicates balancing compute efficiency against peak memory consumption.
  • Variable-Length Training: Packed sequences and packing-aware dynamic batch sizes remove padding and keep batches near a similar number of effective tokens.An additional balancing strategy accounts for heterogeneous transformer-block costs.
  • Variable-Length Training: Attention-cost sorting orders micro-batches in a small-to-large-to-small serpentine pattern to reduce synchronization bubbles across parallelism dimensions.The metric is the mini-batch sum of squared sequence lengths.
  • Dynamic Context Parallelism: Dynamic context parallelism avoids assigning every sequence the fixed CP degree required by the longest sequence, reducing unnecessary bandwidth use for shorter sequences.
  • Memory Management: CPU optimizer offloading evicts optimizer states during forward and backward passes, freeing GPU memory for activations or longer sequences.The states return to the GPU only for parameter updates.
  • Precision: Megatron-Core provides FP16 training with loss scaling and mixed-precision optimizer kernels, reflecting cases where FP16 can be more stable than BF16 in RL training.
  • Router Replay: Router replay records inference-time expert assignments and enforces them during training, decoupling routing variability from weight updates for more consistent RL optimization.

11. Conclusion

Megatron-Core MoE integrates parallelism, memory, communication, computation, and production optimizations to address MoE’s coupled systems challenges and enable trillion-parameter-scale training. The open-source stack delivers high throughput across modern NVIDIA platforms and supports experimentation through production deployment.

  • Integrated Systems Design: Megatron-Core MoE addresses MoE’s parameter-compute and dense-sparse mismatches through integrated solutions across memory, communication, computation, and parallelism.Its design targets the coupled Memory, Communication, and Compute Efficiency walls while decoupling attention and MoE parallelism.
  • Parallelism: Parallel Folding decouples attention and MoE layer configurations, breaking the restrictive EP ≤DP constraint and enabling hardware-aware parallelism mappings.Expert Parallelism integrates with tensor, pipeline, context, and data parallelism.
  • Memory Optimization: Memory optimizations reduce DeepSeek’s per-GPU footprint from 199.5 GB to under 80 GB.The techniques include fine-grained activation recomputation, memory-efficient permutation, precision-aware optimizers, and CPU activation offloading.
  • Performance: 1,233/1,048 TFLOPS per GPU: DeepSeek-V3 on 256 GB300/GB200 GPUs; 974/919 TFLOPS: Qwen3-235B on GB300/GB200.On H100, the corresponding results are 368 TFLOPS per GPU for DeepSeek-V3 and 320 TFLOPS per GPU for Qwen3-235B.
  • Performance: Approximately 3× higher token throughput than H100 is delivered by GB300 and GB200 platforms.The comparison demonstrates the framework’s ability to exploit next-generation hardware.
  • Production Support: Open-sourcing the stack provides production-grade MoE tools spanning rapid prototyping through trillion-parameter production models.The modular architecture supports experimentation and deployment at scale.

A. Notation Reference

The notation reference identifies the report’s notation and abbreviations for consistent interpretation of its technical content.

  • Table 19 summarizes the notation used throughout the report.
  • The reference applies throughout the report rather than to a single benchmark configuration.
  • Table 19 is the report’s notation and abbreviation reference.

B. Detailed Benchmark Configurations

This appendix identifies the parallelism and training settings associated with the reported benchmark throughput results. Each configuration records the system, GPU count, precision, and summarized parallelism and batch layout.

  • The section lists parallelism configurations and detailed settings for reproducing the performance numbers in Table 11.
  • Each benchmark configuration is identified by its system, GPU count, and precision format.
  • The corresponding hyper-parameter string summarizes the parallelism layout and batch configuration.

B.1. Configuration Details

The benchmark configurations are empirical best-found settings documented to support interpretation and reproduction of the reported throughput results. They should not be treated as globally optimal configurations.

  • Table 20 lists benchmark configurations corresponding to the throughput results in Table 11.
  • The listed settings were best-found through empirical tuning at the time of writing.
  • The configurations may not be globally optimal.

B.2. Key Optimizations

The section identifies the optimization features with the greatest throughput impact in benchmark runs. These features vary across workloads and platforms.

  • Performance-critical optimization features have a first-order impact on throughput in the benchmark runs.
  • The full optimization space is larger than the subset summarized in this section.
  • The summarized features are configured differently across workloads and platforms.

DeepSeek-V3

The reported configurations combine dispatcher choice, recomputation, 1F1B overlap, and CUDA Graph settings differently across GB300, GB200, and H100. These choices are intended to balance memory, communication, and kernel efficiency for sustained throughput under hardware-specific constraints.

  • HybridEP is used on GB300/GB200, while DeepEP is used on H100.
  • Recomputation settings vary by platform, including none, mlp, up_proj, moe_act, and layernorm configurations.
  • 1F1B overlap is enabled or disabled differently across GB300, GB200, and H100 configurations.
  • CUDA Graphs target attention, the MoE router, and MoE preprocessing on the listed configurations, with one configuration marking them off on H100.
  • The configuration pattern reflects throughput-first tuning under model- and hardware-specific constraints.

B.3. Reproducibility

Table 11 benchmark numbers can be reproduced either through Megatron-Bridge or directly through Megatron-Core with Megatron-MoE-ModelZoo scripts. Both workflows start from Table 20 and require only a few cluster-specific runtime adjustments.

  • Megatron-Bridge 11 provides a higher-level interface for model, parallelism, and optimization configuration.
  • Megatron-Core can be launched directly with model-specific Megatron-MoE-ModelZoo 12 scripts.
  • Direct Megatron-Core scripts expose low-level launch flags for performance tuning.
  • Both workflows begin from Table 20 and adjust cluster-specific settings such as hostfiles, environment modules, and scheduler options.
Loading 2603.07685v2…