CRAFT: Cost-aware Expert Replica Allocation with Fine-Grained Layerwise Estimations

この論文は、大規模言語モデルの推論における負荷分散を最適化し、既存手法よりも最大 1.2 倍のスループット向上を実現する、メモリ制約下で層ごとの細粒度なエントリ複製を提案する「CRAFT」というフレームワークを紹介しています。

要約

この論文は、**「CRAFT（クラフト）」**という新しい仕組みについて書かれています。これは、超大規模な AI（大規模言語モデル）を動かす際にかかる「コスト（お金や計算リソース）」を節約しながら、処理速度を上げるための工夫です。

難しい専門用語を使わず、**「巨大なレストラン」**に例えて説明しますね。

1. 背景：なぜ「専門家（エキスパート）」が必要なのか？

まず、最近の超高性能な AI は、**「ミクスチャー・オブ・エキスパート（MoE）」**という仕組みを使っています。
これをレストランに例えると、以下のような感じです。

• 普通の AI（密なモデル）： 1 人のシェフが、すべての料理（料理の全工程）を一人で担当します。客が増えると、そのシェフがパンクしてしまいます。
• MoE 方式の AI： 厨房には「魚料理の専門家」「パスタの専門家」「デザート専門の専門家」など、100 人以上のスペシャリストがいます。
  • 客が「パスタ」を注文すると、AI の「配膳係（ルーター）」がパスタ専門のシェフにだけ注文を回します。
  • これにより、AI は「全知識」を持ちながら、1 回の処理で必要な計算量（コスト）を抑えられます。

2. 問題点：「忙しすぎるシェフ」と「暇なシェフ」

しかし、このシステムには大きな問題がありました。それは**「偏り」**です。

• 現実の注文： 客は「パスタ」を注文することが多く、「魚料理」を注文することは少ないかもしれません。
• 結果： パスタ専門のシェフは**「超繁忙」で、料理が山積みになり、待ち時間が発生します。一方、魚料理のシェフは「暇」**で、ただ立っているだけです。
• AI の世界： これを「負荷の偏り」と呼びます。特定の AI の専門家（エキスパート）にアクセスが集中すると、その部分で処理が止まってしまい、全体のスピードが遅くなります。

3. 既存の解決策と「無駄」

この問題を解決するために、これまでのシステム（EPLB など）では、**「繁忙なシェフの真似事をする」**という対策をとっていました。

• 既存の対策（複製）： パスタが忙しすぎるなら、「パスタ専門のシェフ」を全員分コピーして、厨房に並べます。
  • 「パスタ A 担当」「パスタ B 担当」「パスタ C 担当」のように、同じ能力を持つシェフを何人か用意し、注文をバラバラに振り分けるのです。
• 問題点： これには**「スペース（メモリ）」**が必要です。
  • 厨房（GPU メモリ）は限られています。
  • 既存のシステムは、「どのシェフが忙しかろうが、全員をコピーする」という**「一律（ユニフォーム）」**のルールを使っていました。
  • 無駄： 「実は魚料理のシェフは全然忙しくないのに、わざわざコピーしてスペースを奪っている」状態でした。これでは、料理（データ）を置く場所（KV キャッシュ）が足りなくなり、逆に遅くなってしまうのです。

4. CRAFT の登場：「賢い配分」でコスト削減

ここで登場するのが、この論文が提案する**「CRAFT」**です。

CRAFT は、「どのシェフが、どれくらい忙しそうか」を事前にシミュレーションして、必要な分だけコピーするという、非常に賢い仕組みです。

• CRAFT の考え方：
  1. 分析： 「パスタ担当は超忙しいから、4 人コピーしよう」「デザート担当は暇だから、コピーしなくて OK」「スープ担当は少し忙しいから、2 人コピーしよう」と、役割ごとに最適な人数を決めます。
  2. 節約： 無駄なコピーを減らすことで、厨房のスペース（メモリ）を節約します。
  3. 再配分： 節約したスペースを使って、**「注文待ちの客（データ）」**をより多く受け入れられるようにします。

5. 具体的な効果：どうなるの？

この「CRAFT」を使うと、以下のようなメリットがあります。

