Microsoft、MAI-Base-2の事前学習の詳細を公開　複数ラウンドのMoEスケジューリングでメモリ使用量を制御

Microsoft AIは10月15日、MAI-Base-2の事前学習の成果を発表した。この混合エキスパート（MoE）基盤モデルの総パラメーター数は2.8兆で、トークンごとに約800億パラメーターが有効になる。今年8月から9月にかけて、16,000基のNVIDIA GB200 GPUを使い、38日間学習させた。モデルは現在も後学習中で、MicrosoftはFoundryを通じて信頼できるパートナーに重みをライセンス供与する計画だが、一般向けサービスの開始日は発表していない。発表とスケジュール

レポートで重要なエンジニアリング上の変更として挙げられているのは、MoEエキスパート間の負荷の偏りによるメモリの浪費への対処だ。エキスパートごとに受け取るトークン数は異なるため、極端な負荷を想定して領域を確保すると、通常時には未使用のままになる。一方、容量に上限を設けてあふれたトークンを破棄すると、学習品質を損なう可能性がある。Microsoftは代わりに複数ラウンドのスケジューリングを採用した。各エキスパート群に固定容量を割り当て、超過したトークンは後続のラウンドで処理する。また、負荷の高いエキスパートが同じ群の未使用容量を借りられるようにした。容量を小さくすればメモリを節約できる一方、計算ラウンドが増える可能性がある。公式技術レポート

チームはエキスパートの計算と通信をグループ化し、ホストとGPU間の同期を減らすとともに、通信と計算をオーバーラップさせた。Microsoftによると、これらの改善によりMoE層のモデルFLOPs利用率は24%から35%に向上し、すべての層をMoEにできた。完全な学習全体での平均利用率は21%だが、両者は測定範囲が異なる。こうした詳細は、スパース性を高める実際の利点が、メモリ割り当て、通信、カーネル起動のコストによってなお制約されることを示している。エンジニアリング解説

評価では、Microsoftは保持したテキストの予測損失を比較するため、1バイト当たりのビット数（BPB）を用いてtokenizerの違いによる影響を抑えた。また、スケーリング曲線に基づく推定では、DeepSeek-V4-Pro-Baseと同じ損失を達成するために必要なFLOPsを4.2～12.3倍改善できるという。これは学習レシピの推定効率であり、推論の高速化やAPIコストの削減率に換算できるものではない。比較対象の一部モデルには、すでに後学習が施されている。今後、エンジニアリングチームが注目すべきなのは、再現可能な評価、後学習後の能力、そしてデプロイ時のレイテンシとメモリ要件だ。評価方法と制約