Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts

技術要約：計算最適化はクラスター最適化ではない

問題提起
現在の大規模な事前学習ワークフローでは、アルゴリズムのスケール決定とシステム実装が従来、分離されています。「スケーリング則」の段階で、固定された計算量（FLOPs）制約下で損失を最小化するようにモデルアーキテクチャとデータ予算が選択され、その後に、ハードウェア効率のために実装を最適化する、切り離された「システム」の段階が続きます。著者らは、この分離はスパースなMixture-of-Experts (MoE) モデルにおいては欠陥があると考えています。従来のスケーリング則は、モデルのFLOPs ($C_{model}$) を最適化しますが、規定されたアーキテクチャのシステム効率を無視しています。予測される損失と $C_{model}$ が同一である2つのアーキテクチャであっても、モデルFLOPs利用率 (MFU)、通信コスト、およびメモリフットプリントの変動により、トレーニングのスループットは劇的に異なる可能性があります。その結果、「計算最適」な構成は、特定のハードウェアクラスター上で実現可能であったり、あるいは最適であったりするとは限らず、なぜなら、それはクラスターが固定されたトレーニング期間内に実際に維持できる「デリバラブル（提供可能）」な計算量を考慮していないからです。

手法：MOSAIC
著者らは、モデルアーキテクチャとシステムの共同設計を単一の最適化問題として定式化するフレームワークである MOSAIC (Model Optimization via Systems-Aware Training Co-design) を導入します。MOSAICは、予測損失を最小化するために、モデルアーキテクチャ、トレーニング・トークン予算、および分散実行レイアウト（並列化戦略）を共同で選択します。

このフレームワークは、以下の3つのコアコンポーネントで構成されています：

1. 4次元結合スケーリング則： 著者らは、全パラメータ数 ($N_{tot}$)、トレーニング・トークン数 ($D$)、スパース性 ($S$)、およびエキスパート分割係数 ($G$) を含む、スパースMoEモデルのための新しいスケーリング則を適合させました。$G$ は、高密度なフィードフォワード幅と単一のエキスパート幅の比率 ($G = d_{ff}/d_{expert}$) として定義され、エキスパートの粒度を捉えています。この法則は以下のように定式化されます：
   $$L(N_{tot}, S, D, G) = \frac{a}{N_{tot}^\alpha} + \frac{b}{D^\beta} + \frac{c}{(1-S)^\lambda} + \frac{j}{(1-S)^\delta N_{tot}^\gamma G^\eta} + e$$
   この定式化により、抽象的なパラメータ数ではなく、離散的な幾何学的選択に基づいて損失を予測することが可能になります。

2. オペレータレベルのパフォーマンスモデル： システム効率を推定するために、著者らは特定のアーキテクチャと並列化レイアウトに対して MFU とメモリフットプリントを予測する解析的パフォーマンスモデルを開発しました。このモデルは以下を考慮しています：

   • 計算コスト（高密度およびグループ化されたGEMM、アテンション）。
   • 通信オーバーヘッド（エキスパート並列のためのall-to-all、パイプラインバブル、データ並列コレクティブ）。
   • CPUディスパッチレイテンシ（小さなマイクロバッチにおいて重要）。
   • ノード間のストラグラーやノード内オーバーラップのためのハードウェア固有のキャリブレーション係数。
     このモデルは、最大180億のアクティブパラメータを用いたトレーニング実行に対して検証されており、MFU予測における平均絶対パーセント誤差 (MAPE) は15%未満です。

3. 結合最適化 (MINLP)： MOSAICは、要求されるモデルFLOPs ($C_{model}$) が デリバラブル・モデルFLOPs ($C_{deliverable}$) を超えないという制約の下で、予測損失を最小化する混合整数非線形計画問題 (MINLP) を解きます。
   $$C_{deliverable} = C_{peak} \times \text{MFU}(Z, P_\ell, B) \times \eta_{good}$$
   ここで、$C_{peak}$ は生のハードウェア天井値であり、MFUはアーキテクチャに依存します。最適化は、実現可能なジオメトリの離散グリッド ($Z$) と実行プラン ($P_\ell$) を探索し、クラスターの実際の能力内で損失を最小化する構成を見つけ出します。

