Ai2は10月1日、Mixture-of-Experts（MoE）言語モデル向けのオープンな学習フレームワークとしてOlmo-core 3を公開した。システムエンジニアにとって最も有用な結果は、1兆パラメータという見出しより限定的だ。Ai2はコード、参照構成、範囲を限定したGPU測定を公開し、特定のクラスタでこの基盤を試す価値があるかをチームが判断できるようにした。

MoEは各トークンを一部のエキスパートネットワークへ送る。トークンあたりの計算量は抑えられるが、エキスパート全体とオプティマイザー状態は引き続きGPU上に必要であり、エキスパートへのトークン移動にも時間がかかる。Olmo-core 3は、この二つの制約に対するAi2の扱いを変える。DDPベースのシステムはエキスパートをGPUに常駐させ、そこへデータをルーティングする。以前のFSDPベース実装はマイクロバッチごとに重みを繰り返し集約し、再分割していた。新しい基盤はエキスパート並列とパイプライン並列、分散オプティマイザー状態、GPU常駐ルーティング、グループ化したエキスパート計算を組み合わせる。Ai2は一部の演算でMXFP8にも対応する。これらはチームが確認し設定できる設計上の選択であり、別のクラスタでも同じ速度を得られるという保証ではない。

大きな変化

Olmo-core 3はAi2の大規模MoE学習ツールをApache 2.0のコードと参照スクリプトとして公開する。Ai2のテストでは、8台のB300を使う特定の比較で旧基盤より高いスループットが示され、ランダムルーティングの1.2兆パラメータ構成が512台のB300で動作することも示された。これらは、その規模の学習済みモデルの品質や読者のワークロード速度を証明するものではない。実用上の利点は、構成を再現し、ローカルのボトルネックを測り、高コストの学習実行に値する並列化や精度設定を判断するための、検証可能な出発点が得られることだ。

三つの数値、三つの分母

Ai2による47億ではなく470億パラメータMoEの初期比較は、Nvidia B300 GPU 8台で実施された。新基盤は1秒あたり52,000トークンを GPUあたり,従来のFSDPベース実装は19,400だった。比率は52,000 ÷ 19,400 = 2.68で、丸めると2.7倍となる。これは旧基盤と新しい 基盤 の社内比較であり、Megatron-Coreとの測定でも、DDPだけの切り分け実験でもない。タグ付き例で最も近いDDP+エキスパート並列対FSDP/HSDPの比較は、モデル形状と測定時間帯をそろえる一方、ルーター、スパースカーネル、オプティマイザーなど実装の詳細も変えている。

「5%未満」という結果は別の問いに答える。Ai2はエキスパートプールを8から128へ増やし、各トークンで4つを選択した。総容量は46億から470億パラメータに増えた一方、トークンあたりのアクティブパラメータは約32億のままだった。この二つのエキスパートプール構成間で、スループットの低下は5%未満だった この二つのエキスパートプール構成間で。これは新基盤が旧版より5%しか遅くないことも、どの設定でもエキスパートを16倍に増やす費用がかからないことも意味しない。公開情報だけでは、この割合を読者の金額コストに換算できない。

スケール結果には三つ目の分母がある。Ai2は、総パラメータ1.2兆、トークンあたりアクティブパラメータ583.6億の構成をB300 GPU 512台で実行し、最大観測値858 TFLOP/sを GPUあたり 報告した。v3.0.0の参照表では、このUltra-128E行はランダムルーティング、MXFP8、レイヤーごとの再計算とされている。ランダムルーティングはシステム経路を測るもので、学習されたルーターが良いモデルを生むかは検証しない。2.38兆パラメータ構成は短時間の容量試験でDeepEP v2を使用した。Ai2はこれを継続学習ではなく規模の確認と説明する。どちらの数値も、学習済みモデルの精度、有用な出力品質、全工程の安定した経済性を示すものではない。

Ai2はB300 4台の統制比較も報告している。BF16に対しMXFP8はエンドツーエンドの学習スループットを約21%高め、ピーク時のアクティブメモリを103 GiBから95 GiBに下げた。作業はエキスパートへ均等に分散された。この結果が当てはまるのは検証した精度設定とワークロードに限られる。両測定が同じ基準を共有しているかのように2.7倍の数値へ足すことはできない。

