JAX 中无丢弃 MoE 训练加速：NVIDIA Transformer Engine 实现 10 倍吞吐提升

混合专家（Mixture of experts，MoE）已成为大规模 AI 模型训练中最具代表性的架构趋势之一。DeepSeek、Qwen 和 Mixtral 都是 MoE 模型的例子，它们只用一小部分训练算力，就能达到甚至超过同等稠密模型的性能。

MoE 模型通过条件计算实现高效训练。它不再让所有 token 共用一个稠密前馈网络（FFN），而是用许多更小的专家网络加上一个可学习的路由器来取代它，由路由器决定激活哪 Top-K 个专家。

然而，要让 MoE 训练在大规模下保持高效并不容易。在 NVIDIA GB300 上训练 DeepSeek-V3 时，未经优化的基线每块 GPU 只有 103 TFLOPS，GPU 间通信占用了累计 kernel 时间的 84%。借助 JAX Python 库和 NVIDIA Transformer Engine 有针对性的 kernel 优化，这一数字升至每块 GPU 1,068 TFLOPS，提升了 10.4 倍。本文讨论 Transformer Engine（一个在 NVIDIA GPU 上加速 Transformer 模型的库）如何与 JAX 配合，显著提升 MoE 模型运算的性能。

MoE 训练面临哪些挑战？

生产规模的 MoE 训练会带来稠密模型中不存在的瓶颈：token 路由、专家分发与收集、all-to-all 通信，以及不规则的专家 GEMM。

由于路由器是学习得到的，问题会进一步加剧。在整个训练过程中，随着路由器逐渐偏好某些专家，分布可能变得严重倾斜。没有两个批次会产生相同的专家负载，而在同一个批次内，某个专家收到的 token 可能比另一个多得多。每个专家收到的 token 数量不同，因此没有规整的矩形 GEMM 可供批处理和分发，这就产生了不规则张量（ragged tensor）。

在 MoE 中，token 被动态路由到不同的专家。这意味着分配给每个专家的 token 数量不可预测地变化，从而产生不规则张量（图 1）。这是一个挑战，因为大多数库都针对要求统一、矩形数据结构的张量运算做了高度优化。

在专家并行（EP）下，token 必须被分发出去，输出则必须被合并并恢复到原来的 token 顺序。如果分发与合并路径没有优化，通信就会占据主导，GPU 利用率不足。优化不佳的 all-to-all 会迫使 GPU 停下来等待数据，之后才能做任何有用的工作。

解决这个问题需要能原生处理不规则布局的专用 kernel。这正是 Transformer Engine 的 MoE 优化所要解决的问题。

无丢弃 MoE 与基于容量的 MoE 有何不同？

无丢弃（dropless）MoE 和基于容量（capacity-based）的 MoE 是处理 token 路由到专家的两种不同方式。

在无丢弃 MoE 中，无论负载多么不均，每个 token 都由其选定的专家处理。这对模型质量很有吸引力，但对系统要求很高。MegaBlocks: Efficient Sparse Training with Mixture-of-Experts 通过把专家计算重新表述为块稀疏矩阵乘法来解决这一问题，使每个专家可以处理不同数量的 token，而无需丢弃或填充。这需要新的块稀疏 GPU kernel、优化的分组 GEMM，以及专为可变 token 数量设计的分发与合并原语。

相比之下，标准的基于容量的 MoE 训练框架通过限制动态路由来回避其复杂性。每个专家被分配固定的 token 预算，任何溢出部分要么被裁剪，要么被填充以适配。这使计算保持规整、对硬件友好，但迫使模型质量与效率之间直接取舍：丢弃溢出的 token，模型就在不完整的数据上训练；为避免丢弃而进行填充，则要为浪费的计算和内存付出代价。

无丢弃 MoE 需要哪些专门的优化？

选择无丢弃 MoE，意味着训练栈不能再依赖固定的专家形状。所有涉及专家计算的 kernel 都必须高效处理可变的 token 数量。此外，这还意味着每个专家的 token 数量是可变的且依赖数据，因此 kernel 不仅要接受动态形状，还要在这些形状无法在 CPU 上获取时也能工作，以便启用 CUDA graph 并避免重新编译。

Transformer Engine 提供了以下构建块，使这种方法在 JAX 中切实可行：

