{"id":639,"date":"2026-09-15T07:17:44","date_gmt":"2026-09-14T22:17:44","guid":{"rendered":"https:\/\/localmodelwatch.tsuchitsuchi.com\/2026\/09\/15\/accelerating-dropless-moe-training-in-jax\/"},"modified":"2026-09-20T17:37:25","modified_gmt":"2026-09-20T08:37:25","slug":"accelerating-dropless-moe-training-in-jax","status":"publish","type":"post","link":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/2026\/09\/15\/accelerating-dropless-moe-training-in-jax\/","title":{"rendered":"Accelerating Dropless MoE Training in JAX with NVIDIA"},"content":{"rendered":"<p><!-- lmw:facts --><\/p>\n<h2>At a Glance<\/h2>\n<div class=\"lmw-table-scroll\" tabindex=\"0\" style=\"overflow-x:auto;-webkit-overflow-scrolling:touch;max-width:100%;\">\n<table style=\"width:max-content;min-width:100%;border-collapse:collapse;\">\n<thead>\n<tr>\n<th>Item<\/th>\n<th>Value<\/th>\n<\/tr>\n<\/thead>\n<tbody>\n<tr>\n<td>Publisher<\/td>\n<td>NVIDIA Developer<\/td>\n<\/tr>\n<tr>\n<td>Source type<\/td>\n<td>Primary source (the publisher itself)<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n<p><em>Values determined by this site&#8217;s code at collection time. Dates are JST.<\/em><\/p>\n<p><!-- \/lmw:facts --><\/p>\n<h2>Overview<\/h2>\n<p>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.<\/p>\n<h2>Claims and Evidence<\/h2>\n<p>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.<\/p>\n<p>The main optimization techniques and their rationale are as follows:<\/p>\n<ul>\n<li><strong>Handling Ragged Tensors with Grouped GEMM<\/strong>: In MoE, the number of tokens assigned to each expert varies dynamically, resulting in irregularly shaped &#8220;Ragged Tensors.&#8221; Transformer Engine&#8217;s <code>grouped_gemm<\/code> and <code>ragged_dot<\/code> process these in a single kernel call, leveraging cuBLAS\/cuBLASLt to maximize Tensor Core utilization.<\/li>\n<li><strong>Integrating Dispatch and Combine with NCCL EP<\/strong>: Token transfer (Dispatch) and result collection (Combine) in Expert Parallelism (EP) are implemented as a tightly fused kernel path using the communication backend <code>NCCL EP<\/code>. This overlaps communication with computation while saving network bandwidth through token deduplication.<\/li>\n<li><strong>Other Optimizations<\/strong>: These include eliminating memory bottlenecks via JAX host offloading and overlapping NVLink and InfiniBand communication using XLA multistreaming collectives.<\/li>\n<\/ul>\n<p>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.<\/p>\n<h2>Prerequisites<\/h2>\n<p>Based on the documentation, the conditions under which these optimization results and reproduction steps hold are:<\/p>\n<ul>\n<li><strong>Target Model<\/strong>: DeepSeek-V3 671B<\/li>\n<li><strong>Target Hardware<\/strong>: NVIDIA GB300, NVIDIA GB300 NVL72<\/li>\n<li><strong>Software \/ Containers<\/strong>:\n<ul>\n<li>NVIDIA NGC MaxText container (<code>ghcr.io\/nvidia\/jax:maxtext-2026-09-09<\/code> or later) &#8211; JAX &#8211; NVIDIA Transformer Engine &#8211; NCCL EP &#8211; XLA<\/li>\n<\/ul>\n<\/li>\n<li><strong>Key Configuration Parameters<\/strong> (when reproducing DeepSeek-V3):\n<ul>\n<li><code>te_moe_block: true<\/code> &#8211; <code>te_gmm_quantization: \"te_mxfp8\"<\/code> &#8211; <code>ragged_buffer_factor: 2.0<\/code> &#8211; <code>sparse_matmul: true<\/code> &#8211; <code>prefuse_moe_weights: true<\/code> &#8211; <code>weight_dtype: \"bfloat16\"<\/code> &#8211; <code>mu_dtype: \"bfloat16\"<\/code><\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<h2>What Can Be Reproduced Locally<\/h2>\n<p>Readers can reproduce and build the optimized JAX MoE path in their own environments by using the NGC MaxText container provided by NVIDIA.<\/p>\n<h3>Basic Usage<\/h3>\n<p>You can enable <code>TE MoEBlock<\/code> by adding the following flags to the MaxText YAML configuration file or as command-line arguments to the training script:<\/p>\n<pre><code class=\"language-yaml\">te_moe_block: true\nte_gmm_quantization: &quot;te_mxfp8&quot;\nragged_buffer_factor: 2.0\nte_ep_overflow_check_every_n_steps: 20\nsparse_matmul: true\nprefuse_moe_weights: true\n<\/code><\/pre>\n<h3>Reproducing DeepSeek-V3 671B Performance<\/h3>\n<p>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.<\/p>\n<p><strong>MaxText Configuration Example (Excerpt):<\/strong><\/p>\n<pre><code class=\"language-yaml\">model_name: &quot;deepseek3-671b&quot;\nmax_target_length: 4096\nhardware: &quot;gpu_multiprocess&quot;\nper_device_batch_size: 6\ngradient_accumulation_steps: 1\nsteps: 15\nattention: &quot;cudnn_flash_te&quot;\nremat_policy: &quot;custom&quot;\nquantization: &quot;te_fp8_currentscaling&quot;\nte_moe_block: true\nte_gmm_quantization: &quot;te_mxfp8&quot;\nragged_buffer_factor: 2.0\nte_ep_overflow_check_every_n_steps: 20\nprefuse_moe_weights: true\nweight_dtype: &quot;bfloat16&quot;\nmu_dtype: &quot;bfloat16&quot;\n#... (other detailed parameters)\n<\/code><\/pre>\n<p><strong>XLA Flag Configuration Example:<\/strong><\/p>\n<pre><code class=\"language-bash\">xla_gpu_all_reduce_combine_threshold_bytes: 33554432\nxla_gpu_all_gather_combine_threshold_bytes: 6442450944\nxla_gpu_reduce_scatter_combine_threshold_bytes: 201326592\nxla_gpu_experimental_enable_nccl_symmetric_buffers: false\nxla_gpu_enable_command_buffer: &quot;'FUSION,CUBLAS,CUDNN,DYNAMIC_SLICE_FUSION'&quot;\nxla_gpu_experimental_max_unroll_factor: 8\nxla_gpu_memory_limit_slop_factor: 99\n<\/code><\/pre>\n<p><strong>Environment Variable Configuration Example:<\/strong><\/p>\n<pre><code class=\"language-bash\">XLA_PYTHON_CLIENT_MEM_FRACTION: 0.88\nCUDA_DEVICE_MAX_CONNECTIONS: 16\nXLA_PJRT_GPU_HOST_MEMORY_PREALLOCATE: false\nXLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB: 180\n<\/code><\/pre>\n<p>However, these optimizations are specialized for the NVIDIA Blackwell architecture (such as the GB300), and operation in other hardware environments is not explicitly stated.<\/p>\n<h2>What the Material Does Not Cover<\/h2>\n<ul>\n<li>Comparisons with similar Dropless MoE optimizations using major deep learning frameworks other than JAX (such as PyTorch).<\/li>\n<li>Performance in hardware environments other than NVIDIA GPUs.<\/li>\n<li>Specific computational costs for the entire training process or infrastructure setup costs.<\/li>\n<li>Specific performance figures when running traditional capacity-based MoE methods in the same NVIDIA GB300 environment.<\/li>\n<\/ul>\n<p><!-- lmw:related --><\/p>\n<h2>Related Articles<\/h2>\n<ul>\n<li><a href=\"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/2026\/09\/19\/nvidia-aiperf-benchmarking-llm-inference\/\">NVIDIA AIPerf: Benchmarking LLM Inference at Scale<\/a><\/li>\n<\/ul>\n<p><!-- \/lmw:related --><\/p>\n<h2>Sources<\/h2>\n<ul>\n<li><a href=\"https:\/\/developer.nvidia.com\/blog\/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine\/\">https:\/\/developer.nvidia.com\/blog\/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engine\/<\/a><\/li>\n<\/ul>\n<p><!-- lmw:updates --><\/p>\n<h2>Update History<\/h2>\n<ul>\n<li>2026-09-20: Rewrote the article from re-collected sources and restored it from draft to published.<\/li>\n<\/ul>\n<p><!-- \/lmw:updates --><\/p>\n","protected":false},"excerpt":{"rendered":"<p>Learn how NVIDIA achieved a 10.4x throughput boost for DeepSeek-V3 671B training using JAX and Transformer Engine on GB300.<\/p>\n","protected":false},"author":1,"featured_media":638,"comment_status":"closed","ping_status":"closed","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1528],"tags":[1205,1207,165,1794,1209,1547],"class_list":["post-639","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-technical-reports","tag-deepseek-v3-en","tag-jax-en","tag-moe-en","tag-nvidia-gb300-en","tag-transformer-engine-en","tag-verified"],"lang":"en","translations":{"en":639,"ja":637},"_links":{"self":[{"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/posts\/639","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/comments?post=639"}],"version-history":[{"count":7,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/posts\/639\/revisions"}],"predecessor-version":[{"id":2315,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/posts\/639\/revisions\/2315"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/media\/638"}],"wp:attachment":[{"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/media?parent=639"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/categories?post=639"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/localmodelwatch.tsuchitsuchi.com\/en\/wp-json\/wp\/v2\/tags?post=639"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}