推論系統
WIDE、トークンごとにAttentionとFFNを動的にプルーニングし、LLMデコードをエンドツーエンドで1.55倍高速化
WIDEは動的プルーニングをトークン単位のAttentionヘッドグループとFFNチャネルグループにまで細分化し、prefillとdecodeの両方をサポートする。研究では、スパース率50%でコア層のデコードが最大4.95倍高速化した一方、エンドツーエンドの高速化は1.55倍に縮小したと報告している。
大規模言語モデルの構造化プルーニングでは通常、層全体、Attentionヘッド、またはFFNチャネルを固定的に削除する。デプロイは容易だが、トークンの難易度に応じて計算量を調整することはできない。新たに公開されたWIDEは、各層に軽量ルーターを追加し、トークンごとに実行するquery-headグループと連続するFFNチャネルグループを個別に選択する。一方、K、V projectionはdenseのまま維持し、動的な判断によってKV cacheが損なわれるのを防ぐ。
学習は2段階に分かれる。まずbackboneを凍結し、Gumbel-Softmaxを用いてバイナリルーティングのみを学習する。次にLoRAを選択的に導入し、プルーニングによる品質低下を修復する。実行時にはマスクを並べ替え、ハードウェア非依存のブロックスキップと、ハードウェアに合わせて調整するブロック内スキップを組み合わせる。これにより、fine-grained sparsityを単なるFLOPs上の数値に終わらせない。論文によると、スパース率50%において、WIDEは既存のdynamic depth pruningと比べて性能を55.1%向上させた。コア層ではprefillが最大1.98倍、decodeが最大4.95倍高速化し、エンドツーエンドではそれぞれ1.68倍と1.55倍となった。
コードはすでに公開されており、Llama 3.1 8Bの学習パイプライン、FSDP2、評価設定、独立した動的プルーニング用カーネルライブラリが含まれる。リファレンス構成では4基のA100 40GBを使用し、ルーティングのキャリブレーションに約9時間、LoRAによる修復にさらに13時間を要する。つまり、コストゼロで導入できるプラグアンドプレイ型の最適化ではない。
エンジニアリングチームは今後、異なるバッチサイズ、コンテキスト長、GPUアーキテクチャにおけるエンドツーエンド性能に注目すべきだ。論文で示された高い高速化倍率は主にコア層での結果である。実際のサービスでは、依然としてルーティングのオーバーヘッド、メモリ転送、スケジューリング、プルーニング対象外の演算子による制約を受ける。また、量子化モデル、MoE、主要な推論サーバーへそのまま適用できることも、まだ実証されていない。