• 感知分组的 MXFP8 量化
• 用于专家矩阵乘法的 MXFP8 分组 GEMM
• 针对分发与合并的优化 EP 运算

图 3 展示了跨两块 GPU 的专家并行 MoE 层。路由器为每个 token 分配专家，分发操作把 token 移动到其专家所在的 GPU。分组 MLP 在这些变长分组上运行两次分组 GEMM，合并操作则反向执行交换，恢复原来的 token 顺序。

优化 1：分组 GEMM

在稠密 FFN 中，每个 token 都通过同一个权重矩阵。在 MoE 中，路由器把 token 不均匀地分配出去，因此每个专家每步收到的 token 数量各不相同，打破了典型 kernel 所优化的规整 GEMM 形状。

以往的方法包括 GEMM kernel 循环和批量 GEMM。循环需要把 token 数量从设备复制到主机。这位于关键路径上，既带来设备到主机传输的延迟，又会破坏 CUDA graph。批量 GEMM 即使实际使用的 token 较少，也会按最坏情况的 token 容量计算，因为 token 被填充以强制固定的专家计算，导致额外的计算开销。

分组 GEMM 通过在一次 kernel 调用中处理所有专家的矩阵乘法来解决这个问题，每个专家使用其实际的 token 数量。它只计算含有效 token 的区域，因此性能更高。

Transformer Engine grouped_gemm /ragged_dot 以 cuBLAS 和 cuBLASLt 为后盾，直接映射到性能最佳的 NVIDIA GEMM 库，即使专家形状不规则，也能实现 Tensor Core 的充分利用。在 NVIDIA Blackwell GPU 上，这条路径还借助 Transformer Engine 的分组量化 kernel，为专家矩阵乘法开启了 MXFP8 块缩放。

优化 2：用专家并行集成分发与合并

在融合的路由器 kernel 为每个 token 分配专家之后，模型必须把这些 token 实际移动到正确的设备上、处理它们，再把结果带回来。

这个过程分为两个不同的阶段：分发（Dispatch）和合并（Combine）。

• 分发：token 移动发生在这里：token 被重排并跨 GPU 发送到各自被分配的专家，这一步既涉及本地重排，也涉及多 GPU 通信。
• 合并：已处理的 token 在这里被路由回其原始 GPU，并累加各专家的结果。

在朴素实现中，这些阶段作为一串独立运算串行执行，GPU 在步骤之间停顿，数据被多次读写内存，通信在计算运行时大多闲置，反之亦然。

Transformer Engine 的 EP 实现把分发与合并阶段集成为一条紧密融合的 kernel 路径。这一集成由 NCCL EP 驱动，这是一个专门针对专家并行路由所产生的不规则、不均衡流量模式调优的通信后端。

NCCL EP 还采用了 token 去重机制：当一个 token 被分发到同一 rank 上的多个专家，或远端 IB 节点上的多个 rank 时，它只穿越网络一次，并在接收节点上复制，从而节省网络带宽。EP 与分组 GEMM 相辅相成：分组 GEMM 负责每个专家内部发生的事，EP 负责其周围的一切。

其他优化

其他优化包括 JAX 主机卸载和 XLA 多流集合通信。

JAX 主机卸载

中间激活值不必在整个前向传播期间都保存在设备上。JAX 提供重物化（rematerialization）API，可将激活值卸载到主机内存。要在 DSv3 训练中节省显存，可以把查询（query）和值（value）投影的结果卸载到主机。更多信息请参阅 Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading。

XLA 多流集合通信

EP 由 Transformer Engine NCCL EP 驱动，而优化后的 FSDP 则由 XLA 原生处理。默认情况下，XLA 在单个流上运行通信，因此本可并行执行的集合通信被串行化，其中一部分还暴露在关键路径上。多流集合通信让编译器把相互独立的集合通信并发调度到不同的 CUDA 流上，使跨节点的 InfiniBand 传输与节点内的 NVIDIA NVLink 通信相互重叠，从而同时利用两种互连，而不必在单个串行流上等待。

延迟隐藏调度器（Latency Hiding Scheduler，LHS）通过分析各集合通信的副本组并检查死锁风险，来判断哪些集合通信可以安全地重叠，因此带宽方面的收益是自动获得的，无需手动标注。这显著降低了 DSv3 训练中暴露的集合通信所占的比例。

