NVIDIA Transformer Engine Boosts Dropless MoE Training in JAX by 10.4x

NVIDIA reports a 10.4x throughput improvement for Mixture of Experts (MoE) training in JAX using Transformer Engine, raising DeepSeek-V3 performance from 103 to 1,068 TFLOPS/GPU on GB200 hardware. The unoptimized baseline spent 84% of accumulated kernel time on inter-GPU communication. The optimized stack sustains 97% scaling efficiency at 1,024 GPUs on GB300 NVL72 hardware when training DeepSeek-V3 671B.

Dropless MoE and the ragged tensor problem

MoE models route each token to a learned set of experts, but the router distribution shifts during training, so expert loads vary per batch. Dropless MoE processes every token without dropping or padding, which preserves model quality but creates ragged tensors where each expert receives a different token count. Capacity-based MoE avoids this by assigning fixed token budgets, but that forces a tradeoff between quality and efficiency: overflow tokens are either dropped or padded, wasting compute and memory.

Handling ragged tensors requires kernels that accept dynamic shapes. Previous approaches used a loop of GEMM kernels, which required Device-to-Host copies of token counts on the critical path and broke CUDA graphs, or batched GEMMs that padded to worst-case capacity and added extra compute. Transformer Engine's grouped GEMM handles all expert matmuls in a single kernel call with actual token counts, computing only regions with valid tokens. It maps onto cuBLAS and cuBLASLt and, on Blackwell GPUs, enables MXFP8 block scaling for expert matmuls.

Fused dispatch and combine with NCCL EP

Expert parallelism requires two stages: dispatch moves tokens to their assigned GPUs, and combine routes processed tokens back and accumulates results. A naive implementation runs these as a serial chain, stalling the GPU and leaving communication idle while compute runs. Transformer Engine integrates both stages into a fused kernel path powered by NCCL EP, a communication backend tuned for imbalanced expert-parallel traffic. NCCL EP deduplicates tokens: when a token goes to multiple experts on the same rank or to multiple ranks on a remote InfiniBand node, it traverses the network once and is replicated on the receiving node, conserving bandwidth.

Additional optimizations include JAX host offloading, which moves intermediate activations like query and value projection results to host memory to reduce high-bandwidth memory pressure, and XLA multistreaming collectives. By default, XLA runs communication on a single stream, serializing collectives that could run in parallel. Multistreaming lets the compiler schedule independent collectives across separate CUDA streams, overlapping cross-node InfiniBand transfers with intra-node NVLink communication. The Latency Hiding Scheduler analyzes replica groups and checks for deadlock risk to decide which collectives are safe to overlap, requiring no manual annotation.

Performance and next steps

The end-to-end gain on DeepSeek-V3 671B came from stacking cuBLAS GroupedGEMM, XLA multistream collectives, MXFP8 GroupQuant, host activation offloading, and the optimized EP implementation. NVIDIA plans to add NVFP4, quantization fused with GEMM, and A2A overlap. The company points to the NGC MaxText container with Transformer Engine enabled to reproduce the optimized JAX MoE path, and to the MaxText MoE Configuration guide for setup details.