主な貢献

• 計算予算の再定式化： 本論文は、固定された $C_{model}$ 予算の下では、最適化がデータのサポートの上限（すなわち、可能な限りスパースなモデル）へと向かう「境界探索」の挙動を示すことを実証しています。これは、スケーリング則がシステムコストを無視して、固定FLOPsにおいてよりスパースなモデルの方が低い損失を予測するためです。
• システムを意識した共同設計： スケーリング則をパフォーマンスモデルと結合させることで、MOSAICは予算を $C_{model}$ から $C_{deliverable}$ へとシフトさせます。これにより、最適なスパース性は、境界ではなく、クラスターの特定の制約（メモリ、帯域幅、並列化オーバーヘッド）によって決定される内部点であることが明らかになります。
• パフォーマンスモデルの検証： 著者らは、スパースMoEのための最適な並列レイアウト（テンソル、エキスパート、パイプライン、データ並列）を正確に予測し、最適なレイアウトが特定のジオメトリやクラスターサイズに強く依存することを示す、検証済みのオペレータレベルのパフォーマンスモデルを提供しています。
• 実証的検証： フレームワークは、最大250Bの総パラメータを持つAWS p6ノード（NVIDIA B200 GPU）上でテストされました。結果は、モデルFLOPsに基づく構成のランキングが、ハードウェア・計算効率（ウォールクロック時間）に基づくランキングと大きく異なることを裏付けました。

結果

• 境界 vs 内部最適： 固定された $C_{model}$ の下では、適合された損失はスパース性の増加とともに単調に減少し、最適解を境界 ($S_{max}$) へと押し上げます。しかし、固定されたハードウェア予算 ($C_{deliverable}$) の下では、損失曲線は内部に最小値を示します。極端にスパースなモデルのシステムコスト（通信、メモリ）は、最終的にその理論的な損失メリットを上回ります。
• 損失 vs スループットのトレードオフ： 32ノード、20日間の予算を用いた実験では、最も多くのモデルFLOPsを実現した構成（より高密度で粗いエキスパート・ジオメトリ）が、最低の損失を達成したわけではありません。逆に、損失最適の構成（よりスパースで細かいエキスパート・ジオメトリ）は、より低いMFUを維持しましたが、デリバラブルFLOPsの制約内に収まりつつ好ましい損失軌道を維持できたため、最終的に低い損失を達成しました。
• 実世界での検証： 8ノードのクラスターを用いた段階的な事前学習実行により、フレームワークが検証されました。測定されたMFUの順序は、予測された順序と一致しました。さらに、構成を「ピーク等価ハードウェア計算量」（MFUを考慮したもの）に対してプロットすると、構成のランキングが生のモデルFLOPsに対するランキングと比較して逆転しており、「計算最適」なモデルが「クラスター最適」なモデルではないことが確認されました。

意義
本論文は、最先端の言語モデルのトレーニング方法における根本的な転換を提唱しています。すなわち、切り離された「スケーリング則による選択」と「システムチューニング」の段階から、統一されたアーキテクチャ・システムの共同設計への移行です。著者らは、スパースMoEモデルにおいて、モデルFLOPs予算のみに依存することは不十分であると主張しています。なぜなら、それはアーキテクチャのシステム効率を価格に反映できていないからです。MOSAICを導入することで、彼らは、最適なモデル・ジオメトリは普遍的な定数ではなく、特定のハードウェアクラスター、トレーニング期間、および並列化の制約に依存することを証明しています。このアプローチにより、選択されたアーキテクチャが理論的に損失最適であるだけでなく、ターゲットとなるハードウェア上で実用的かつ効率的であることを保証します。