NVIDIA reports optimized JAX stack for dropless MoE training with Transformer Engine

The documented access path is through MaxText plus NVIDIA Transformer Engine, not a new foundation model release. The relevant user-facing switch is te_moe_block: true, which MaxText documentation describes as using Transformer Engine’s fused expert-parallel MoEBlock for routing, dispatch, grouped GEMMs, and combining expert outputs. The MaxText documentation says this requires sparse_matmul=True, prefuse_moe_weights=True, and Transformer Engine JAX with expert-parallel MoE support, listed as TE 2.19+. maxtext/docs/reference/core_concepts/moe_configuration.md at main · AI-Hypercomputer/maxtext · GitHub

The NVIDIA blog’s minimal MaxText flags are:

te_moe_block: true
te_gmm_quantization: "te_mxfp8"
ragged_buffer_factor: 2.0
te_ep_overflow_check_every_n_steps: 20
sparse_matmul: true
prefuse_moe_weights: true

NVIDIA says the exact DeepSeek-V3 reproduction path additionally uses model_name: "deepseek3-671b", max_target_length: 4096, per_device_batch_size: 6, steps: 15, attention: "cudnn_flash_te", quantization: "te_fp8_currentscaling", te_gmm_quantization: "te_mxfp8", a custom rematerialization/offload policy, and a 128-GPU parallelism layout with FSDP and expert parallelism. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

A notable documentation wrinkle: the Transformer Engine 2.19 API docs expose JAX mesh concepts including an ep_resource axis for expert parallelism, but the specific “JAX: Expert Parallelism with TransformerEngine” tutorial page currently says “TODO — Coming soon.” Jax — Transformer Engine 2.19.0 JAX: Expert Parallelism with TransformerEngine — Transformer Engine 2.19.0 For developers, that means the capability appears to exist in the stack, but the public tutorial documentation is incomplete as of retrieval.

Architecture: why dropless MoE is hard

In a standard dense Transformer feed-forward block, every token uses the same MLP weights. In MoE, a router sends each token to one or more expert MLPs. That gives conditional computation: more total parameters can be available without activating all parameters for every token. Switch Transformer described MoE as selecting different parameters per input, yielding sparse activation but also noted that MoE adoption is hindered by complexity, communication cost, and training instability. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

The hard systems problem is that expert assignment is data-dependent. Each expert may receive a different number of tokens per step, creating ragged expert batches rather than uniform rectangular matrices. NVIDIA frames the critical bottlenecks as token routing, expert dispatch/gather, all-to-all communication, and ragged expert GEMMs. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

Capacity-based MoE avoids some of this irregularity by imposing a fixed per-expert capacity. Overflow tokens may be dropped or padding may be used to preserve regular shapes. MegaBlocks identified the same core tradeoff: existing frameworks restricted dynamic routing, forcing users to choose between dropping tokens or wasting compute/memory on padding. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

Optimization 1: grouped GEMM / ragged dot

NVIDIA’s first core optimization is grouped GEMM: rather than running one GEMM per expert or padding every expert to a worst-case capacity, grouped GEMM handles all expert matmuls in one call using each expert’s actual token count. NVIDIA says Transformer Engine’s grouped_gemm / ragged_dot path uses cuBLAS and cuBLASLt and enables MXFP8 block scaling on Blackwell GPUs. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

This is conceptually aligned with prior dropless MoE work, but the implementation path differs. MegaBlocks reformulated MoE computation as block-sparse matrix multiplication and built block-sparse GPU kernels; the NVIDIA JAX path emphasizes grouped GEMM, ragged buffers, MXFP8 quantization, and Tensor Core utilization through NVIDIA libraries. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

Optimization 2: expert-parallel dispatch and combine

MoE expert parallelism requires moving tokens to the devices hosting their selected experts, running the expert MLPs, and combining results back into the original token order. NVIDIA says Transformer Engine integrates dispatch and combine into a fused kernel path powered by NCCL EP, including token deduplication so a token routed to multiple experts on the same rank or remote InfiniBand node traverses the network once and is replicated on receipt. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

This targets the communication-heavy part of MoE. Prior work such as Tutel also focused on MoE all-to-all and dynamic workload handling, including adaptive pipelining and a two-dimensional hierarchical all-to-all algorithm. [](https://proceedings.mlsys.org/paper_files/paper/2023/file/5616d34cf8ff73942cfd5aa922842556-Paper-mlsys2023.pdf) The repeated pattern across systems papers is that MoE performance is rarely about matmul alone; routing, packing, communication, and load imbalance are first-order constraints.

Optimization 3: host offloading and XLA collectives

NVIDIA says its DeepSeek-V3 configuration uses host offloading for selected activations and XLA multistream collectives. In the blog, NVIDIA specifically says intermediate activations need not remain on device for the entire forward pass, and mentions offloading query and value projection results for DSv3 training. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog NVIDIA’s separate host-offloading blog reports that QKV activation offloading plus the Latency Hiding Scheduler improved throughput for Llama 3.1 405B, but that is a different model and should be treated as supporting context, not proof of the DeepSeek-V3 MoE result. Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading | NVIDIA Technical Blog

OpenXLA’s own flags guidance says xla_gpu_enable_latency_hiding_scheduler enables schedulers to overlap asynchronous communication with computation, and that optimization level -O1 enables advanced GPU passes including collective pipelining and latency hiding scheduling. XLA Flags Guidance | OpenXLA Project NVIDIA/JAX Toolbox documentation similarly says LHS is now enabled by O1 and documents XLA flags for asynchronous collectives. JAX-Toolbox/docs/GPU_performance.md at main · NVIDIA/JAX-Toolbox · GitHub

Evaluation methodology and reproducibility

The public method is configuration-centric, not a full benchmark artifact. NVIDIA provides a DeepSeek-V3-specific MaxText configuration, XLA flags, and environment variables, and says different models require different tuning. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog The disclosed run uses steps: 15, which is enough to reproduce a short performance experiment but not, by itself, evidence of long-horizon training stability, convergence, or final model quality. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

There is also a hardware-label ambiguity in the public text: the article body reports the 103-to-1,068 TFLOPS/GPU result for NVIDIA GB200, while later figure captions describe end-to-end DeepSeek-V3 671B performance and scaling on GB300 NVL72 hardware. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog A careful reproduction report should specify exact GPU type, rack topology, CUDA/cuDNN/NCCL versions, Transformer Engine commit/version, MaxText commit, XLA build, fabric configuration, and measurement window.

MoE is an old idea, but its modern LLM role is to increase parameter capacity while activating only a subset of experts per token. Switch Transformer simplified routing and showed sparse models could train faster than dense baselines under matched resources, while also highlighting complexity, communication costs, and instability. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity GLaM and later MoE models further established the scaling appeal of sparsely activated experts. GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

DeepSeek-V3 is a relevant stress test because it is itself a large MoE model. Its technical report describes 671B total parameters with 37B activated per token, using MLA and DeepSeekMoE architectures, an auxiliary-loss-free load-balancing strategy, and a multi-token prediction objective. DeepSeek-V3 Technical Report

MegaBlocks is the closest prior dropless MoE systems reference surfaced in search. It directly addressed the problem of dynamic expert token counts without token dropping, using block-sparse operations and new GPU kernels. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts Tutel is another important systems predecessor, focusing on adaptive MoE parallelism and optimized all-to-all communication. [](https://proceedings.mlsys.org/paper_files/paper/2023/file/5616d34cf8ff73942cfd5aa922842556-Paper-mlsys2023.pdf)

The main limitation is evidentiary: the performance numbers are NVIDIA’s own measurements. No independent reproduction, profiler traces, raw logs, or third-party benchmark report were found. The configuration is useful, but a 15-step reproduction setting is not equivalent to full pretraining evidence. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

The second limitation is portability. The stack is tightly coupled to NVIDIA GPUs, Transformer Engine, cuBLAS/cuBLASLt, NCCL EP, MaxText, and XLA GPU behavior. Teams on AMD, TPU, Trainium, or generic PyTorch stacks should not assume comparable gains.

The third limitation is model-quality evidence. Dropless routing avoids dropping routed tokens, but the NVIDIA post does not present downstream quality or convergence comparisons for the optimized path. Prior work motivates avoiding token dropping, but it does not independently validate NVIDIA’s DeepSeek-V3 JAX result.

Safety implications are indirect. This is training infrastructure, not a model capability release. The main safety-relevant effect is economic: faster large-MoE training can lower the cost of frontier-scale or near-frontier-scale model development, increasing the number of actors able to train large sparse models if they also have hardware access.

For AI infrastructure leaders, this is a reminder that MoE economics depend heavily on the software stack. Buying more accelerators is insufficient if routing, all-to-all, grouped matmuls, memory pressure, and compiler scheduling leave GPUs idle.

For JAX users, the announcement is strategically important: it suggests NVIDIA is investing in JAX as a first-class large-scale training stack alongside PyTorch. Transformer Engine officially supports both PyTorch and JAX. Getting Started — Transformer Engine 2.19.0

For developers, the practical adoption path is to start with the NVIDIA MaxText container, verify correctness on a smaller MoE, instrument step time, grouped GEMM latency, dispatch/combine latency, MFU, and exposed collective time, then scale. NVIDIA itself recommends starting small before scaling and tracking these metrics. Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine | NVIDIA Technical Blog

For procurement and strategy teams, the main takeaway is that Blackwell-era MoE training performance may increasingly depend on vendor-integrated software features such as MXFP8/NVFP4, NCCL EP, and compiler-level communication overlap. That can improve time-to-train, but it may also deepen platform lock-in.