Ai2は10月1日、Olmo-core 3を公開した。混合エキスパート(MoE、入力ごとに一部のエキスパートモジュールだけを使うモデル構造)の学習システムを再設計したフレームワークだ。
Ai2はこれを次世代Olmoの中核基盤と紹介した。今回の発表に新たに学習したモデル重みは含まれない。
あるテストでは、トークンごとに4つのエキスパートだけを選びつつ、エキスパート数を8から128に増やした。アクティブパラメータは約3.2Bに保たれ、総パラメータは4.6Bから47Bに増えたが、学習スループットの低下は5%未満だったという。
以前のMoE実装は完全シャード化データ並列(FSDP)で、小さなバッチごとに重みを集めて再分割していた。Olmo-core 3は分散データ並列(DDP)ベースに切り替え、エキスパートをGPUに常駐させてデータをそのエキスパートへ送る。
Ai2がNVIDIA B300 8基で行った予備テストでは、47BパラメータのMoEがGPUあたり毎秒52,000トークンを処理した。以前の実装の19,400トークンの約2.7倍だ。
モデルと学習状態は、エキスパート並列、パイプライン並列、分散オプティマイザーで複数のGPUに分割される。分散オプティマイザーは、更新の計算に使うオプティマイザー状態をGPUごとに丸ごと持たず、分けて保存する。
低精度形式のMXFP8にも対応する。B300 4基のテストではBF16よりスループットが約21%高く、ピークのアクティブメモリは103GiBから95GiBに減った。
コードはGitHubのOlmo-coreリポジトリでApache License 2.0のもと公開されている。
注意点:システム測定とモデル品質の区別
Ai2はB300 512基で、トークンあたり58.36Bがアクティブになる1.2兆パラメータモデルをテストし、最大858 TFLOP/s/GPUを観測した。このテストは学習済みモデルの品質ではなくシステム性能を測るため、ランダムルーティングを使った。
DeepEP v2で到達した2.38兆パラメータの構成は、完全な学習ではなく短時間の容量テストだった。Ai2もこれを持続的な学習性能ではなく、到達可能な規模を示す結果と説明している。
数値はすべてAi2の測定で、この記事では再現していない。Ai2自身も、通信と計算を重ねる方式が一部のテストでかえって全体の実行を遅くしたと報告している。
実装を評価するときに併せて見る指標
自ら評価するなら、目標とするルーティング、精度、ネットワークで小さな代表的学習ジョブを組み、メモリと実効スループットを併せて比較できる。
容量の数値や個々のカーネル速度だけでは、学習全体が実用的かを判断しにくい。
Olmo-core 3の兆規模の数値はオープンな学習基盤のシステム限界を示すものであり、次世代Olmoの品質はまだ別の問題だ。