Fast MoE Inference via Predictive Prefetching and Expert Replication

技術的概要：予測的プリフェッチとエキスパート複製による高速 MoE 推論

問題定義
Mixture of Experts (MoE) アーキテクチャは、計算オーバーヘッドの比例増加なしに膨大なパラメータ数を可能にすることで、大規模言語モデル (LLM) のスケーリングの基盤となっています。しかし、MoE 推論は実用的な展開を妨げる重大な非効率性に直面しています：

1. GPU 利用率の最適化不足: エキスパート活性化の固有のスパース性により、任意の時点でアクティブなエキスパートはごく一部のみであり、GPU リソースの大部分がアイドル状態のままです。
2. 負荷の偏りとレイテンシ: 動的ルーティングはしばしば不均一な負荷分散をもたらし、複数のトークンが同じ物理的エキスパートをキューイングすることで、トークンのストールとレイテンシの増加を引き起こします。
3. ルーティングオーバーヘッド: 従来の softmax ベースのルーティングメカニズムは、計算複雑性と通信オーバーヘッド（オールツーオールルーティング）を導入し、エキスパート数の増加に伴って有利にスケーリングしません。

既存の最適化技術（蒸留やカーネルレベルの強化など）は、限られた改善しか提供しないか、モデルの精度を犠牲にすることが多いです。

手法：MoE-MPMC
著者は、予測的プリフェッチと動的エキスパート複製を通じてこれらの非効率性に対処するフレームワーク「MoE-MPMC（Multiple Producers Multiple Consumers を備えた Mixture of Experts）」を提案します。このシステムは 2 つの並行スレッドを使用して動作します：

1. 推論スレッド:

   • 現在のバッチに対するエキスパート割り当てと複製カウントを含む事前計算されたハッシュテーブルを取得します。
   • 予測されたアクティブなエキスパートのみを GPU メモリに動的にロードし、未使用のエキスパートを CPU メモリにオフロードします。
   • アクティブなエキスパートが GPU 上で物理的に複製されるように修正された MoE レイヤーを使用して推論を実行します。これにより、同じエキスパートを必要とする複数のトークンを、異なる物理的レプリカによって並列に処理でき、キューイング遅延を排除します。

2. ハッシュ構築スレッド:

   • 推論スレッドと並行して実行され、次のバッチのトークンを処理します。
   • スパースアテンションを備えた Simple Recurrent Unit (SRU) を使用して、次のバッチで活性化されるエキスパートを予測します。
   • 先行研究（SiDA-MoE）で使用されていた LSTM ベースの予測器を SRU に置き換え、相互依存の減少と単純なゲーティングメカニズムにより、より高速なトレーニングおよび推論速度を実現します。
   • トークンの需要に基づいて各予測エキスパートの複製カウントを計算し、ハッシュテーブルを更新します。その後、このハッシュテーブルは推論スレッドのためにキューに追加されます。

主要な貢献
本論文は、3 つの主要な貢献を概説しています：

1. SRU ベースの予測: 各トークンに対するエキスパート選択を迅速に予測するための Simple Recurrent Unit (SRU) の導入。これにより、遅い LSTM 予測器が置き換えられ、精度を維持しながら予測フェーズが大幅に高速化されます。
2. 深い物理的複製: トークンの需要に基づいて、予測されたエキスパートを GPU 上で物理的に複製する戦略。これにより、バッチ処理中に待機遅延を経験するトークンがほぼなくなり、並列性と GPU 利用率が最大化されます。
3. 容量キャッピング: 定義済みの GPU メモリ制約に応じてエキスパート複製カウントを制限する実用的なメカニズム。これにより、オーバーフローを引き起こすことなくハードウェア制限内で複製が最適化され、ほぼ最適なリソース利用率が達成されます。

実験結果
著者は、SST2 や MRPC などのデータセットを使用して、大規模な SwitchTransformers モデル（Switch-base-8, 64, 128, 256）上で MoE-MPMC を評価しました。主な知見は以下の通りです：

• GPU 利用率: この手法は、ほぼ完全な GPU 利用率（一部の構成では約 100%）を達成し、ベースラインの SwitchTransformers（128 エキスパートで利用率が 1.3% まで低下）や SiDA-MoE アプローチよりも大幅に改善されました。
• 推論速度: このアプローチは、ベースラインアーキテクチャと比較して最大3 倍の推論速度の向上をもたらします。
• 精度: この手法は、ベースラインアーキテクチャの性能の約90-95% を維持し、SRU の LSTM に対する優れた予測能力により、SiDA-MoE よりも精度が良好に維持されました。
• スループット: エキスパート複製による複数のトークンの並列処理能力により、エキスパート数が増加するにつれてスループットが大幅に増加します。

意義と主張
本論文は、MoE-MPMC が MoE ベースの大規模言語モデルの実用的かつ経済的に viable な展開に向けた重要な一歩であると主張しています。反応的な負荷分散から、予測的プリフェッチと複製による先行的戦略への転換により、この手法は：

• メモリアクセスのボトルネックとトークンのストールを軽減します。
• 基盤となる MoE 計算グラフを変更することなく GPU リソース利用率を最適化します。
• トークン予測タスクと並行して微調整と事前学習が行われる場合、エキスパート選択の予測が実行可能かつ効果的であることを実証します。

著者は、この作業を、将来の需要を予測しリソース割り当てを動的に調整することで、スパースアーキテクチャにおけるアイドル GPU メモリや負荷の偏りといった典型的な落とし穴に対処する、反応的な手法（SiDA-MoE など）に対する進歩として位置づけています。