利用 NVIDIA Transformer Engine 加速 JAX 下的无丢弃 MoE 训练

混合专家模型（MoE）已成为当前大模型训练的主流架构趋势，如 DeepSeek、Qwen 和 Mixtral 等模型均采用了该方案。针对 MoE 训练过程中的计算与通信开销，技术团队探讨了如何利用 NVIDIA Transformer Engine 在 JAX 深度学习框架下实现无丢弃（Dropless）MoE 模型的高效训练加速，进一步优化大语言模型的训练性能与资源利用率。

原标题：Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine