Auroraスーパーコンピュータ上で大規模Mixture of Experts言語モデルをスケーラブルに事前学習する

要旨: 大規模言語モデル（LLM）をスクラッチから事前学習するには、大量の計算資源が必要です。Aurora スーパーコンピュータは、127,488 個の Intel PVC（Ponte Vechio）GPU タイルを備えた ExaScale マシンです。本研究では、1000 台規模のGPUタイルにおいて Aurora 上で LLM の事前学習を行うことを示します。この取り組みに向けて、標準的な大規模モデル学習手法をサポートする社内トレーニングライブラリである Optimus を開発しました。Optimus を用いて、まず 3072 台の GPU タイル上で、OLMoE-mix-0924 データセットの全 4 兆トークンに対し、密な 10 億（1 Billion）パラメータモデルである Mula-1B、ならびに 70 億（7 Billion）パラメータの Mixture of Experts（MoE）モデルである Mula-7B-A1B を、スクラッチから学習しました。その後、同じデータセット上で、3 つの大規模 MoE モデル（Mula-20B-A2B、Mula-100B-A7B、Mula-220B-A10B）を 1000 億トークンまで事前学習することで、モデルスケーリングを実証しました。最大規模のモデルである Mula-220B-A10B では、計算スケールを 384 から 12288 GPU タイルへと拡大し、12288 GPU タイルにおけるスケーリング効率がおよそ 90% であることを観測しました。エキスパート計算のためのカスタム GPU カーネルを用いることで、MoE モデルの実行性能を大幅に改善し、さらに新規の EP-Aware シャーディング最適化（sharded optimizer）により、最大 1.71 倍の学習速度向上を達成しました。Optimus ライブラリの一部として、スケールにおける学習の安定性と継続性を高めるための、堅牢な信頼性およびフォールトトレランス機能のセットも開発しました。