---
title: "Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine"
description: "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."
---

# Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine

[NVIDIA](https://yomu.fyi/company/nvidia-developer-blog) · Tanya Lenz · Sep 14, 2026

**Type:** Explainer

## 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.

**Tags:** [GPU](https://yomu.fyi/topic/gpu), [JAX](https://yomu.fyi/topic/jax), [NVIDIA](https://yomu.fyi/topic/nvidia), [Parallelism](https://yomu.fyi/topic/parallelism), [Quantization](https://yomu.fyi/topic/quantization)

- Source: [NVIDIA](https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine)
- Source URL: https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine
- Ingested by Yomu: 2026-09-14T20:02:57.048Z

[Read original post](https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine)