NVIDIA’s Transformer Engine Boosts MoE Training in JAX by 10x

NVIDIA has unveiled optimizations to its Transformer Engine that dramatically accelerate the training of Mixture-of-Experts (MoE) models using JAX. By addressing bottlenecks like token routing and inter-GPU communication, the company claims a 10x speed-up in end-to-end training throughput for DeepSeek-V3, a 671-billion-parameter MoE model. These advancements also enable 97% scaling efficiency for clusters with up to 1,024 GPUs.

MoE architectures have gained traction in recent years as a way to scale AI models efficiently. Unlike dense models, which activate all parameters for every token, MoE selectively routes each token to a subset of specialized subnetworks, or “experts.” This approach allows models to reach trillions of parameters without linearly increasing computational costs, making them particularly attractive for large-scale language and multimodal models. NVIDIA’s latest work integrates these principles with its high-performance hardware and software stack.

Breaking Through MoE Bottlenecks

Training MoE models comes with unique challenges, particularly in token routing and handling “ragged tensors,” where data distributions are uneven across experts. NVIDIA’s Transformer Engine introduces specialized kernels, such as grouped GEMM (general matrix multiplication) and expert-parallel operations, which address these inefficiencies. For example, the grouped GEMM kernel optimizes irregular tensor shapes by processing variable token counts in a single operation, avoiding the overhead of padding or splitting data.

The engine also leverages NCCL EP, a communication protocol tailored to MoE’s complex traffic patterns. By fusing token dispatch and combine operations into a single kernel, GPUs maintain high utilization, minimizing idle time caused by waiting for data transfers. These optimizations, combined with techniques like MXFP8 quantization and JAX host offloading, push the hardware performance envelope.

Dropless MoE: A Quality-First Approach

NVIDIA’s framework focuses on “dropless” MoE, where every token is processed regardless of load imbalances across experts. While computationally demanding, this method improves model quality by avoiding data loss. To support this, the Transformer Engine relies on block-sparse matrix operations and dynamically shaped kernels, ensuring that even uneven token distributions are efficiently handled.

This contrasts with capacity-based MoE approaches, which trim or pad tokens to fit fixed budgets, sacrificing model fidelity for hardware simplicity. NVIDIA’s advancements allow developers to prioritize performance without making such compromises.

Massive Scaling with JAX

NVIDIA’s efforts extend beyond single-GPU performance. The company demonstrated that its MoE training stack sustains 97% efficiency even on clusters with 1,024 GPUs. Achieving this level of scalability is critical for training trillion-token datasets, where inter-node communication overhead can quickly become a bottleneck. Innovations like XLA multistreaming collectives and Latency Hiding Scheduler (LHS) further minimize inefficiencies in large-scale deployments.

Why It Matters

MoE architectures are a core focus for the AI industry as models grow in size and complexity. Recent research by Microsoft and the Allen Institute for AI has highlighted MoE’s potential to balance performance and efficiency, with applications ranging from language models to multimodal systems. NVIDIA’s Transformer Engine enhancements position it as a key enabler of these advances, reducing the cost and time required to train next-generation AI models.

For developers, the optimizations are available in the NVIDIA NGC MaxText container, which includes pre-configured tools for reproducing NVIDIA’s results with DeepSeek-V3. As MoE adoption grows, these tools could become essential for companies looking to scale their AI infrastructure without breaking the bank.

Looking ahead, NVIDIA plans to integrate additional features like NVFP4 quantization and advanced kernel fusions, further improving performance. For now, its Transformer Engine represents a significant step forward in making MoE training practical at scale.

Image source: Shutterstock