在 JAX 中结合 Transformer Engine 使用 MoE，对训练性能有何影响？

通过在 JAX 中结合 Transformer Engine 优化的 MoE，我们在 DeepSeek-V3 671B 上观察到端到端吞吐量提升了 10 倍。

回想一下，基线的 JAX 训练栈浪费了大部分硬件潜力。我们逐层攻克这一技术栈，依次加入了 cuBLAS GroupedGEMM、XLA 多流集合通信、MXFP8 GroupQuant、主机端激活值卸载，最后是优化的 EP 实现。

我们计划加入 NVFP4、与 GEMM 融合的量化，以及 A2A 重叠。要了解 Transformer Engine JAX 绑定将来会支持的更多内核融合，请参阅 Boosting MoE Training Throughput with Advanced Fusion Kernels。

JAX 实现出色的多机架扩展性能

大规模训练大模型需要激进的优化。在生产规模下，这意味着数万亿词元和巨大的批大小低效问题。这些问题在单节点上可以忽略，但在数千块 GPU 上会迅速累积，使得计算、内存和通信中的每一个瓶颈都必须解决。

多机架扩展是大多数系统的难点，因为通信开销的增长往往快于计算。应用 JAX MoE 与 Transformer Engine 技术栈后，这种性能衰减被控制得非常出色。系统在 1,024 块 GPU 上保持了 97% 的效率，这一结果直接说明了底层通信优化在集群扩大时保持吞吐量的有效性。

如何开始无丢弃 MoE 训练

这些优化已随内置 Transformer Engine 的 NVIDIA NGC MaxText 容器发布，你可以直接复现并在此基础上构建。要开始使用，请通过启用了 Transformer Engine 的 NVIDIA NGC MaxText 容器体验优化后的 JAX MoE 路径。

先从参考配置开始，在小型 MoE 模型上验证正确性，再逐步放大规模，同时跟踪步时间、TFLOPS/GPU、MFU、分组 GEMM 延迟，以及 MoE dispatch/combine 延迟。

基本用法配置：在 MaxText 中使用 TE MoEBlock

要在 MaxText 中启用 TE MoEBlock，请将以下标志添加到 MaxText 的 YAML 配置中，或作为命令行参数传给训练脚本。

容器

请使用 2026 年 9 月 10 日的容器（ghcr.io/nvidia/jax:maxtext-2026-09-10）或更新版本。更多详情请参阅 NVIDIA/JAX-Toolbox 的容器镜像部分 GitHub 仓库。

MaxText 配置（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

DeepSeek V3 的性能复现

要精确复现本文给出的 DeepSeek-V3 671B 结果，请在基本用法配置的基础上，加上以下 MaxText 配置标志、XLA 标志和环境变量。请注意，该配置专用于 DeepSeek-V3；不同的模型需要不同的调优。使用 TE MoEBlock 本身并不要求这样做。

MaxText 配置

MaxText 配置如下。参数详情请参阅 MaxText MoE 配置指南。

# 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 标志调优

有关 XLA 标志调优的指导，请参阅 XLA GPU 标志指南和 JAX Toolbox GPU 性能指南。

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

环境变量

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

了解更多

无丢弃 MoE 训练保持了模型质量，而 Transformer Engine 的分组 GEMM 和 EP 内核让它在大规模下依然高效。该方法在 DeepSeek-V3 671B 上带来约 10 倍的吞吐量提升，以及 1,024 块 GPU 下 97% 的扩展效率。这些优化已随内置 Transformer Engine 的 NVIDIA NGC MaxText 容器发布，你可以直接复现并在此基础上构建。

有关在 MaxText 中使用 Transformer Engine MoE 块的信息，请参阅 MaxText MoE 配置指南。有关 Transformer Engine 的更多信息，请参阅 Transformer Engine 文档。

致谢

特别感谢 Abhinav Goel、MD Fahim Faysal Khan、Jane Liu、Terry Sun、Tj Xu、Ming Huang、Chase Roberts 和 Oleg Goncharov 对 JAX、XLA 和 Transformer Engine 中 MoE 支持与优化所做的贡献。感谢 Artem Polyakov、Ke Wen 和 Subhadeep Bhattacharya 对 NCCL EP 的贡献，以及 Igor Safanov 对 cuBLASLt 的贡献。