• スピードアップ： 無駄なコピーを減らして、必要な部分にリソースを集中させたため、料理（AI の回答）が出るまでの時間が約 1.14 倍（最大 1.2 倍）速くなりました。
• コスト削減： 無駄なコピー（メモリ使用）を減らすので、同じ性能を出すのに必要な GPU（計算機）の数が減ります。
• 柔軟性： どの AI モデルを使っても、追加の学習なしでそのまま使えます。

まとめ：一言で言うと？

「CRAFT」は、AI の厨房で「忙しすぎる人」だけを選んでコピーし、「暇な人」はコピーしないという、 **無駄を徹底的に省いた「賢い人員配置システム」**です。

これにより、限られた厨房スペース（メモリ）の中で、より多くの客（データ）を、より速く、安くサービスできるようになります。まるで、レストランの経営者が「客の傾向を分析して、必要なスタッフだけを増員する」ことで、お店の回転率を劇的に上げたようなものです。

技術概要

CRAFT: 推論スループット向上のためのコスト意識型エキスパートレプリカ割り当て技術の概要

本論文は、大規模言語モデル（LLM）の推論における**エキスパート混合（Mixture-of-Experts: MoE）アーキテクチャの負荷分散問題を解決するための新しいフレームワーク「CRAFT」**を提案しています。CRAFT は、既存のレプリケーション手法が過剰なメモリ消費を引き起こす課題を克服し、細粒度の層別推定に基づいて最適なレプリカ割り当てを行うことで、メモリ制約下での最大のスループット実現を目指します。

以下に、問題定義、手法、主要な貢献、評価結果、および意義について詳細にまとめます。

---

1. 背景と問題定義

MoE アーキテクチャの課題
MoE は、トランスフォーマーの Dense 層を「ルーター」と「エキスパートのプール」に置き換えることで、推論コストを増大させずにモデルパラメータを拡張する技術です。しかし、大規模な MoE モデルを分散環境（Expert Parallelism: EP）でデプロイする際、以下の問題が発生します。

• エキスパート負荷の偏り（Load Imbalance）: 入力トークンの分布は Zipf 分布に従うため、一部の「ホットなエキスパート」にトークンが集中し、他の「コールドなエキスパート」はアイドル状態になります。これにより、GPU 間の負荷が不均一になり、通信の輻輳や計算の待機が発生します。
• 既存の解決策の限界:
  • エキスパート配置（Placement）: ホットとコールドのエキスパートを同じデバイスに配置することで負荷を平準化しますが、負荷の偏りが極端な層では効果が限定的です。
  • エキスパートレプリケーション（Replication）: ホットなエキスパートを複製して複数のデバイスに分散させる手法（例：EPLB）は負荷分散に有効ですが、メモリオーバーヘッドが巨大です。特に大規模モデルでは、レプリカ保存に必要なメモリが KV キャッシュの容量を圧迫し、結果としてスループットが低下する「過剰レプリケーション」の問題が顕在化しています。

核心となる課題
「どの層に、どの程度のレプリカを割り当てるべきか」を決定する際、既存手法は均一なレプリケーション（すべての層に同じ数のレプリカを割り当てる）を採用しており、負荷分散の恩恵が小さい層でもメモリを浪費しています。

---

2. 提案手法：CRAFT

CRAFT（Cost-Aware Replica Fine-grained Throughput）は、メモリ予算の制約下で、推論スループットを最大化するためのコスト意識型エキスパートレプリカ割り当てフレームワークです。

2.1 主要な洞察

著者は、大規模 MoE モデルの特性について以下の 2 つの重要な洞察を得ました。

1. レプリケーションの恩恵は層によって異なる: 負荷が偏っている層（High-skew）ではレプリケーションによる負荷分散効果が大きい一方、負荷が均一な層（Low-skew）では効果がほとんどありません。
2. レプリカ数の増加に対する限界効用逓減: レプリカ数が増えるにつれて、負荷分散の改善度（Balancedness gain）は急速に減少します。既存の均一レプリケーション（例：EPLB）は、多くの場合、メモリコストに見合わない過剰なレプリカを割り当てています。

2.2 システムのワークフロー

CRAFT は以下の 3 つのステップで動作します。

1. 層別レプリケーション恩恵の推定（Fine-Grained Estimation）:

   • オフラインで収集したエキスパート負荷分布に基づき、各層ごとに異なるレプリカ数（1, 2, 4, 8...）を適用した場合の「負荷分散の改善度（Balancedness gain）」を推定します。
   • これにより、各層がレプリケーションから得られるベネフィットを定量化します。

