Ai2、MoE学習のエキスパート配置とデータルーティングを再設計したOlmo-core 3をオープンソース化

Ai2は10月1日、オープンソースの混合エキスパートモデル（MoE）学習システムを再設計したOlmo-core 3を公開した。今回の更新はエキスパートの配置とGPU間のデータ転送に重点を置いており、次世代Olmoの学習基盤となる。研究者が変更や実験を行えるよう、コードも公開された。公式発表

主な変更点は、従来の完全シャーディングデータ並列（FSDP）を中心とした実装から、分散データ並列（DDP）を基盤とする設計への移行だ。エキスパートの重みをGPU上に保持し、トークンデータを対応するエキスパートに送ることで、各ミニバッチで重みを繰り返し集約し、再シャーディングするコストを削減する。公式はNVIDIA B300 8基で470億パラメータのMoEをテストし、GPUあたりのスループットが毎秒19,400トークンから52,000トークンへ向上したと報告している。従来実装の約2.7倍にあたる初期テストの結果であり、他のクラスターにそのまま当てはめることはできない。テストとアーキテクチャの説明

スケーリング能力には明確な限界もある。公式はB300を512基使い、総パラメータ数1.2兆、トークンあたりの有効パラメータ数583億6,000万の構成をテストしたが、システム性能の測定にはランダムルーティングを使用した。この結果は、大規模な構成をフレームワークが扱えることを示すものの、同規模モデルの収束品質、長時間稼働時の安定性、実際の学習コストを証明するには至っていない。公式テストの制約

実運用にあたっては、依存関係や環境の違いへの対応も必要だ。リポジトリではai2-olmo-coreパッケージとソースからのインストール方法が案内されている。一部機能には追加のカーネルパッケージが必要で、READMEによると、dropless MoEで使うgrouped_gemmは自分でコンパイルしなければならない場合がある。公式Dockerイメージにはコアおよびオプションの依存関係が含まれるが、Olmo-core本体は事前インストールされていない。また、ハードウェア、ドライバー、CUDAのバージョンによってはイメージの再構築が必要になる。リポジトリのインストール手順

新バージョンの採用を検討するチームは、次に自社のハードウェアと実際のルーティング負荷で、スループット、ピークメモリ使用量、学習損失を比較し、アーキテクチャ変更を導入する価値があるか判断するとよい。これはテスト範囲に基づくエンジニアリング上の提案だ。コードの公開によって調査や改変はしやすくなったが、異なる環境での効果は再現して確かめる必要がある。