NVIDIA 在 JAX 中加速无丢弃 MoE 训练

用途与变化

NVIDIA 开发者博客（2026 年 9 月 14 日）介绍在 JAX 中使用 NVIDIA Transformer Engine 加速无丢弃（dropless）MoE 训练。来源称在 NVIDIA GB300 NVL72 上训练 DeepSeek-V3 671B 时，优化前基线为 103 TFLOPS/GPU，采用该技术栈后达到 1068 TFLOPS/GPU，约 10.4 倍提升；在 1024 GPU 规模下维持 97% 扩展效率。

关键机制

• 无丢弃 MoE 让每个 token 都经其选中的专家处理，不做丢弃或填充，但会产生不规则张量；容量式 MoE 则用固定 token 预算换取计算规整。
• Transformer Engine 提供分组 GEMM（底层对接 cuBLAS/cuBLASLt），可在单次 kernel 调用中处理各专家不同数量的 token，避免逐专家循环和固定容量填充。
• 专家并行通过 NCCL EP 融合 dispatch 与 combine，并做 token 去重以减少网络流量。
• 其他优化包括 JAX 主机内存卸载和 XLA 多流集合通信，以重叠 NVLink 与 InfiniBand 传输。

适用场景与获取方式

来源称相关优化随内置 Transformer Engine 的 NVIDIA NGC MaxText 容器发布，文中给出了 MaxText MoE 配置标志、XLA 标志与环境变量示例，并建议先在小型 MoE 模型上验证正确性，再跟踪 step time、TFLOPS/GPU、MFU 及分组 GEMM、dispatch/combine 延迟逐步扩展。

限制与待核实

本文为厂商博客的自我报告结果，未提供独立复现或第三方验证；所述 10.4 倍提升、97% 扩展效率与 1024 GPU 数据均来自来源，具体配置针对 DeepSeek-V3，其他模型需要不同调优。文中未说明测试协议、精度校验细节与失败场景，实际效果有待核实。