📖 AI 总结

NVIDIA 发布技术博客，介绍如何借助 JAX 与 NVIDIA Transformer Engine 加速无丢弃（Dropless）混合专家（MoE）模型训练。文章指出，MoE 通过条件计算提升训练效率，但大规模训练面临令牌路由、专家分发与聚合、全对全通信等瓶颈。在 GB200 上训练 DeepSeek-V3 时，未优化基线仅达 103 TFLOPS/GPU，其中 GPU 间通信占内核累计时间的 84%。通过 Transformer Engine 的分组 GEMM 内核、NCCL EP 融合分发与合并阶段并去重令牌，以及 JAX 主机卸载和 XLA 多流集合通信等优化，吞吐量提升至 1068 TFLOPS/GPU，实现 10.4 倍加速。在 GB300 NVL72 上以 1024 块 GPU 训练 671B 参数模型时，仍保持 97% 的扩展效率。