Ai2がOlmo-Core 3をリリース、兆パラメータMoE向けオープン・トレーニングスタック

Allen Institute for AIは2026年10月1日にOlmo-core 3を発表しました。これは、再設計されたオープンなMixture-of-Expertsトレーニングシステムを中心に構築された大規模言語モデル開発フレームワークのアップグレードで、Ai2は総パラメータ数が1兆を超えるベンチマークを達成したと述べています。

Ai2はOlmo-core 3がMoEトレーニングを兆パラメータ規模にスケールさせつつ計算効率を維持するよう設計されており、次世代Olmoの主要システムの一つであると説明しています。このリリースには技術報告書、インタラクティブデモ、オープンコードが同梱されています。

MoEモデルは、すべての入力がすべてのパラメータを使用しなくても多数のパラメータを保持できますが、完全なモデルは依然としてGPUメモリ全体に保存され、トレーニング中に更新されます。また、クラスタ全体で適切なエキスパートに入力をルーティングすることは、独自の通信・調整コストを生じさせます。Ai2は、これらのコストがMoEが拡大するにつれて計算上の優位性を大きく削ぐ可能性があると指摘し、Olmo-core 3はそのギャップを埋めるよう構築されたと述べています。あるベンチマークでは、エキスパートプールが8から128に増加したにもかかわらず、トークンあたり4つのエキスパートを選択し続け、アクティブパラメータは約32億にほぼ固定されたままで、総パラメータ容量は46億から470億に増加し、トレーニングスループットは5％未満の低下にとどまりました。

FSDPからDDPベースのトレーニングスタックへ

Ai2の以前のOlmo-coreにおけるMoE実装は、完全にシャーディングされたデータ並列（FSDP）を使用し、各小バッチのトレーニングデータに対してモデル重みを集約・再分配するよう構成されていました。Olmo-core 3は、エキスパートをGPU上に常駐させ、関連データをそれらにルーティングする分散データ並列（DDP）ベースのシステムに切り替え、重みの繰り返し集約を回避します。8枚のNVIDIA B300 GPUでの予備テストでは、47億パラメータのMoEがGPUあたり秒間52,000トークンを処理したとAi2は報告しており、以前の実装での19,400トークンと比較して約2.7倍のスループット向上となっています。

Ai2のスパースモデルに関する取り組みはOlmoEに遡ります。OlmoEは64個のルーティングエキスパートを持つMoEアーキテクチャを使用していましたが、Olmo 3はほぼすべてのトークンでモデル全体がアクティブになる密なアーキテクチャを採用しました。Olmo-core 3は、はるかに大規模なMoEモデル向けに設計されたトレーニングシステムでフレームワークを拡張します。Ai2は、NVIDIAのMegatron-Coreが大規模MoEのトレーニングにおける確立されたオプションであることに言及し、Olmo-core 3をOlmoの背後にあるフレームワークへ統合されたMoEトレーニングスタックをもたらすものとして説明しています。

並列化、精度、メモリ技術

ハードウェア上でモデルとトレーニング状態をどのように分割するかは、3つの技術で決まります。エキスパート並列はエキスパートをGPU間に分散させ、パイプライン並列はモデル層をGPUグループに跨がって分割し、分散オプティマイザはオプティマイザ状態を全GPUにフルコピーする代わりにGPU間に分散させます。Olmo-core 3は、データを適切なエキスパートへルーティングするコストも削減します。行指向エキスパート並列はルーティングデータを直接エキスパート入力バッファへ配置し、GPU常駐ルーティングはルーティングメタデータをGPU上に保持してCPUがコピー待ちせずに作業をキューイングできるようにし、グループ化GEMMは多数の小規模エキスパート計算をまとめてGPUが効率的に実行できるようにします。technical reportでは、NVSHMEMベースの行指向エキスパート並列設計が各トークンを直接エキスパートのバッファへ書き込み、デバイススケジュールされたグループ化行列乗算によりエキスパート並列が完全に同期不要になることが記述されています。

