Accelerating Dropless MoE Training in JAX with NVIDIA

September 15, 2026September 20, 2026

At a Glance

Values determined by this site’s code at collection time. Dates are JST.

Overview

NVIDIA has announced that combining JAX and NVIDIA Transformer Engine significantly accelerates Dropless MoE (Mixture of Experts) training. This approach improves throughput by 10.4x for DeepSeek-V3 671B training on NVIDIA GB300 environments, achieving 97% scaling efficiency at a 1,024 GPU scale.

Claims and Evidence

According to measurements by the presenters, optimization using JAX and Transformer Engine drastically enhances DeepSeek-V3 training performance. In the conventional unoptimized JAX-based stack, performance on NVIDIA GB300 was limited to 103 TFLOPS/GPU, with inter-GPU communication accounting for 84% of cumulative kernel time. However, applying targeted kernel optimization via Transformer Engine achieved a 10.4x throughput increase, reaching 1,068 TFLOPS/GPU.

The main optimization techniques and their rationale are as follows:

• Handling Ragged Tensors with Grouped GEMM: In MoE, the number of tokens assigned to each expert varies dynamically, resulting in irregularly shaped “Ragged Tensors." Transformer Engine’s grouped_gemm and ragged_dot process these in a single kernel call, leveraging cuBLAS/cuBLASLt to maximize Tensor Core utilization.
• Integrating Dispatch and Combine with NCCL EP: Token transfer (Dispatch) and result collection (Combine) in Expert Parallelism (EP) are implemented as a tightly fused kernel path using the communication backend NCCL EP. This overlaps communication with computation while saving network bandwidth through token deduplication.
• Other Optimizations: These include eliminating memory bottlenecks via JAX host offloading and overlapping NVLink and InfiniBand communication using XLA multistreaming collectives.

Regarding scaling performance, it has been demonstrated that DeepSeek-V3 671B training using NVIDIA GB300 NVL72 hardware maintains 97% scaling efficiency even when utilizing 1,024 GPUs.

Prerequisites

Based on the documentation, the conditions under which these optimization results and reproduction steps hold are:

• Target Model: DeepSeek-V3 671B
• Target Hardware: NVIDIA GB300, NVIDIA GB300 NVL72
• Software / Containers:
  • NVIDIA NGC MaxText container (ghcr.io/nvidia/jax:maxtext-2026-09-09 or later) – JAX – NVIDIA Transformer Engine – NCCL EP – XLA
• Key Configuration Parameters (when reproducing DeepSeek-V3):
  • te_moe_block: true – te_gmm_quantization: "te_mxfp8" – ragged_buffer_factor: 2.0 – sparse_matmul: true – prefuse_moe_weights: true – weight_dtype: "bfloat16" – mu_dtype: "bfloat16"

What Can Be Reproduced Locally

Readers can reproduce and build the optimized JAX MoE path in their own environments by using the NGC MaxText container provided by NVIDIA.

Basic Usage

You can enable TE MoEBlock by adding the following flags to the MaxText YAML configuration file or as command-line arguments to the training script:

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

Reproducing DeepSeek-V3 671B Performance

The documentation discloses detailed configurations to accurately reproduce the benchmark results for DeepSeek-V3 671B. This includes MaxText settings, along with specified XLA flags and environment variables.

MaxText Configuration Example (Excerpt):

model_name: "deepseek3-671b"
max_target_length: 4096
hardware: "gpu_multiprocess"
per_device_batch_size: 6
gradient_accumulation_steps: 1
steps: 15
attention: "cudnn_flash_te"
remat_policy: "custom"
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"
#... (other detailed parameters)

XLA Flag Configuration Example:

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 Variable Configuration Example:

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

However, these optimizations are specialized for the NVIDIA Blackwell architecture (such as the GB300), and operation in other hardware environments is not explicitly stated.

What the Material Does Not Cover

• Comparisons with similar Dropless MoE optimizations using major deep learning frameworks other than JAX (such as PyTorch).
• Performance in hardware environments other than NVIDIA GPUs.
• Specific computational costs for the entire training process or infrastructure setup costs.
• Specific performance figures when running traditional capacity-based MoE methods in the same NVIDIA GB300 environment.

Related Articles

• NVIDIA AIPerf: Benchmarking LLM Inference at Scale

Sources

• https://developer.nvidia.com/blog/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine/

Update History

• 2026-09-20: Rewrote the article from re-collected sources and restored it from draft to published.