Ai2、MoE学習基盤Olmo-core 3を公開 1.2Tモデルで性能測定

Ai2は2026年10月1日、大規模MoE向けに再設計したオープンな学習インフラ「Olmo-core 3」を公開しました。8基のNVIDIA B300で47B MoEを学習した予備テストでは、従来実装の19,400 tokens/秒/GPUに対して52,000 tokens/秒/GPUを記録し、スループットは約2.7倍になりました。

何が新しいか

Olmo-core 3は次世代Olmoの基盤となるMoE学習フレームワークです。従来のOlmo-coreのMoE実装はFSDPベースで、小さなバッチごとに重みのgatherとreshardを繰り返していました。Olmo-core 3はDDPベースの方式に移行し、専門家をGPU上に常駐させたまま、該当するデータをそこへルーティングします。これにより重みの反復的なgatherが不要になります。

モデルと学習状態の分散には、次の3つの手法を組み合わせます。

• expert parallelism:専門家をGPU間に分散し、各GPUは専門家プールの一部だけを保持します
• pipeline parallelism:レイヤーをGPUグループ間で分割します
• distributed optimizer:オプティマイザー状態を全GPUに複製せず分散します

ルーティングと計算の最適化には次の3つを使います。

• rowwise expert parallelism:ルーティングしたデータを専門家の入力バッファへ直接配置します
• GPU-resident routing:ルーティングのメタデータをGPU上に保持し、CPUはその情報のコピーを待たずに処理をキューに積めます
• grouped GEMM:多数の小さな専門家計算をまとめて実行します

低精度形式としてMXFP8にも対応しています。

ベンチマーク

項目	条件	結果
専門家数の拡大	1トークンあたり4専門家を選択、活性パラメータ約3.2Bで固定	専門家8→128、総パラメータ4.6B→47B、スループット低下5%未満
47B MoEのスループット	NVIDIA B300 8基	52,000 tokens/秒/GPU(従来19,400、約2.7倍)
MXFP8 vs BF16	B300 4基、専門家間で作業を均一に分配	スループット約21%向上、ピーク活性メモリ103 GiB→95 GiB
1.2Tモデル	B300 512基、活性58.36B/トークン、ランダムルーティング	最高858 TFLOP/s/GPU
2.38T構成	DeepEP v2による短時間の容量テスト	構成を実現(完全な学習実行ではない)

MXFP8による向上の大部分は、attention単体ではなく、feed-forward計算と専門家間のデータ移動から得られました。1.2Tモデルの測定はシステム性能を見るためにランダムルーティングを使っており、学習済みモデルの品質を測ったものではありません。

技術レポートには、ほかに次の知見も記録されています。

• ルーティングの均衡化を促すスコアが改善しても、実際の負荷はかえって偏ることがあります。Ai2はこれを「token gerrymandering」と呼んでいます
• 専門家の学習率を下げても、テストしたモデル群では結果は改善しませんでした
• 行列の次元が同じでも、入力値によってGPUの計算時間が変わることがあります
• 通信と計算を別ストリームでオーバーラップさせると、かえって遅くなる場合がありました

提供と使い方

Ai2はOlmo-core 3をGitHubで公開しています。研究者や開発者は、独自のMoEの学習、異なるハードウェアへの適応、ルーティングや並列化の実験に利用できます。比較対象としては、大規模MoE学習の既存の選択肢であるNVIDIAのMegatron-Coreを挙げています。

背景

Ai2のスパースモデルの取り組みは、64個のルーティング専門家を持つOlmoEにさかのぼります。一方、Olmo 3はdenseアーキテクチャでした。Ai2は次世代OlmoをMoEアーキテクチャで開発し、同社最大のデータセットと最長のコンテキストウィンドウで学習する予定だと説明しています。