2. 最適レプリカ数の計画（Optimization）:

   • 全体のメモリ予算（許容されるレプリカ総数）を制約条件として、**多選択ナップサック問題（Multiple-Choice Knapsack Problem: MCKP）**として定式化します。
   • 動的計画法（Dynamic Programming）を用いて、各層に割り当てるレプリカ数を決定し、総負荷分散改善度を最大化します。これにより、恩恵の大きい層には多くのレプリカを、恩恵の小さい層にはレプリカを割り当てない（または最小限にする）という最適解を導き出します。

3. 容量意識型エキスパート割り当て（Capacity-Aware Assignment）:

   • 層ごとにレプリカ数が異なる場合、デバイス間のメモリ使用量や KV キャッシュのサイズが不均一になるリスクがあります。
   • CRAFT は、各デバイスのエキスパート容量を均等に保つためのグリーディアルゴリズムとインターリーブ割り当てを採用し、メモリ断片化を防ぎつつ、負荷分散を維持します。

---

3. 主要な貢献

1. レプリケーションのスケーリング特性の初詳細な解明:
   • スループットと負荷分散性がレプリカ数とともにどのように変化するかを体系的に分析し、メモリオーバーヘッドと負荷分散のトレードオフ、および層固有の振る舞いが効率化の鍵であることを示しました。
2. CRAFT フレームワークの提案:
   • 大規模 MoE デプロイメント向けに設計されたエンドツーエンドのレプリカ割り当てフレームワークです。既存の LLM サービングフレームワーク（SGLang など）に追加のトレーニングやモデル変更なしで統合可能です。
3. 広範な評価と実証:
   • 数百億から 1 兆パラメータ規模のモデル（DeepSeek-R1, Kimi-K2 など）を用いた大規模評価を行い、既存手法（EPLB）と比較して平均 1.14 倍（最大 1.2 倍）の推論スループット向上を実現することを示しました。

---

4. 評価結果

実験は AWS p4de.24xlarge インスタンス（8 枚の NVIDIA A100 GPU）を用いたクラスター環境で行われ、DeepSeek-R1-671B や Kimi-K2-1000B などの大規模モデルを対象に評価されました。

• スループット向上:
  • CRAFT は、EPLB（均一レプリケーション）と比較して、平均 1.14 倍、最大 1.2 倍のグッドプット（実質スループット）向上を達成しました。
  • 特に負荷の偏りが激しいデータセットでは、EPLB が過剰なレプリケーションにより KV キャッシュ容量を圧迫しスループットが低下するのに対し、CRAFT はメモリを効率的に利用し、高いスループットを維持しました。
• メモリ効率:
  • CRAFT は、EPLB が使用するレプリカ数の約 1/7〜1/8 のレプリカ数で、ほぼ同等の負荷分散効果を得ています。これにより、KV キャッシュの容量を確保でき、より長いシーケンスや大きなバッチサイズを処理可能になりました。
• スケーラビリティ:
  • クラスターサイズ（6 ノード、8 ノード、12 ノード）が増加しても、CRAFT は EPLB を上回るスケーラビリティを示しました。特に大規模クラスターでは、負荷偏りが顕著になるため、CRAFT の層別最適化の効果がより顕著に現れました。
• オーバーヘッド:
  • 推論時のオーバーヘッドはゼロです。初期化時の推定プロセスは約 10 秒程度であり、フレームワークの起動時間と比較して無視できるレベルです。

---

5. 意義と結論

CRAFT は、大規模 MoE モデルのデプロイにおいて、「メモリ制約」と「負荷分散」のジレンマを解決する重要な技術です。

• コスト削減: 不要なレプリカを排除することで、GPU メモリリソースを効率的に利用し、推論コストを削減します。
• 実用性: 既存のサービングフレームワークに容易に統合でき、モデルの再学習を必要としないため、すぐに導入可能です。
• 将来への示唆: 「均一なリソース割り当て」ではなく、「データと負荷の特性に応じた細粒度の最適化」が、次世代の大規模 AI インフラにおいて不可欠であることを示唆しています。

本論文は、大規模言語モデルの運用効率化において、単なるハードウェアの拡張ではなく、ソフトウェアレベルの高度なリソース管理がスループット向上の鍵となることを実証しました。