NVIDIA Accelerates Dropless MoE Training in JAX Using Transformer Engine

NVIDIA Accelerates Dropless MoE Training in JAX Using Transformer Engine NVIDIA has introduced targeted optimizations using Transformer Engine and JAX to accelerate dropless Mixture of Experts (MoE) training, achieving a 10.4x throughput improvement on DeepSeek-V3 671B running on GB300 systems. Unoptimized DeepSeek-V3 training on GB300 achieved 103 TFLOPS/GPU with inter-GPU communication consuming 84% of kernel time, which rose to 1,068 TFLOPS/GPU after optimizations. Dropless MoE routes all tokens to designated experts without dropping or padding, relying on grouped GEMM and specialized kernels to manage resulting ragged tensors. Expert parallelism stages are integrated into a fused path using NCCL EP, incorporating token deduplication across network ranks. Training DeepSeek-V3 maintained 97% scaling efficiency at 1,024 GPUs on NVIDIA GB300 NVL72 hardware. Optimizations are available via the NVIDIA NGC MaxText container with Transformer Engine. Dropless MoE Bottlenecks and Ragged Tensors Unlike capacity-based MoE frameworks that constrain routing by enforcing fixed token budgets through truncation or padding, dropless MoE processes every token via its chosen expert to preserve model quality. Because routing is dynamic and learned, experts receive uneven and unpredictable token allocations per batch. This imbalance creates ragged tensors that break standard uniform GEMMs and causes unoptimized all-to-all communications to stall GPUs during dispatch and gather stages. Grouped GEMM and MXFP8 Block Scaling To bypass the overhead of sequential GEMM loops or padded batched GEMMs, Transformer Engine employs grouped GEMM backed by cuBLAS and cuBLASLt. This executes all expert matrix multiplications in a single kernel call while calculating only valid token regions. On NVIDIA Blackwell GPUs, this interface supports MXFP8 block scaling through grouped quantization kernels, ensuring Tensor Core utilization without requiring CPU-dependent dynamic shape transfers. Fused Expert Parallelism via NCCL EP Dispatch and combine stages—which permute and exchange tokens across GPUs and accumulate results—are fused into an integrated path powered by NCCL EP. The backend is designed for irregular traffic distributions and uses token deduplication, transferring tokens dispatched to multiple experts on the same remote rank across the network only once before replicating them locally. XLA Multistreaming and Memory Offloading Complementary optimizations include JAX host offloading of intermediate forward-pass activations, such as query and value projections, to conserve high-bandwidth memory. In addition, XLA multistreaming collectives employ a Latency Hiding Scheduler that analyzes replica groups to overlap independent cross-node InfiniBand transfers with intra-node NVLink communications across separate CUDA streams. Performance and Multirack Scaling Applying grouped GEMM, XLA multistream collectives, MXFP8 GroupQuant, host activation offloading, and optimized EP achieved a ~10x end-to-end throughput increase on DeepSeek-V3 671B. When scaling from 128 GPUs to 1,024 GPUs on NVIDIA GB300 NVL72 racks, the implementation sustained 97% scaling efficiency. Future additions are planned for NVFP4, GEMM-fused quantization, and all-to-all overlap.