エンジニアが実際に試せること

10月1日のリポジトリ v3.0.0 はApache 2.0で公開されている。READMEは、対象OSとハードウェア向けのPyTorchを先にインストールし、PyPIから ai2-olmo-core を導入するか、開発用に複製したリポジトリを使うよう勧めている。タグ付き src/examples/olmo_ddp/README.md は報告された動作点をスクリプトに対応づけ、torchrun を一般的な起動方法として示す。本番スクリプトはノードあたりGPU 8台を前提とし、ノード数を REF_NUM_NODES で指定する。分散GPU環境をすでに持つ人向けの例であり、Ai2のB300結果を一つのコマンドで再現する手段ではない。

レポート用スクリプトには重要なアクセス要件がある。READMEは追加の beaker、FlashAttention 4、書き込み可能な /workspace、Ai2の s3://ai2-llm データミックスへの読み取りアクセスを挙げる。DeepEP v2経路にはDeepEPリポジトリのチェックアウトも必要だ。Ai2のメインREADMEによると、公開Dockerイメージには依存関係は含まれるがOlmo-coreパッケージは含まれず、異なるハードウェアやドライバー/CUDAバージョンではイメージの調整が必要になる場合がある。v3.0.0では最低対応PyTorchが2.10.0に引き上げられた。これらの文書とタグ付き参照表を確認したが、BIG CHANGEはパッケージをインストールせず、データミックスへのアクセスを得ず、GPUを予約せず、学習ジョブも実行していない。

ローカル評価では、意図するモデル規模と利用可能なハードウェアに最も近いタグ付き参照行から始め、計算資源を予約する前に各依存関係とアクセス不能なデータ経路を確認する。同じシーケンス長、バッチ、ルーティング、精度、測定時間帯でGPUあたりの毎秒トークンを比較し、ピークメモリとチェックポイント動作は別途監視する。公開された本番行はシーケンス長8,192、top-fourルーティング、共有エキスパート1つを使用する。TinyとSmallには学習ルーティング用の補助スクリプトがある一方、見出しにある大規模行はランダムルーティングだ。チームの学習ルーティング用データ分布は通信バランスとスループットを変える可能性があるため、全工程を予測する前に代表的な試験が必要となる。

コストの境界はソフトウェアだけでなく物理的な資源にもある。注目されたGPU 8台の比較にはB300が8台必要で、最高スループットの結果には512台が使われた。Ai2は一般的な時間貸し価格や学習全体の予算を公開していない。Apache 2.0はこのコードのソフトウェアライセンス料をなくすが、GPU、互換ドライバー、指定データへのアクセス、運用担当者の時間は提供しない。参照READMEによると、レポートの2バッチ重複測定はマージ前のブランチを使っており、タグ付きスクリプトでは再現できない。

オープンなシステムは、実際の構成と照らして主張を確認できる点で有用だ。Ai2の公開資料は、見かけ上改善したルーティングバランスのスコアが不均衡な作業を隠していた例や、通信と計算の重複が学習を遅くした例をまとめている。だからローカル測定計画が重要になる。高速なカーネルや有利な単独指標は、エンドツーエンドの学習結果ではない。Ai2によれば次世代OlmoはMoEアーキテクチャを使う予定であり、本リリースは将来モデルの品質評価に先立って学習基盤とシステム測定を提供する。

出典・参考資料

  • Ai2の10月1日発表 はアーキテクチャを説明し、範囲を限定したスループット、メモリ、規模の数値を示す。数値はAi2の測定であり、BIG CHANGEは再現していない。
  • Olmo-core v3.0.0リリース はタグ付きバージョン、10月1日の公開日、PyTorch最低要件、コード変更を示す。 リポジトリREADME はインストール、任意依存関係、Dockerイメージの注意点を記載する。
  • タグ付きMoE参照スクリプトガイド はレポート構成をGPU数、ルーティング、精度に対応づけ、旧新基盤の比較上の限界を説明する。 技術レポート はシステム実験の主要資料だが、調査環境ではPDFを取得できなかったため、詳細な主張は閲覧可能なAi2の発表とタグ付きスクリプトに基づく。