用 NVIDIA Transformer Engine 在 JAX 中加速无丢弃 MoE 训练
原标题：Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine
AI 摘要
NVIDIA 技术博客介绍如何利用 JAX 和 NVIDIA Transformer Engine 加速无丢弃（dropless）MoE 训练。文章指出 MoE 训练面临动态路由、不规则张量、all-to-all 通信等瓶颈，在 DeepSeek-V3 于 GB200 上的基线仅 103 TFLOPS/GPU，通信占内核时间 84%。通过 Transformer Engine 的组感知 MXFP8 量化、MXFP8 分组 GEMM 和优化的专家并行 dispatch/combine 操作，性能提升至 1,068 TFLOPS/GPU，实现 10.4 倍加速。
正文节选
Mixture of experts (MoE) has become one of the defining architectural trends in large-scale AI model training. DeepSeek, Qwen, and Mixtral are examples of MoE models that match or exceed the performance of dense model counterparts at a fraction of the training compute. MoE models provide efficient training through conditional computation. Instead of one dense feed-forward network (FFN) shared by all tokens, MoE replaces it with many smaller expert networks and a learned router that decides which