Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

技術要約：タスク・アグノスティックを超えて：通信効率の高いマルチタスクMoE推論のためのタスク・アウェアなグルーピング

1. 問題提起

疎に活性化するMixture-of-Experts (MoE) モデルは、条件付き計算を通じて効率的に容量をスケールさせることができるが、分散推論において重大なシステム上の課題、具体的には過度なクロスGPU通信オーバーヘッドと**深刻なロードインバランス（負荷不均衡）**に直面する。既存のエキスパート配置戦略は、頻繁に同時活性化されるエキスパートを近接配置することでこれらのコストを軽減しようと試みているが、それらは通常、タスク・アグノスティック（グローバルに集計された）なルーティング・トレースから単一のグローバルなデプロイメント計画を導出している。

著者らは、このアプローチにおける決定的な欠陥を指摘している。すなわち、エキスパートの同時活性化は強く**タスクに依存（task-conditioned）**しているということである。本論文の分析（図1）で示されているように、あるタスクファミリー（例：コード生成）において密接に結合しているエキスパートのペアは、別のタスク（例：数学的推論）においては相関がないか、あるいは競合することさえある。これらの異質なパターンを単一の集約されたアフィニティ・グラフに平均化してしまうと、タスク固有のエキスパート・グループが崩壊し、通信の局所性が活用されず、オンライン・ワークロードにおける分布シフトへの適応にも失敗する。

2. 手法

提案されたフレームワークであるTask-Aware Coactivation Grouping (TACG) と Generic Expert Shared Replication (GESR) は、以下の2段階のデプロイメント時のプロセスを通じてこれらの問題に対処する。

A. Task-Aware Coactivation Grouping (TACG)
TACGは、タスク条件付きの同時活性化統計量を、容量制約を満たす配置計画へと変換する：

1. タスク・ファミリーの選好度のスコアリング: キャリブレーション・トレースを用いて、システムは各タスク・ファミリーごとの使用ベクトルと同時活性化行列を算出する。各エキスパートについて、特定のタスク・ファミリーに対する「ソフトな選好スコア」($s^{(f)}_{l,e}$) を計算し、これらのスコアを正規化して、エキスパートのファミリーに対する親和性を表す確率分布 ($p^{(f)}_{l,e}$) を作成する。
2. タスク変調同時活性化グラフ: 集約された同時活性化グラフを使用する代わりに、TACGはタスク変調グラフ ($\tilde{A}_l$) を構築する。このグラフは、エキスパートのファミリー選好度の積 ($p^{(f)}_{l,e} p^{(f)}_{l,e'}$) によって標準的な同時活性化エッジを再重み付けする。これにより、同時活性化し、かつ強いファミリー選好を共有するエキスパート間のエッジは増幅され、クロス・ファミリーのエッジは減衰される。
3. 容量制約を満たすグルーピング: フレームワークは、厳格な容量制約の下で、タスク変調グラフ上のグループ内アフィニティを最大化することにより、エキスパートを $M$ 個のグループ（1つにつき1つのGPU）に分割する。これは、2段階のスペクトラル・クラスタリングに続く、正確なGPU負荷制限を確保するためのローカルな「容量修復（capacity repair）」ステップを介して解決される。

B. Generic Expert Shared Replication (GESR)
「ジェネリック（汎用的）」なエキスパート（すべてのファミリーにわたって活性化し、ホットスポットになりやすいもの）によって引き起こされるロードスキューに対処するため、GESRは軽量なコンパニオン・メカニズムを導入する：

1. 特定: グローバルな中心性とクロス・ファミリーの一貫性に報酬を与え、強いタスク固有の偏差にペナルティを与えるスコアを用いて、ジェネリック・エキスパートを特定する。
2. 複製: これらのエキスパートを、構造的な適合性（同時活性化のアフィニティ）に基づいて選択された少数のセカンダリGPUへと複製する。
3. オンライン・ガード: サービング時には、局所性と負荷を考慮した選択ルールにより、指数関数的に減衰する負荷信号を用いて、MoEルーターを変更することなく、ジェネリック・エキスパートにとって最適な候補GPUを選択する。

3. 主な貢献

• 洞察: 本論文は、マルチタスクMoEサービングにおけるエキスパートの同時活性化は強くタスクに依存していることを確立し、効果的な配置は単一のグローバルに平均化されたアフィニティ・グラフではなく、タスク・アウェアな同時活性化によって駆動されるべきであると主張している。
• TACG: エキスパートとファミリーの選好度を用いて同時活性化グラフを再重み付けし、タスク・アグノスティックなグルーピングを厳密に包含する配置を実現する、タスク・アウェアな戦略。
• GESR: クロス・ファミリーのジェネリック・エキスパートを選択的に複製し、ワークロードの変動下でもバランスを維持するために負荷を考慮したオンライン選択を行う、軽量な複製モジュール。
• 実証的検証: 3つの代表的なオープンソースMoEモデル（DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B）を用いた広範な評価により、強力なベースラインに対する一貫した改善を実証した。

4. 実験結果

フレームワークは、4つのタスク・ファミリー（Code, Query, Math, Complex Reasoning）を持つ3つのモデルに対し、16-GPUのエキスパート並列セットアップで評価された。

• 通信の削減: フルフレームワーク（TACG + GESR）は、最も強力な従来の配置戦略（ExFlow）を 12.58% 上回り、ベースラインと比較して平均パー・トークン通信コストを 31.39% 削減した。
• ロードバランス: 本手法は、0.9975 という高い平均Jain公平性指数と、0.0736 という低い平均最大負荷違反（MaxVio）を維持し、通信を大幅に削減しながら、ベースラインと同等以上のロードバランスを実現した。
• 堅牢性: 深刻なタスク分布シフト（Out-of-Distributionのシングルタスク・ワークロード）の下でも、タスク・アウェアなアプローチは安定した通信削減を維持したが、タスク・アグノスティックなバリアントは大きな性能変動とホットスポット形成（CodeタスクにおいてMaxVioが約7倍に急増）に見舞われた。
• アブレーション解析: タスク変調の同時活性化信号を除去すると、通信の利得が31.39%から16.51%に減少した。これは、タスク・コンディショニングが支配的な要因であることを裏付けている。また、GESRは、通信効率を犠牲にすることなく負荷違反（MaxVio）を制御するために不可欠であることが示された。

5. 意義と主張

本論文の意義は、MoEエキスパートの配置のパラダイムを、静的なタスク・アグノスティックな視点から、動的なタスク・アウェアな視点へと移行させたことにあると主張している。異なるタスク・ファミリー間におけるエキスパートの同時活性化の異質性を明示的にモデル化することで、本フレームワークは、グローバルな集約手法では見落とされる通信の局所性を解き放つ。

著者らは、このアプローチが通信と負荷のバランスのフロンティアにおける厳密な改善を提供すると断言している。通信の向上のためにロードバランスを犠牲にする（あるいはその逆を行う）ことが多い従来の手法とは異なり、TACG + GESRは、ほぼ完璧なロードバランスを維持しながら、大幅な通信削減を達成する。本研究は、効果的なマルチタスク・サービングには、それらを平均化してしまうのではなく、タスク固有のエキスパート間の協調という内在的な構造を尊重するデプロイメント戦略が必要であることを強調している。

なお、本論文は、焦点は静的なオフライン配置にあり、リクエスト・ウィンドウ内での動的な再グルーピング、メモリ共同設計（プリフェッチ/キャッシング）、およびドメイン分割モデルにおける特化型エキスパートに関する今後の課題を残しているとして、その範囲について謙虚な姿勢を示している。