Loading post…
NVIDIATanya LenzExplainer
Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine
Summary
Production-scale mixture-of-experts (MoE) training must handle uneven, data-dependent expert loads, creating ragged tensors and expensive dispatch, combine, and all-to-all communication. Dropless MoE processes every selected token, preserving model quality but requiring kernels that support variable token counts without padding or dropping. The post presents JAX and NVIDIA Transformer Engine optimizations centered on group-aware MXFP8 quantization, grouped GEMM, and fused expert-parallel (EP) dispatch and combine using NCCL EP. Grouped GEMM computes each expert’s actual workload in one kernel path, while EP reduces movement overhead and can deduplicate tokens crossing networks. For DeepSeek-V3 training on NVIDIA GB200, the reported throughput increased from 103 to 1,068 TFLOPS/GPU, a 10.4x improvement; the post also reports roughly 10x throughput improvement and 97% scaling efficiency at 1,024 GPUs.
Context
MoE routing produces different token counts for different experts, making expert computation ragged rather than rectangular. Expert parallelism adds dispatch and combine operations across GPUs, and poorly optimized communication can leave GPUs waiting for data. Dropless routing avoids discarding overflow tokens but increases the need for dynamic-shape kernels and efficient communication.
Approach / What changed
The implementation combines Transformer Engine’s group-aware MXFP8 quantization, grouped GEMM or ragged_dot for variable-length expert matmuls, and optimized expert-parallel dispatch and combine operations backed by NCCL EP. The source also identifies JAX host offloading and XLA multistreaming collectives as additional optimizations, and provides MaxText, XLA, and environment settings for reproducing DeepSeek-V3 results.
Takeaways
Dropless MoE sends every token to its selected expert, avoiding overflow drops and padding but requiring kernels that operate efficiently on variable token counts.
Grouped GEMM handles all expert matmuls in one kernel call using each expert’s actual token count, avoiding device-to-host token-count transfers and worst-case padded computation.
NCCL EP supports fused dispatch and combine paths and deduplicates tokens sent to multiple experts on the same rank or across a remote IB node, reducing network traffic.
FAQ
Why do dropless MoE workloads produce ragged tensors?
The learned router assigns different numbers of tokens to each expert, and the distribution changes between batches. Because expert inputs do not share a fixed shape, conventional rectangular GEMM paths may require padding or inefficient handling.
How does grouped GEMM avoid wasted MoE computation?
It processes all expert matmuls in a single kernel call while using each expert’s actual token count. This avoids computing a worst-case capacity padded for experts that received fewer tokens and avoids device-to-host transfers on the critical path.
What does expert parallelism optimize?
Expert parallelism manages token dispatch to the correct GPUs and combines results back in the original token order. Transformer Engine integrates these stages with NCCL EP, which can replicate deduplicated tokens at the receiving node instead of sending repeated network traffic.
Source and freshness
- Source
- developer.nvidia.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.