暗い緑色の背景に、地球儀から伸びる線が複数の重み行列のパネルへとつながる抽象的な図

記事

Ai2、「Olmo-core 3」公開——兆パラメータ級MoEを2.7倍速く学習

Allen Institute for AIが、オープンなMoE学習基盤「Olmo-core 3」を公開しました。NVIDIA Megatron-Core採用で学習スループットを最大2.7倍に高め、1兆パラメータ超の構成でもベンチマークを実施しています。

非営利AI研究機関のAllen Institute for AI(Ai2)が、オープンソースの大規模言語モデル「Olmo」シリーズを支える学習基盤の最新版「Olmo-core 3」を公開しました。Hugging Face Blogで公開された技術ブログによると、Olmo-core 3はMoE(Mixture of Experts)モデルの学習を1兆パラメータ級まで効率的にスケールさせるために設計されており、NVIDIAのB300 GPUを使ったベンチマークでは、従来実装と比べて最大2.7倍のスループットを記録したとしています。Ai2はこの基盤を、次世代Olmoモデルの中核システムの一つと位置付けており、学習ツールとインフラを公開し続けるという同社の方針の一環としています。

MoEモデルの効率化という課題

MoEモデルは、入力ごとにすべてのパラメータを使うのではなく、一部の「エキスパート」と呼ばれる専門コンポーネントだけを呼び出すことで、計算量を抑えながらモデルの総容量を増やせる手法です。しかし、モデル全体はGPUメモリ上に保持し続け、学習中も更新し続ける必要があり、さらに入力をクラスタ内の適切なエキスパートへ振り分ける処理自体が通信や調整のコストを生みます。Ai2によると、MoEが大規模化するほど、こうしたコストが「一部のパラメータしか使わない」という本来の計算上の利点を相殺してしまう傾向があるといいます。Olmo-core 3は、この課題を解消することを目的に設計されました。

Googleのブログ記事バナー。「Scaling up training」というタイトルと「How do 512 GPUs train one AI model?」という副題、GPUサーバーラックのイラストが描かれている 画像引用元: Hugging Face Blog

ベンチマークで示された効率化の規模

Ai2が示したベンチマークの一つでは、エキスパートの数を8個から128個に増やしながら、1トークンあたりに選択するエキスパート数は4個のまま据え置いたとしています。これにより、1トークンあたりの有効パラメータ数はおよそ32億でほぼ一定に保たれた一方、総パラメータ容量は46億から470億まで拡大し、学習スループットの低下は5%未満に収まったとのことです。同じ学習基盤は、1兆パラメータを超える規模でも動作確認が済んでいるとAi2は説明しています。さらに、NVIDIAのB300 GPU 512基を使った構成では、1トークンあたり583.6億パラメータが有効になる1.2兆パラメータモデルで、GPU1基あたり最大858TFLOP/sのスループットを観測したとしており、これはランダムなルーティングでシステム性能を測定した結果であり、学習済みモデルの品質を示すものではないと注記されています。

NVIDIA Megatron-Coreへの移行とMXFP8

これまでOlmo-coreのMoE実装は、学習データの小バッチごとにモデルの重みを収集・再分散する「Fully Sharded Data Parallelism(FSDP)」という方式を使っていました。Olmo-core 3では、大規模MoEの学習で実績のあるNVIDIAのMegatron-Coreをベースにしたシステムへ切り替え、ルーティングと計算の効率を高める複数の最適化を組み合わせています。NVIDIA B300 GPU8基を使った予備テストでは、470億パラメータのMoEにおいて、新しいスタックが1GPUあたり毎秒5万2,000トークンを処理し、従来のFSDPベース実装の1万9,400トークンと比べて約2.7倍のスループットを達成したと報告されています。

また、Olmo-core 3は低精度の数値形式「MXFP8」にも対応しました。4基のNVIDIA B300 GPUを使った制御されたベンチマークでは、ベースラインとして使われた高精度形式「BF16」と比べ、MXFP8を主要な箇所で有効にすることで学習スループットが約21%向上し、ピーク時のアクティブメモリ使用量は103GiBから95GiBへ減少したとのことです。Ai2は、こうした高速化の多くがフィードフォワード計算とエキスパート間のデータ移動によってもたらされたものであり、アテンション機構単体の効果ではないと説明しています。

オープンな学習基盤としての意義

Ai2は、Olmo-core 3の技術レポートにおいて、どのMoE学習手法を採用し、どれを見送ったかという判断の過程も含めて公開するとしています。計算の高速化がデータ移動の増加を招いたり、ビット数を減らす変換処理そのものに時間がかかりすぎて効果が相殺されたりするなど、最適化同士のトレードオフを踏まえてシステム全体を設計した点を強調しており、これはAi2自身だけでなく、オープンな学習スタックを利用する外部の研究者にとっても、学習プロセス全体の構成要素をどう組み合わせるかを把握しやすくする狙いがあるとみられます。大学や小規模な研究機関では、計算コストの高さが最先端モデルの開発を阻む要因になりやすいとされる中、Olmo-core 3のような効率化されたオープンインフラの公開は、MoEモデル開発への参入障壁を下げる取り組みの一つとして注目されそうです。