Back to News
InnovationAI Understanding briefing

NVIDIA reports 10x faster dropless MoE training in JAX

NVIDIA says its Transformer Engine and JAX optimizations increased DeepSeek-V3 671B training throughput by about 10x and achieved 97% efficiency across 1,024 GPUs. The optimized path is included in an NVIDIA NGC MaxText container dated September 9, 2026, or newer.

4 min readRead the primary source
Source-provided image accompanying NVIDIA reports 10x faster dropless MoE training in JAX
Primary-source documentSource recorded
Publisher
developer.nvidia.com
Source link
developer.nvidia.comhttps://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine/
Source type
Primary document — an official announcement, paper, filing, or first-party page we read directly.
ContextUnderstand this in 60 seconds

Start here

Key terms

Mixture of Experts (MoE)
An architecture with specialized subnetworks where only selected experts run per input.
Memory (Agent Memory)
Stored context an AI agent uses across steps or sessions to improve continuity.
Quantization
Converting model weights to lower precision formats such as 8-bit or 4-bit.
Test yourselfAI Models Explained Quiz

What happened

NVIDIA describes a set of JAX and Transformer Engine optimizations for dropless mixture-of-experts training, including grouped GEMM, fused expert-parallel dispatch and combine operations, NCCL EP, XLA multistream collectives, MXFP8 grouped quantization, and host activation offloading. The company says an unoptimized DeepSeek-V3 training baseline reached 103 TFLOPS per GPU, compared with 1,068 TFLOPS per GPU after optimization, and that end-to-end throughput improved by about 10x. NVIDIA also reports 97% scaling efficiency at 1,024 GPUs. The optimized components are included in the NVIDIA NGC MaxText container, with instructions for testing and reproducing the configuration.

NVIDIA’s post focuses on dropless mixture-of-experts training, in which every token is processed by its selected expert even when routing creates uneven expert loads. Unlike capacity-based approaches that trim overflow or pad experts to fixed sizes, dropless training requires kernels and communication paths to handle variable token counts. NVIDIA says these irregular shapes produce ragged tensors and can leave GPUs waiting on dispatch, all-to-all communication, and expert computation.

The reported changes operate at several layers. Transformer Engine’s grouped GEMM processes variable-length expert groups in one kernel call, while fused dispatch and combine operations and the NCCL EP backend address movement of tokens between GPUs. NVIDIA also cites XLA multistream collectives, host activation offloading, and MXFP8 grouped quantization as contributors to the result.

The company reports that its DeepSeek-V3 671B configuration increased a baseline from 103 to 1,068 TFLOPS per GPU and produced an approximately 10x end-to-end throughput gain. It also reports 97% efficiency at 1,024 GPUs. These figures are presented by NVIDIA as observations from its optimized stack, not as independently verified results.

The implementation is available to try through the NVIDIA NGC MaxText container with Transformer Engine built in. The post specifies ghcr.io/nvidia/jax:maxtext-2026-09-09 or newer and provides configuration guidance for DeepSeek-V3. It cautions that different models require different tuning. Pricing and licensing details are not provided.

Source details: developer.nvidia.com

Why it matters

Large mixture-of-experts models reduce computation by activating only selected expert networks, but their uneven token routing creates difficult communication and memory bottlenecks. If NVIDIA’s results hold beyond its reported configuration, the changes could make training very large MoE systems more efficient without dropping routed tokens or padding every expert to a fixed capacity. That matters primarily to organizations operating large NVIDIA GPU clusters, where communication overhead and wasted computation can materially affect training time and cost. The results are NVIDIA’s own claims; the source does not provide independent benchmarking or a comparison with every alternative implementation.

MoE training is increasingly important for large AI models because conditional computation can reduce the active computation per token while retaining many parameters. The engineering difficulty shifts toward routing, irregular matrix operations, memory use, and interconnect traffic. Improving those parts could reduce the time and hardware required to train models at production scale.

The practical audience is specialized: researchers and companies already using JAX, MaxText, NVIDIA GPUs, and multi-GPU or multi-node infrastructure. The source does not establish that the same gains apply to smaller systems, non-NVIDIA hardware, dense models, or workloads outside the reported DeepSeek-V3 configuration.

NVIDIA’s figures should be treated as vendor-reported performance claims. The source supplies no independent tests, detailed cost analysis, or full methodology sufficient to determine how much of the improvement comes from each optimization.

What to watch next

The main question is whether the reported gains reproduce across different MoE architectures, cluster layouts, model sizes, and GPU generations. NVIDIA plans to add NVFP4, quantization fused with GEMM, and additional all-to-all overlap. Users should also watch whether these features become standard in JAX and MaxText workflows, and whether the optimizations require substantial model-specific tuning. The source documents a container-based access path but does not state pricing, licensing terms, support commitments, or general availability beyond the cited NGC container.

Reproduction across other MoE models and hardware configurations will indicate whether the reported gains are broadly transferable or tightly coupled to DeepSeek-V3 and NVIDIA’s cluster setup.

NVIDIA says future work will add NVFP4, fused quantization with GEMM, and additional all-to-all overlap. Those additions could further affect the performance and memory tradeoffs of the stack.

The source recommends tracking step time, TFLOPS per GPU, model FLOP utilization, grouped GEMM latency, and dispatch/combine latency. Those measurements will help separate compute gains from communication improvements.

The post does not state the container’s price, licensing conditions, support level, or whether all cited components are equally mature for production use.

Related guides & quizzes

AI Models ExplainedTransformersAI TrainingFuture of AITest what you know — try a free AI quizLook up an AI term in our glossary
Found this useful?