- AI2が大規模MoE向けの学習基盤「Olmo-core 3」をオープンに公開
- エキスパートをGPUに常駐させるDDPベースの設計へ変更
- 8台のB300で47B MoEのスループットが約2.7倍に向上
何が発表されたのか
AI2は、大規模言語モデルを開発するためのフレームワーク「Olmo-core 3」を公開しました。MoE(mixture-of-experts)の学習システムを再設計したもので、学習を1兆パラメータ規模まで拡張しつつ計算効率を維持することを狙っています。同社は次世代の「Olmo」モデルの中核システムと位置づけており、モデルの重みだけでなく学習基盤も公開するという方針の一環です。
MoEは、入力の単位ごとに一部の専門家(エキスパート)だけを使う仕組みで、全体のパラメータ数を増やしても計算量を抑えられます。ただしモデル全体はGPUメモリに保持して更新し続ける必要があり、入力を正しいエキスパートへ振り分ける通信と調整のコストも発生します。規模が大きくなるほどこのコストが、一部だけを使うことによる利点を打ち消してしまう点が課題でした。
何が新しくなったのか
従来のOlmo-coreのMoE実装は、fully sharded data parallelism(FSDP)を使い、学習データの小さなバッチごとにモデルの重みを集約・再分割していました。Olmo-core 3はdistributed data parallelism(DDP)ベースの方式に切り替え、エキスパートをGPU上に常駐させ、必要なデータをそちらへ配送します。これにより重みの繰り返し集約を避けていると説明されています。
8台のNVIDIA B300 GPUによる予備テストでは、470億パラメータのMoEが新スタックで毎秒52,000トークン/GPUを処理し、従来実装の19,400と比べて約2.7倍だったと報告されています。ルーティングの負荷を下げる工夫として、rowwise expert parallelism、GPU常駐のルーティング、grouped GEMMが挙げられています。加えてMXFP8という低精度の数値形式に対応し、4台のB300での検証ではBF16比でスループットが約21%向上し、ピークのアクティブメモリは103GiBから95GiBに減ったとしています。
どこまで拡張できるのか
Olmo-core 3は、エキスパート数を8から128に増やしても1トークンあたり4エキスパートのみを選ぶ構成で検証されました。トークンあたりのアクティブパラメータを約3.2Bにほぼ固定したまま、総パラメータ数は4.6Bから47Bへ増え、学習スループットの低下は5%未満だったと報告されています。同じ基盤は1兆パラメータ超でもベンチマークされているということです。
大規模MoEの学習ではNVIDIAのMegatron-Coreが確立した選択肢として挙げられていますが、Olmo-core 3が優位だと述べているのは自社の従来FSDP実装との比較であり、Megatron-Coreとの直接比較は示されていません。1.2兆パラメータでトークンあたり58.36Bがアクティブという構成を512台のB300で測定し、最高858 TFLOP/s/GPUを観測したとしています。専門家間の通信を扱う別方式DeepEP v2では2.38兆パラメータの構成に到達しましたが、これは短時間の容量テストであり、継続的な学習性能を示すものではないと明記されています。
現時点の注意点
これらの数値はランダムなルーティングを使ったシステム性能の測定であり、学習済みモデルの品質を評価したものではないと原文は断っています。示されているのはインフラとしての速度と規模であって、その構成で学習したモデルがどれだけ良い性能になるかではありません。
技術レポートでは、ルーティングの偏りを抑えるための指標が改善しても実際の負荷分散は悪化しうる現象(token gerrymandering)や、トークン数が少ないエキスパートの学習率を下げても結果が改善しなかったこと、同じ行列サイズでも値が変わるとGPUの計算時間が変わること、通信と計算のオーバーラップが必ずしも高速化につながらないことなどが報告されています。採用しなかった手法まで含めて公開されている点は、実務でMoEを扱う際の参考になりそうです。
| 項目 | 従来のMoE実装 | Olmo-core 3 |
|---|---|---|
| 並列化方式 | FSDP | DDPベース |
| 重みの扱い | バッチごとに集約・再分割 | エキスパートをGPUに常駐させデータを配送 |
| 47B MoEのスループット(8×B300) | 19,400トークン/秒/GPU | 52,000トークン/秒/GPU(約2.7倍) |
A 47-billion-parameter MoE processed 52,000 tokens per second per GPU with the new stack, compared with 19,400 using our earlier implementation—about 2.7× the throughput.
470億パラメータのMoEは、新しいスタックで毎秒52,000トークン/GPUを処理し、従来実装の19,400と比べて約2.7倍のスループットでした。
今後の見通し
次世代のOlmoはMoEアーキテクチャを採用し、これまでで最大のデータセットと最長のコンテキストウィンドウで学習される予定だと述べられており、Olmo-core 3はその基盤になります。ただし今回示された数値の多くは、ランダムルーティングによるシステム性能の測定や短時間の容量テストです。実際のモデル学習を長時間回したときのスループットや安定性はまだ明らかになっていません。次世代Olmoの学習結果が公開されれば、実運用での実力が判断できるとみられます。加えて、Megatron-Coreとの直接比較やNVIDIA以外のハードウェアへの対応状況が示されれば、採用を検討する際の材料になると考えられます。
日本の開発者・IT企業にとっての意味
MoEは、同じ計算予算でより多くのパラメータを持つモデルを扱えるため、限られたGPU資源でモデルを開発したい企業にとって現実的な選択肢になりつつあります。しかし大規模MoEの学習基盤はこれまで一部の実装に限られ、中身が公開されているものは多くありませんでした。Olmo-core 3は設計と、うまくいかなかった実験まで含めて公開される点に意味があります。日本の開発者にとっては、数十億パラメータ級のMoEであれば8台規模のGPUで検証されている点が参考になります。一方、1.2兆パラメータの検証は512台のB300を前提としており、国内で同じ規模を再現できる組織は限られるとみられます。まずは小規模クラスタでルーティングや並列化の設定を試し、自社のワークロードでスループットがどう変わるかを確かめる使い方が現実的でしょう。
気になる点
- Olmo-core 3を動かすには何台のGPUが必要ですか
- 原文のベンチマークでは、470億パラメータのMoEを8台のNVIDIA B300 GPUで測定し、1.2兆パラメータの構成は512台のB300で測定しています。ただし動作に必要な最低構成や推奨要件は明示されていないため、小規模な環境でどこまで動くかは現時点では公表されていません。
- NVIDIA製以外のGPUでも使えますか
- 原文は、研究者や開発者がOlmo-core 3を異なるハードウェアに適応させられるとしており、オープンであることを強調しています。一方で示されているベンチマークはすべてNVIDIA B300 GPUで実施されたものです。他社製GPUでの動作実績や対応状況については、現時点では公表されていません。
- Megatron-Coreから移行する価値はありますか
- 原文はMegatron-Coreを大規模MoE学習の確立した選択肢として挙げていますが、両者の直接比較は示していません。Olmo-core 3が高速だと述べているのは、あくまで自社の従来FSDP実装との比較です。移行の判断には、同じハードウェアでの比較データが待たれるところです。
用語解説
- MoE(mixture-of-experts)
- 複数の専門家サブネットワークのうち、入力ごとに一部だけを使うモデル構成。全体のパラメータを増やしても計算量を抑えやすい。
- FSDP
- fully sharded data parallelism。モデルの重みをGPU間で分割し、バッチごとに集約・再分割する並列化手法。
- DDP
- distributed data parallelism。各GPUがモデルのコピーを持ち、データを分割して学習する並列化手法。
- エキスパート並列
- 専門家を複数のGPUに分散配置し、各GPUが全体の一部だけを保持するようにする手法。
- MXFP8
- 一部の値をより少ないビット数で表す低精度の数値形式。計算量とGPU間のデータ移動を減らせる可能性がある。
- BF16
- 16ビットの浮動小数点形式。MXFP8より高精度で、比較の基準として使われている。
出典
Introducing Olmo-core 3: Open, scalable training infrastructure for large MoEs
AI導入も顧問も、お任せください
何から手をつけるか、どこまでAIに任せるか。自社の業務に合わせて整理し、導入から運用まで伴走します。相談だけでも構いません。
AI導入について相談する