JAXでのDropless MoEトレーニング加速
2026-09-15
本記事では、NVIDIAのTransformer Engineを用いたJAXでのDropless MoE(Mixture of Experts)トレーニングの効率化を解説します。MoEは、トークンを動的に異なる専門家にルーティングし、各専門家が異なるトークン数を処理しますが、これには特有の課題があります。Transformer Engineは、これらの課題を解決するための最適化を提供し、トレーニング性能を約10倍向上させ、1,024 GPUでのスケーリング効率を97%に保ちます。