NVIDIA 借助 Transformer Engine 在 JAX 中将无丢弃 MoE 训练吞吐量提升超 10 倍 在基于 NVIDIA GB300 NVL72 的 DeepSeek-V3 671B 模型训练中，NVIDIA 将 JAX 与 Transformer Engine 结合，针对无丢弃（dropless）MoE 的动态不规则张量与通信开销进行软硬件栈优化，使单 GPU 吞吐量从未优化的 103 TFLOPS 提升至 1,068 TFLOPS（提升 10.4 倍），并在 1,024 块 GPU 规模下保持 97% 的扩展效率。 在 NVIDIA GB300 运行 DeepSeek-V3 671B 训练时，未经优化的基线性能仅为 103 TFLOPS/GPU，GPU 间通信占累计内核时间的 84%；经系统级优化后提升至 1,068 TFLOPS/GPU。 基于容量（capacity-based）的 MoE 会对溢出 Token 进行裁剪或填充，导致质量损失或算力浪费；而无丢弃（dropless）MoE 会处理所有 Token，但带来了动态非规整张量（ragged tensors）的计算与调度挑战。 Transformer Engine 引入了基于 cuBLAS/cuBLASLt 的分组通用矩阵乘法（Grouped GEMM）及 MXFP8 分组量化，单次内核调用即可处理各专家动态长度的矩阵乘法，无需回传 CPU 即可支持 CUDA 图。 专家并行（EP）整合了分发（Dispatch）与聚合（Combine）阶段，结合 NCCL EP 后端实现通信融合，并通过 Token 去重机制减少跨节点网络带宽消耗。 搭配 JAX 主机激活值卸载（Host Offloading）与 XLA 多流通信重叠（LHS 调度器），系统在 1,024 块 GPU 上实现了 97% 的线性扩展效率，相关配置已在 NVIDIA NGC MaxText 容器中提供。 无丢弃 MoE 面临的系统级瓶颈 传统基于固定容量的 MoE 通过限制每个专家的 Token 预算来简化计算，但强制丢弃或填充 Token 会导致模型质量下降或计算资源浪费。无丢弃 MoE 保证每个 Token 均得到完整处理，然而由于路由器分配具有倾斜性和不可预测性，各专家接收的 Token 数量各不相同，产生了不规则张量（ragged tensors）。在缺乏专门优化的情形下，GPU 往往因等待 All-to-All 通信和动态形状重新编译而陷入停顿。 Grouped GEMM 与 MXFP8 算子优化 为消除循环调用 GEMM 带来的设备到主机（Device-to-Host）同步开销以及规避填充带来的额外计算，Transformer Engine 采用了 Grouped GEMM（ragged_dot）技术。该算子基于 cuBLAS 和 cuBLASLt 构建，能够在单个 GPU 内核调用中直接处理所有专家实际分配的 Token 区域，保持 Tensor Core 的高利用率，并支持在 Blackwell 架构上应用 MXFP8 块缩放分组量化。 融合分发与聚合的专家并行通信 Transformer Engine 将 Token 路由后的物理分发（Dispatch）与聚合（Combine）整合为紧密融合的内核执行路径。底层采用专门针对不规则、不均衡通信流量优化的 NCCL EP 后端。该机制还包含了 Token 去重技术，即同一 Token 发往同一 rank 的多个专家或远程节点的多个 rank 时，仅在网络中传输一次，到达目标节点后再进行复制，大幅节省了网络带宽。 JAX 主机卸载与 XLA 多流隐藏延迟 为了进一步降低显存与通信瓶颈，训练流程利用 JAX 的重计算与主机卸载 API，将 Query 和 Value 投影结果等中间激活值暂存至主机内存。针对跨节点通信，XLA 将过去串行执行的单流通信改为多流并行调度，配合延迟隐藏调度器（LHS）自动分析死锁风险并重叠 InfiniBand 与 NVLink 数据传输，显著减少了关键路径上暴露的通信开销。 性能表现与多机扩展验证 在 NVIDIA GB300 NVL72 机架硬件上，DeepSeek-V3 671B 的全流程吞吐量实现了 10.4 倍的增长。在扩展至 1,024 块 GPU 时，系统保持了相对 128 卡基线 97% 的扩展效率，克服了大规模集群下通信开销快于计算开销的痛点。未来 NVIDIA 还计划引入 NVFP4、量化与 GEMM 融合算子以及 All-to-All 重叠优化。

