NVIDIA Accelerates MoE Training 10x with TE for JAX

NVIDIA significantly improved Dropless MoE (Mixture of Experts) training performance through Transformer Engine (TE) for the JAX framework. While existing libraries achieved only about 105 TFLOPs per GPU for DeepSeek-V3 model training due to Ragged Tensor processing from routing imbalance and inter-GPU communication bottlenecks, applying TE achieved 1,025 TFLOPs per GPU in a GB200 environment, marking an approximately 10x performance improvement.

Key Technical Improvements

• Grouped GEMM and MXFP8 Quantization: Processes expert-specific matrix multiplications in a single kernel call and supports MXFP8 block scaling on the Blackwell architecture, eliminating padding waste.
• EP (Expert Parallelism) Communication Optimization: Integrates dispatch and combine operations into fused kernels based on NCCL EP, eliminating CPU critical paths and saving bandwidth.
• Memory Offloading: Prevents OOM (Out of Memory) and achieves 57% faster speeds compared to recomputation by offloading activations using Grace Blackwell's NVLink-C2C bandwidth.

Deployment and Application Conditions

The optimization is available in the NVIDIA NGC MaxText container (images from 2026-09-09 onwards) and TE JAX v2.19 or later. To reproduce the DeepSeek-V3 671B model, settings such as EP=8 and FSDP=16 must be used along with enabling options like te_moe_block=true and te_gmm_quantization="te_mxfp8". Future plans include supporting NVFP4 and further kernel fusion based on CuTe DSL to improve DeepSeek-V3 pre-training speed by up to 8% and GPT-OSS by up to 93%.