NVIDIA has achieved a 10.4x improvement in the performance of Mixture of Experts (MoE) AI training, jumping from 103 TFLOPS/GPU to 1,068 TFLOPS/GPU on NVIDIA GB200 GPUs, the company says. The leap resulted from combining the JAX Python library with NVIDIA Transformer Engine, tackling a significant bottleneck in scaling this increasingly popular AI architecture. Researchers found that, in an unoptimized DeepSeek-V3 MoE training run, inter-GPU communication consumed 84% of accumulated kernel time. Solving the challenge required specialized kernels to handle ragged tensors created when a learned router distributes tokens unevenly across expert networks.
JAX and NVIDIA Transformer Engine Accelerate Dropless MoE Training
NVIDIA and its partners achieved a 10.4x improvement in TFLOPS/GPU. This substantial gain stems from targeted kernel optimizations within Transformer Engine, specifically designed to address bottlenecks unique to MoE model operations. Production-scale MoE training presents challenges absent in dense models, including efficient token routing, expert dispatch and gathering, and all-to-all communication between GPUs. This communication bottleneck hindered GPU utilization, forcing processors to idle while awaiting data transfers.
The team addressed this issue by developing specialized kernels capable of handling ragged layouts, a core difficulty in MoE training where the learned router creates an uneven distribution of tokens to experts, resulting in “no clean rectangular GEMM to batch and dispatch.” Transformer Engine’s MoE optimizations were specifically designed to resolve this problem, enabling more efficient data flow and maximizing GPU processing time. Dropless MoE differs from capacity-based MoE in its approach to token routing, though the specific distinctions are not detailed in the published findings.
Additional optimizations further contributed to the performance increase, including JAX host offloading and XLA multistreaming collectives. JAX host offloading allows intermediate activations to be temporarily stored in host memory, reducing the demand on high-bandwidth memory during the forward pass.
Specifically, the query and value projection results were offloaded to host memory during DeepSeek-V3 training to conserve resources, as described in a related publication on reducing high-bandwidth memory bottlenecks. XLA multistreaming collectives also played a role, optimizing the execution of collective operations like all-reduce, and reducing the percentage of exposed collectives in DSv3 training. Demonstrating the scalability of the optimized system, the team achieved 97% efficiency up to 1,024 GPUs, indicating that the communication optimizations effectively preserved throughput as the cluster size increased.
DeepSeek-V3 Achieves 10.4x Speedup on NVIDIA GB200 with Optimized Kernels
This communication overhead significantly limited performance, as GPUs spent the majority of their time waiting for data rather than performing computations. Addressing this issue required a fundamental shift in how MoE models handle data distribution and processing across multiple GPUs. The team focused on specialized kernels capable of natively managing the irregular data layouts inherent in MoE training. Previous attempts to optimize these irregular layouts relied on looping through GEMM kernels or employing batched GEMMs, but these methods introduced device-to-host copies of token counts, adding latency and hindering scalability.
The new approach utilizes cuBLAS GroupedGEMM, XLA multistreaming collectives, and MXFP8 GroupQuant to minimize data movement and maximize GPU utilization. These optimizations collectively reduce the percentage of exposed collectives during DeepSeek-V3 training, allowing for more efficient computation.
Dropless MoE and capacity-based MoE represent distinct strategies for managing this token routing. The team’s implementation of dropless MoE, combined with the optimized kernels, enables efficient handling of these ragged tensors, according to NVIDIA. This approach drives a 10.4x improvement in TFLOPS/GPU and 97% scaling efficiency up to 1,024 GPUs on DeepSeek-V3 671B.
Ragged Tensors and Communication Bottlenecks in MoE Training
Mixture of Experts models, despite their efficiency gains through conditional computation, presented a significant communication bottleneck in early NVIDIA GB200 training runs; an unoptimized baseline consumed 84% of accumulated kernel time solely on inter-GPU communication. Because each expert receives a varying number of tokens, standard libraries struggle to efficiently process the resulting non-uniform data. Transformer Engine addresses this directly by providing specialized kernels designed to natively handle these ragged layouts, an important step in unlocking MoE’s potential.
This approach, combined with the JAX Python library, propelled performance from 103 TFLOPS/GPU to 1,068 TFLOPS/GPU, a 10.4x improvement, demonstrating the effectiveness of targeted kernel optimizations. Beyond simply processing the irregular data, efficient dispatch and combination of tokens across experts is paramount. If this “dispatch and combine path is not optimized, communication dominates and GPUs are underutilized,” leading to stalled processing as GPUs wait for data.
The team’s solution involved optimizing the all-to-all communication patterns inherent in expert parallelism (EP), ensuring that data transfer doesn’t become the limiting factor. This optimization isn’t merely about speed; it’s about maximizing GPU utilization, allowing them to perform more useful work. This high level of sustained performance indicates that the underlying architecture effectively scales, preserving throughput even as the cluster size increases.
The configuration also highlights the importance of tuning parameters specific to MoE routing, such as the number of routing groups and the top-k routing strategy . These parameters influence how tokens are distributed to experts, impacting both model quality and training efficiency.
The team also emphasizes the use of specific XLA flags to further optimize performance, including settings for all-reduce, all-gather, and reduce-scatter operations, as well as flags to control command buffer behavior and memory limits, the company says. Environment variables, with settings for XLA’s Python client memory fraction and CUDA device connections , influence resource allocation and performance. The team notes that dropless MoE training, combined with Transformer Engine’s grouped GEMM and EP kernels, not only preserves model quality but also delivers significant efficiency gains at scale.
Dropless MoE vs Capacity-Based MoE: Handling Dynamic Token Routing
Committing to dropless Mixture of Experts (MoE) training necessitates a fundamental shift in how kernels handle expert computation, demanding efficient processing of variable token counts within each expert. Unlike traditional approaches, dropless MoE processes every token through its selected expert, regardless of load imbalances, a strategy favored for maintaining model quality but challenging for system performance. This approach directly contrasts with capacity-based MoE, where fixed token budgets are assigned to each expert, requiring trimming or padding to maintain regularity and hardware compatibility, potentially sacrificing model accuracy for efficiency.
MegaBlocks, an efficient sparse training technique utilizing MoE, reformulated expert computation as block-sparse matrix multiplication, enabling each expert to operate on a varying number of tokens without resorting to artificial padding or token dropping. Achieving this required the development of new block-sparse GPU kernels, optimized grouped GEMM operations, and specialized dispatch and combine primitives designed to accommodate these variable token counts.
A key challenge lies in ensuring these kernels can handle dynamic shapes not only efficiently but also while remaining accessible on the CPU to facilitate CUDA graphs and prevent recompilation delays, a critical factor for sustained performance. The team discovered that a batched GEMM, while seemingly straightforward, computed based on the worst-case token capacity, even when fewer tokens were actually utilized, resulting in unnecessary computation, NVIDIA reports.
This inefficiency is resolved by employing a grouped GEMM, which processes all expert matrix multiplications within a single kernel call, utilizing the actual token count for each expert. “A grouped GEMM solves this by handling all expert matmuls in a single kernel call, each with its actual token count,” the team reports, highlighting the performance gains achieved through this optimization.
This approach avoids wasted compute cycles and contributes to the overall efficiency of the dropless MoE training process. The MaxText configuration, detailed in the accompanying guide, provides specific parameters for implementing these optimizations with the DeepSeek-V3 model.
See today’s quantum computing news on Quantum Zeitgeist for the latest breakthroughs in qubits, hardware, algorithms, and industry deals.
