11月12日、PyTorchが「Training a 1T-Parameter MoE on EU public infrastructure with PyTorch and TorchTitan」と題した記事を公開した。EUの公共スーパーコンピュータ上で約1兆パラメータのMixture-of-Expertsモデルを事前学習するにあたり、スループットを90トークン/秒/GPUから約1,300トークン/秒/GPUへ、約14倍に引き上げることに成功した。通信コストの削減・メモリ断片化の解消・CPUオフロードの活用という3つの柱を軸に、PyTorchおよびTorchTitanへの具体的な実装貢献とともにその最適化の全貌が詳述されている。

背景：「公共インフラでは無理」を覆す

フロンティア規模のAI学習は通常、開発者が専有契約で借り切った大規模クラスターで行われる。EUが主導するEuroHPC Frontier AI Grand Challengeは、それを公共インフラで実現できるかを問うプロジェクトだ。

DomynとFraunhofer IAISおよびIISから構成されるEUROPAコンソーシアムは、EU公式24言語対応・EU AI Act準拠・400B以上のオープンMoEモデルという条件のもと選出された。使用したのはドイツ・フォルシュングスツェントルム・ユーリッヒに設置されたスーパーコンピュータJUPITERだ。JUPITERはエクサスケール（1エクサFLOPS＝10¹⁸回/秒の浮動小数点演算）クラスの性能を持つマシンであり、現時点で世界的にも数少ない稼働中のエクサスケール機の一つである。商業クラウドではなく公共機関が運用するこの種のインフラでフロンティア規模の学習を実施したことは、AIインフラの民主化という観点でも注目に値する。

学習したモデルの主要構成は以下のとおり：

JUPITERの各ノードは4基のNVIDIA GH200 Grace Hopperスーパーチップを搭載する。各スーパーチップはHopper GPU（HBM3 96GB）とGrace CPU（LPDDR5X 120GB、72コアArm）が900 GB/sのNVLink-C2Cで接続される構成だ。一方でノード間はInfiniBandのNDR200（GPU当たり約25 GB/s）であり、NVLink比で36倍遅い。この非対称なバンド幅がすべての制約の根本にある。

384基のGH200（96ノード）をDP（Data Parallel）レプリカ単位として使用し、4レプリカで合計1,536 GPUで動作している。

最も効いた最適化：Muon通信をほぼゼロにする「FlexShard」

+35%のスループット改善をもたらしたのがFlexShardによるMuonオプティマイザの通信削除だ。

Muonは各重み行列の更新をNewton-Schulz反復で直交化するオプティマイザで、Adamと比べてfp32状態を半分（パラメータ当たり1つ）しか必要としない。これが65,536トークンのシーケンスをメモリに収めるうえで重要だった。

問題はFSDP（Fully Sharded Data Parallel）との組み合わせだ。FSDPは各行列を384シャードに分割するため、素朴な実装ではNewton-Schulz反復のたびに行列全体を1ランクに集約（gather）→計算→再分散（scatter）する必要がある。InfiniBand越しのこの通信コストは実用に耐えない。

FlexShard（現在TorchTitanに統合済み）は「どのランクがどのパラメータを所有するか」を制御できる仕組みだ。エキスパートをエキスパート単位でシャーディングすることで、各ランクが完全なエキスパート行列を所有し、Newton-Schulz反復を通信なしで実行できる。Attentionの射影など密な行列はオーナーランクへ集約するが、その通信はエキスパートのNewton-Schulz計算に重ねてオーバーラップする。

さらに、CPUオフロード時にNewton-Schulz反復をGrace CPU（72コアArm）からHopper GPU側で実行するようTorchTitanに貢献した結果、オプティマイザステップが41.7秒から2.2秒に短縮され、スループットは441→597トークン/秒/GPU（+35%）となった。

Grace CPUへのアクティベーションオフロード（+37%）

FSDPがHBMを占有しているため、65,536トークンのシーケンスでは境界アクティベーションだけで約55 GiB必要になる計算だ（74レイヤー × 65,536トークン × 6,144 × 2バイト）。これはHBMに収まらない。

GH200のGPU-CPU間900 GB/sリンクを活かし、フォワードパス後に各境界アクティベーションをGrace CPUメモリへオフロードし、バックワードパスの1レイヤー前にHBMへ返す戦略を採用した。TorchTitanのアクティベーションチェックポイントラッパーにオフロード用フックを追加貢献している。これによりスループットは370→508トークン/秒/GPU（+37%）となった。

CUDAアロケータのメモリ断片化問題の解消（+21%）

ソフトウェアスタックのアップグレード後、HBMの空き領域が25〜30 GiBあるにもかかわらずメモリアロケーションが失敗する現象が発生した。

原因はFSDP2のストリーム設計にあった。FSDPはコンピュートストリームと4本のサイドストリームを使い分けるが、CUDAのキャッシングアロケータはストリームごとに独立したメモリプールを管理し、あるストリームが解放したメモリは他のストリームから再利用できない。ピーク時にサイドストリームが保持していた27.3 GiBのうち、実際に使用中だったのは6.4 GiBのみ。残り20.9 GiBがコンピュートストリームから見えない「幽霊メモリ」となり、MoEバッファ用の4 GBブロックを確保できなかった。

FSDP2のunshard_async_opフラグを使い、all-gatherのセットアップをコンピュートストリーム上で実行させることで解決した。ギャザーバッファがコンピュートストリームのプールへ戻るため断片化が解消され、アロケーションリトライが発生するランク数が384→0になった。ピーク予約メモリは88.1→74.6 GiBに減少し、スループットは996→1,210トークン/秒/GPU（+21%）向上した。このフラグは現在公開FSDP APIへの移行が進行中だ。

その他のアップストリーム貢献

記事ではこれら3つ以外にもTorchTitanとPyTorchへの貢献を列挙している：

• RoPEキャッシュの共有（シーケンス長65,536で2.56 GiB/GPU削減）
• バックワードパスでの埋め込みre-gatherのスキップ
• PyTorch本体でのCPU初期化によるRNG状態リセットのバグ修正（#196072）

データロードバランシング、InfiniBand送信バイト削減、ノードローカルリシャーディング、SonicMoEでの動的タイルスケジューリングなどのさらなる詳細は全文ポスト（Domyn）で公開されている。

詳細はTraining a 1T-Parameter MoE on EU public infrastructure with PyTorch and TorchTitanを参照していただきたい。