PyTorch、欧州の公共スパコンで約1兆パラメータMoEを訓練

導入

約1兆パラメータ規模のMixture-of-Experts（MoE）モデルの学習は、通常ならAI専用の大規模クラスタを必要とする。EUROPAコンソーシアムは、EuroHPC Frontier AI Grand Challengeの一環として、研究プロジェクトが共有する公共スーパーコンピューターJUPITERでこの課題に取り組んだ。利用したのはPyTorchとTorchTitanである。

実験モデルの総パラメータ数は9519億で、1トークン当たりに活性化されるのは約328億パラメータ。74層、384個のルーティング対象エキスパート、top-8ルーティング、隠れ層サイズ6144、語彙数25万6000、事前学習系列長65536という構成だ。計算には96ノード、384基のNVIDIA GH200 GPUを用いた。

ボトルネックは計算量だけではない

JUPITERのBoosterノードは、4基のGH200スーパーチップを搭載する。各GPUは96 GBのHBM3を備え、72コアのGrace CPUと120 GBのLPDDR5XメモリにNVLink-C2Cで接続される。一方、ノード間通信はInfiniBandに依存し、ローカルのGPU-CPU接続より大幅に遅い。

この不均衡を踏まえ、チームは384-way FSDPを採用し、各ノード内に4-wayのエキスパート並列を配置した。パイプライン並列は使わず、MoEのtoken all-to-all通信は可能な限りNVLink内に留めた。さらに長い系列長を使い、計算時間を増やして通信を計算の裏に隠す余地を確保した。

効果を生んだ三つの工夫

• FlexShardでMuonの通信を削減。 ほぼすべての重み行列にMuonを使うと、Newton-Schulz更新のために完全な行列を集約・分散する必要が生じる。FlexShardではエキスパート単位でパラメータを配置し、各rankが完全なエキスパート行列をローカル処理できるようにした。密行列の通信中にはエキスパート計算を進める。Newton-SchulzをGrace CPUからGPUへ移すことで、オプティマイザの処理時間は約41.7秒から2.2秒へ短縮され、GPU当たりの性能は441から597 tokens/sへ向上した。
• 活性化チェックポイントをGraceへ退避。 系列長65536では、層境界の活性化だけでGPU当たり約55 GiBを消費する。TorchTitanのチェックポイント機構を拡張し、順伝播後に一部の活性化をGraceメモリへ移し、逆伝播の再計算前に戻した。系列長32768で完全再計算を使うと性能は90から370 tokens/sへ伸び、オフロード後は508 tokens/sに達した。
• FSDPの通信と計算を重ねる。 FSDPは現在の層を計算している間に次の層をall-gatherできる。unshard_async_opを計算ストリーム上で利用し、GPU当たり約14 GiBのHBMを解放した。最終的に384基のGH200で約1300 tokens/s/GPUを実現した。

意義と限界

この事例の意義は、単なる性能値にとどまらない。AI専用に設計されたとは限らない公共研究基盤でも、メモリ階層、並列化、オプティマイザ、通信スケジューリングを一体として設計すれば、最先端規模の事前学習を実行できることを示した。PyTorchの柔軟性とTorchTitanの分散学習機能は、JUPITER向けの細かな調整と、成果のオープンソースへの還元を可能にした。

ただし、この結果をすべての1兆パラメータMoEにそのまま適用できるわけではない。GH200のトポロジー、Graceメモリ容量、系列長、並列化方式に大きく左右されるからだ。より一般的な教訓は、ノード間帯域が制約になるほど、モデル構造とシステムソフトウェアを別々ではなく、統合された設計対象として扱う必要があるという点にある。

出典：PyTorch Blog