Olmo-core 3はMXFP8という低精度数値フォーマットをサポートしています。4枚のNVIDIA B300 GPUでエキスパート間に均等に作業を分配した制御ベンチマークでは、Ai2はBF16ベースラインと比較してトレーニングスループットが約21％向上し、ピークアクティブメモリは103 GiBから95 GiBに減少したと報告しています。この向上の大部分はフィードフォワード計算とエキスパート間データ転送によるものです。レポートは、トポロジーに依存しないチェックポイントが並列レイアウトに関係なくグローバルFP32テンソルを記録できるため、異なるトポロジー上で実行を再開できること、また制御比較においてアクティベーション再計算がアクティベーションメモリの約3分の2を削減し、ピークメモリを約4分の1減らす代わりにスループットが約5分の1低下することを付記しています。

兆パラメータベンチマークと明示された制限

Ai2はNVIDIA B300 GPU上でさまざまな構成でOlmo-core 3をベンチマークしたと述べています。その中には、512 GPUで1.2兆パラメータモデル（トークンあたり583.6億パラメータがアクティブ）を走らせ、最高観測スループットがGPUあたり858 TFLOP/sに達したケースが含まれます。Ai2は、これらのテストがランダムルーティングを用いてシステム性能を測定するものであり、訓練済みモデルの品質を評価したものではないと明言しています。technical reportは、16 GPU上の129億パラメータモデルから512 GPU上の1.2兆パラメータモデルまで測定された動作点を列挙し、ヘッドラインレートはランダムルーティング下で測定されたシステムリファレンスであり、制御されたスケーリング曲線や時間対品質の主張ではないと記しています。

Ai2はDeepEP v2という、GPU間のエキスパート間通信の代替バックエンドでも実験を行い、総パラメータ数2.38兆の構成に到達しました。Ai2はこの結果を、持続的な学習性能ではなく、Olmo-core 3が到達可能な規模を示す短期容量テストと説明しています。技術報告書のタイトルは “Supercharging Olmo-core for Efficient and Scalable MoE Training” で、2026年10月付けです。著者はAllen Institute for AIとUniversity of Washingtonのメンバーで、Tianhua Taoがリードオーサー兼主要開発者として記載されています。

記録された所見と次世代 Olmo の計画

レポートは、Ai2が「トークン・ゲリマンダリング」と呼ぶ失敗パターンを記録しています。このパターンでは、バランスの取れたルーティングを促すことを意図したスコアが、実際のワークロードがより不均衡になるにもかかわらず改善することがあります。その他の記録された所見には、エキスパートが処理するトークンが少なくなるため学習率を下げても、テストされたモデルファミリーでは結果が改善しなかったこと、同一の行列サイズでも処理する値が変わると GPU の計算時間が異なること、別々の GPU ストリーム上で通信と計算をオーバーラップさせても、常に学習が高速化するわけではなく、いくつかのテストではエンドツーエンドの実行が遅くなったこと、が含まれます。レポートはまた、テストされたものの採用されなかったアプローチについても述べています。その中には、ホストリンクが搬送できなかった活性化の CPU オフロードや、エキスパート計算と GPU リソースを争合する二つのオーバーラップ方式が含まれます。

Ai2は、次世代の Olmo が MoE アーキテクチャを採用し、これまでで最も高性能な Olmo になることを目指し、最大規模のデータセットで学習し、最長のコンテキストウィンドウを持たせると述べました。組織は、Olmo-core 3 が完全にオープンであるため、研究者や開発者は自分たちの MoE を訓練したり、さまざまなハードウェアに適応させたり、ルーティングや並列処理、システムの他の部分を実験できると述べています。Olmo-core GitHub リポジトリは、プロジェクトを OLMo エコシステム向けの PyTorch ビルディングブロックとして説明し、Apache-2.0 ライセンスの下で提供され、ソースからまたは PyPI（ai2-olmo-core）からインストール可能です。