Mixture of experts (MoE) has become one of the defining architectural trends in large-scale AI model training. DeepSeek, Qwen, and Mixtral are examples of MoE models that match or exceed the performance of dense model counterparts at a fraction of the training compute. MoE models provide efficient training through conditional computation. Instead of one dense feed-forward network (FFN) shared by all tokens, MoE replaces it with many smaller expert networks and a learned router that decides which Top-K experts to activate. However, making MoE training efficient at scale is challenging. In DeepSeek-V3 training on NVIDIA GB300, an unoptimized baseline achieved just 103 TFLOPS/GPU with inter-GPU communication consuming 84% of accumulated kernel time. With the JAX Python library and NVIDIA Transformer Engine targeted kernel optimizations, that number rose to 1,068 TFLOPS/GPU, a 10.4x improvement. This post discusses how Transformer Engine, a library for accelerating Transformer models on NVIDIA GPUs, with JAX leads to significant performance improvement in MoE model operations. What are the challenges involved in MoE training? Production-scale MoE training introduces bottlenecks that don’t exist with dense models: token routing, expert dispatch and gather, all-to-all communication, and ragged expert GEMMs. The problem compounds because the router is learned. Throughout training, the distribution can become heavily skewed as the router develops preferences for certain experts. No two batches produce the same expert loads, and within a single batch, one expert might receive many more tokens than another. Each expert receives a different number of tokens, so there is no clean rectangular GEMM to batch and dispatch. This results in ragged tensors. In MoE, tokens are routed dynamically to different experts. This means that the number of tokens assigned to each expert varies unpredictably, which results in ragged tensors (Figure 1). This is a challenge because most libraries are highly optimized for tensor operations that expect uniform, rectangular data structures. Figure 1. Ragged tensors in MoE training With expert parallelism (EP), tokens must be dispatched and outputs must be combined and restored to original token order. If the dispatch and combine path is not optimized, communication dominates and GPUs are underutilized. A poorly optimized all-to-all forces GPUs to stall and wait for data before doing any useful work. Solving this requires specialized kernels that can natively handle ragged layouts. This is precisely the problem that Transformer Engine MoE optimizations are designed to solve. How is dropless MoE different from capacity-based MoE? Dropless and capacity-based MoE are two different ways to handle token routing to experts. In dropless MoE, every token is processed by its selected expert no matter how uneven the load. This is attractive for model quality but demanding on the system. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts addressed this by reformulating expert computation as block-sparse matrix multiplication, allowing each expert to operate on a different number of tokens without dropping or padding. This requires new block-sparse GPU kernels, optimized grouped GEMM, and dispatch and combine primitives all designed specifically for variable token counts. In comparison, standard capacity-based MoE training frameworks sidestep the complexity of dynamic routing by constraining it. Each expert is assigned a fixed token budget, and any overflow is either trimmed or padded to fit. This keeps computation regular and hardware-friendly, but it forces a direct tradeoff between model quality and efficiency: drop the overflow tokens and the model trains on incomplete data, or pad to avoid dropping and pay the cost in wasted compute and memory. Figure 2. A comparison between capacity-based and dropless MoE What specialized optimizations are required for dropless MoE? Committing to dropless MoE means the training stack can no longer rely on fixed expert shapes. Every kernel that touches expert computation has to handle variable token counts efficiently. Additionally, it means that each expert’s token count is variable and data-dependent, so the kernels must not only accept dynamic shapes but also work when those shapes are inaccessible on the CPU to enable CUDA graphs and avoid recompilation. Transformer Engine provides the following building blocks that make this approach practical in JAX: A group-aware MXFP8 quantization An MXFP8 grouped GEMM on expert matmuls Optimized EP operations for dispatch and combine Figure 3 shows an expert-parallel MoE layer across two GPUs. The router assigns each token to an expert, dispatch moves tokens to their expert’s GPU. The grouped MLP runs two grouped GEMMs on those variable-length groups, and combine reverses the exchange to restore the original token order. Figure 3. An expert-parallel MoE layer across two GPUs Optimization 1: Grouped GEMM In a dense FFN, every token passes through the same weight matrix. In MoE, the router distributes tokens unevenly so each expert receives a different number of tokens per step, breaking the regular GEMM shape that typical kernels are optimized for. Previous approaches included a loop of GEMM kernels and batched GEMMs. The loop required Device-to-Host copies of token counts. This is on the critical path, which incurs the latency of the Device-to-Host transfer and breaks CUDA graphs. The batched GEMM computed the worst-case token capacity even if fewer tokens were used because they are padded to force fixed expert computation, leading to extra compute. A grouped GEMM solves this by handling all expert matmuls in a single kernel call, each with its actual token count. It computes only the regions with valid tokens and is more performant as a result. Transformer Engine grouped_gemm /ragged_dot backs this with cuBLAS and cuBLASLt, mapping directly onto the best-performing NVIDIA GEMM libraries to deliver full Tensor Core utilization even with irregular expert shapes. On NVIDIA Blackwell GPUs, this path also opens up MXFP8 block scaling for expert matmuls utilizing the Transformer Engine grouped quantization kernels. Optimization 2: Expert parallelism to integrate Dispatch and Combine After the fused router kernels assign each token to its experts, the model must physically move those tokens to the correct devices, process them, and bring the results back. This process breaks into two distinct stages: Dispatch and Combine. Dispatch: Where the token movement occurs: tokens are permuted and sent across GPUs to their assigned experts, a step that involves both local reordering and multi-GPU communication. Combine: Where processed tokens are routed back to their original GPUs and their per-expert results are accumulated. In a naive implementation, these stages run as a serial chain of separate operations, with the GPU stalling between steps, data getting read and written to memory multiple times, and communication sitting mostly idle while compute runs and vice versa. The Transformer Engine EP implementation integrates the Dispatch and Combine stages into a tightly fused kernel path. This integration is powered by NCCL EP , a communication backend tuned specifically for the irregular, imbalanced traffic patterns that expert-parallel routing produces. NCCL EP also employs a token deduplication mechanism: when a token is dispatched to multiple experts on the same rank or to multiple ranks on a remote IB node, it traverses the network only once and is replicated on the receiving node, conserving network bandwidth. EP is the counterpart to grouped GEMM: grouped GEMM handles what happens inside each expert; EP handles everything around it. Additional optimizations Additional optimizations include JAX host offloading and XLA multistreaming collectives. JAX host offloading Intermediate activations don’t have to be saved on device for the entire forward pass. JAX provides rematerialization APIs for offloading activations to host memory. To save memory in DSv3 training, offload the query and value projection results to host. To learn more, see Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading . XLA multistreaming collectives While EP is driven by Transformer Engine NCCL EP, optimized FSDP is handled natively in XLA. By default, XLA runs communication on a single stream, so collectives that could execute in parallel are serialized and some end up exposed on the critical path. Multi-stream collectives let the compiler schedule independent collectives concurrently across separate CUDA streams, overlapping cross-node InfiniBand transfers with intra-node NVIDIA NVLink communication to draw on both fabrics at once rather than waiting on one serialized stream. The Latency Hiding Scheduler (LHS) decides which collectives are safe to overlap by analyzing their replica groups and checking for deadlock risk, so the memory-bandwidth gains are automatic and require no manual annotation. This reduces the percentage of exposed collectives in DSv3 training significantly. What is the training performance impact of MoE in JAX with Transformer Engine? We observed a 10x end-to-end throughput gain on DeepSeek-V3 671B through MoE in JAX with Transformer Engine optimizations. Recall that the baseline JAX training stack was leaving most of the hardware potential on the table. Tackling the stack at each layer, we added cuBLAS GroupedGEMM, XLA multistream collectives, MXFP8 GroupQuant, host activation offloading, and finally an optimized EP implementation. We plan to add NVFP4 , quantization fused with GEMM, and A2A overlap. To learn more about future kernel fusions that will be supported in Transformer Engine JAX bindings, see Boosting MoE Training Throughput with Advanced Fusion Kernels . Figure 4. End-to-end training performance of DeepSeek-V3 671B on MaxText, measured in TFLOPs/GPU/sec and tokens/GPU/sec across four progressive optimization stages on NVIDIA GB300 NVL72 rack hardware High multirack scaling performance with JAX Training large models at scale demands aggressive optimization. At production scale, this amounts to trillions of tokens and massive batch size inefficiencies. While these are negligible on a single node, they can compound quickly across thousands of GPUs, making every bottleneck in compute, memory, and communication critical to address. Multirack scaling is where most systems struggle, as communication overhead tends to scale faster than compute. With JAX MoE and Transformer Engine stack applied, this degradation stays remarkably in check. The system sustains 97% efficiency at 1,024 GPUs, a result that speaks directly to the effectiveness of the underlying communication optimizations in preserving throughput as the cluster grows. Figure 5. Multirack scaling performance of DeepSeek-V3 671B on NVIDIA GB300 NVL72. TFLOPs/GPU/sec (green bars, left axis) and scaling efficiency relative to 128-GPU baseline (black line, right axis) across five GPU counts from 128 to 2,048 How to get started with dropless MoE training The optimizations ship in the NVIDIA NGC MaxText container with Transformer Engine built in, so you can reproduce and build on them directly. To get started, try the optimized JAX MoE path using the NVIDIA NGC MaxText container with Transformer Engine enabled . Start with the reference configuration, validate correctness on a small MoE model, then scale up while tracking step time, TFLOPS/GPU, MFU, grouped GEMM latency, and MoE dispatch/combine latency. Basic usage configuration: TE MoEBlock with MaxText To enable the TE MoEBlock in MaxText, add the following flags to your MaxText YAML config or pass them as command-line arguments to the training script. Container Use the container from September 10, 2026 ( ghcr.io/nvidia/jax:maxtext-2026-09-10 ) or newer. For more details, refer to the container images section of the NVIDIA/JAX-Toolbox GitHub repo. MaxText configuration ( MaxText moe_configuration.md ): 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 Performance reproduction for DeepSeek V3 To exactly reproduce the DeepSeek-V3 671B results presented in this post, extend the basic usage configuration with the following MaxText config flags, XLA flags, and environment variables. Note that this configuration is specific to DeepSeek-V3; different models will require different tuning. It is not required to use the TE MoEBlock itself. MaxText configuration The MaxText configuration is provided below. For parameter details, refer to the MaxText MoE Configuration guide . # Model parameters model_name: "deepseek3-671b" max_target_length: 4096 hardware: "gpu_multiprocess" # Training settings per_device_batch_size: 6 gradient_accumulation_steps: 1 steps: 15 attention: "cudnn_flash_te" remat_policy: "custom" # Transformer Engine MoEBlock with MXFP8 grouped GEMMs quantization: "te_fp8_currentscaling" te_moe_block: true te_gmm_quantization: "te_mxfp8" ragged_buffer_factor: 2.0 te_ep_overflow_check_every_n_steps: 20 prefuse_moe_weights: true weight_dtype: "bfloat16" mu_dtype: "bfloat16" # Features pgle: true profiler: "xplane" scan_layers: true zero_one: false shardy: true use_segment: false skip_first_n_steps_for_profiler: 4 custom_remat_enabled: true logits_dot_in_fp32: false use_iota_embed: false custom_remat_config: mlpwi: device mlpwi_0: device mlpwi_1: device mlpwo: device moe_mlpwi_0: offload #remat moe_mlpwi_1: offload #remat moe_mlpwo: device query_proj: remat #offload key_proj: remat value_proj: remat #offload query_wa_proj: device kv_wa_proj: device out_proj: device context: device # MoE routing parameters n_routing_groups: -1 topk_routing_group: -1 capacity_factor: 1.0 megablox: false # 128 GPUs: total FSDP=16 (ICI 8 × DCN 2) × EP=8. nodes: 32 ici_data_parallelism: 1 ici_fsdp_parallelism: 8 ici_tensor_parallelism: 1 ici_expert_parallelism: 8 dcn_data_parallelism: 1 dcn_fsdp_parallelism: 2 dcn_tensor_parallelism: 1 dcn_expert_parallelism: 1 shard_optimizer_over_data: false shard_exp_on_fsdp: false XLA flag tuning For guidance on XLA flag tuning, refer to the XLA GPU flags guide and JAX Toolbox GPU performance guide . xla_gpu_all_reduce_combine_threshold_bytes: 33554432 xla_gpu_all_gather_combine_threshold_bytes: 6442450944 xla_gpu_reduce_scatter_combine_threshold_bytes: 201326592 xla_gpu_experimental_enable_nccl_symmetric_buffers: false xla_gpu_enable_command_buffer: "'FUSION,CUBLAS,CUDNN,DYNAMIC_SLICE_FUSION'" xla_gpu_experimental_max_unroll_factor: 8 xla_gpu_memory_limit_slop_factor: 99 Environment variables XLA_PYTHON_CLIENT_MEM_FRACTION: 0.88 CUDA_DEVICE_MAX_CONNECTIONS: 16 XLA_PJRT_GPU_HOST_MEMORY_PREALLOCATE: false XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB: 180 Learn more Dropless MoE training preserves model quality, while Transformer Engine grouped GEMM and EP kernels make it efficient at scale. This approach drives a ~10x throughput improvement and 97% scaling efficiency to 1,024 GPUs on DeepSeek-V3 671B. These optimizations ship in the NVIDIA NGC MaxText container with Transformer Engine built in, so you can reproduce and build on them directly. For information on using the Transformer Engine MoE block in MaxText, refer to the MaxText MoE Configuration guide . For more information on Transformer Engine, refer to the Transformer Engine documentation . Acknowledgments Special thanks to Abhinav Goel, MD Fahim Faysal Khan, Jane Liu, Terry Sun, Tj Xu, Ming Huang, Chase Roberts, and Oleg Goncharov for their contributions to MoE enablement and optimization in JAX, XLA, and Transformer Engine. Thanks to Artem Polyakov, Ke Wen, and Subhadeep Bhattacharya for their contributions to NCCL EP and to Igor Safanov for cuBLASLt contributions.