技術要約: PagedWeight

問題提起

Mixture-of-Experts (MoE) 大規模言語モデル (LLM) は、トークンごとにエキスパートのサブセットのみを活性化することで、高い効率性と精度を実現します。しかし、サービングのシナリオにおいて、MoEモデルは深刻なメモリの緊張状態に直面します。すなわち、GPUは膨大なモデルの重みと、増大し続けるKey-Value (KV) キャッシュの両方を保持しなければなりません。PagedAttentionのような技術はKVキャッシュの断片化を管理しますが、Mo0エキスパートの重みが占める大幅なメモリ・フットプリントの問題は解決していません。これらはGPUメモリの60%以上を占めることがあります。

既存の量子化手法の多くは静的であり、推論開始前に精度を固定してしまいます。一部の実行時精度適応手法も存在しますが、それらは精度変更を利用して、GPUメモリを重みとKVキャッシュの間で動的に再割り当てすることまでは行いません。その結果、タスクの精度を損なったり推論を中断したりすることなく、KVキャッシュが拡大するにつれて重要度の低いエキスパートの重みから動的にメモリを解放できるシステムが存在しない状況にあります。

手法: PagedWeight

著者らは、エキスパートの重みの精度とKVキャッシュのサイズとのバランスを取るために、実行時にモデルの重みを動的に量子化する、MoE LLMサービングのための新しい管理システムであるPagedWeightを提案しています。本システムは、エキスパートの重みをページ化されたKVキャッシュブロックと同様に扱い、動的なオフロードとリロードを可能にします。

コアコンポーネント

1. Paged Weight 表記:

   • PagedWeightは、層（layer）、エキスパート（expert）、およびモジュールタイプ（gate_up または down）によって識別される、エキスパート単位のリニアブロックの粒度で動作します。
   • 本システムは、関連するルックアップテーブル (LUT) を伴うオーバーレイ・ビットプレーン形式で重みを表現する、Any-Precision LLM (APL) テクノロジーを利用しています。
   • Weight Page Table は、各リニアブロックの所望のビット幅 ($d_i$) と確定済みのビット幅 ($q_i$) を追跡します。ページは、GPU常駐、CPU常駐、または転送状態のいずれかの状態で存在できます。

2. 品質認識型ランタイムプランナー (Quality-Aware Runtime Planner):
   プランナーは、精度低下を最小限に抑えるために、以下の3つの要因に基づいてどの重みを量子化（ビット幅を削減）するかを決定します。

   • オフライン・グローバル感度 (Offline Global Sensitivity): ヘッセ行列による重み付けスコアを使用して、各リニアブロックの量子化に対する固有の感度を推定します。
   • オンライン・ルーティング統計 (Online Routing Statistics): エキスパートの「ルーティング質量」（選択の頻度と重み）を追跡します。頻繁にルーティングされる（「ホットな」）エキスパートは、より高いビット幅の下限値とダメージ倍率によって保護されます。
   • プロンプト残差 (Prompt Residuals): 線形回帰ヘッドを介して予測されるプロンプト固有の補正項です。これにより、量子化の影響は入力シーケンスによって異なることを考慮し、グローバルな感度推定を調整します。

   プランナーは、潜在的なビット幅削減に対する「ダメージ」スコアを計算します。KVキャッシュの圧力がメモリ目標をトリガーすると、プランナーは、目標を達成するために、最もダメージの少ない削減（節約されるバイトあたりの品質が最も高いもの）を貪欲に選択します。

3. 非同期実行とカーネル最適化:

   • 非同期ページ移動: レイテンシを隠蔽するために、システムは重みページをCPU RAMにオフロードし、非同期にリロードします。ビット幅の確定は、安全な境界（例：プランが完全に実行された後）でのみ更新され、推論が中断されないようにします。
   • 融合混合精度カーネル (Fused Mixed-Precision Kernel): カスタムCUDAカーネルは、Any-PrecisionビットプレーンとLUTを直接読み取り、単一の融合操作内で各リニアブロックに対して異なるビット幅をサポートすることで、オーバーヘッドを削減します。

主な貢献

1. システム設計: オンラインAny-Precision MoE重みのためのページ化された重みシステムであるPagedWeightの提案。これは、確定済み/所望のビット幅、ページ状態、およびメモリ消費量を動的に管理します。
2. プランニングアルゴリズム: オフラインの感度、オンラインのルーティング統計、およびプロンプト残差を統合し、メモリ圧迫下で低ダメージの重み削減戦略を選択する、品質認識型のランタイムプランナー。
3. 実行戦略: スループットの損失を最小限に抑えつつ、オフロード/リロードのレイテンシを隠蔽する非同期ページ移動の実装、および融合混合精度MoEカーネル。
4. 実証的検証: 3つのMoEモデル（Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B）にわたる包括的な評価により、既存のベースラインに対する優れた性能を実証。

結果

著者らは、言語モデリング、推論、およびロングコンテキスト・タスクにおいて、PagedWeightをFP16、Any-Precision LLM (APL)、DP-LLM、およびMxMoEと比較評価しました。

• 品質とメモリのトレードオフ: PagedWeightは、ベースラインと比較して、最大72.0%のGPUメモリ節約と1.94倍のスループット向上を実現しつつ、FP16と同等の精度を達成します。
• 固定予算における品質: 同様のメモリ予算において、PagedWeightは量子化手法よりも最大39.3%向上したタスク品質を実現し、スループットの低下は最大**4.1%**に留まりました。
• ロングコンテキスト性能: 厳しいメモリ予算（例：10 GB）の下でも、PagedWeightはFP16に匹敵するロングコンテキストの品質（Passage Retrieval, NarrativeQA）を維持しますが、静的な量子化ベースライン（APLなど）は大幅な劣化を起こします。
• スループット: PagedWeightは一様量子化ベースラインのスループットに匹拓しており、観測された最大の低下はバッチサイズ4において4.1%でした。
• アブレーション研究: ルーティング統計、プロンプト残差、またはページ移動などのコンポーネントを取り除くと、一貫してパフォーマンスが低下し、フルシステム設計の必要性が確認されました。

意義

本論文は、PagedWeightが、モデルのタスク精度、メモリ消費量、およびスループット/レイテンシの間の複雑なトレードオフを効果的にナビゲートすることを主張しています。エキスパートの重みを静的な資産ではなく、ページ化された管理可能なユニットとして扱うことで、PagedWeightは、タスク精度と、重みおよびKVキャッシュ間のGPUメモリ割り当ての間の、未開拓のトレードオフ空間を切り拓きます。本システムは、動的で品質を認識した量子化が、静的な量子化や実行時のメモリ圧力に適応しない手法と比較して、MoEサービングにおける実行可能な、かつ優れた戦略であることを示しています。