Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads

要約

大規模言語モデル (LLM) は、長い文脈の理解とタスクの処理のサポートにおいて目覚ましい進歩を示しています。
ただし、LLM の生成推論をこのような長いコンテキストに拡張すると、大幅な追加の計算負荷が発生し、トランスフォーマーベースの LLM のキーバリュー (KV) キャッシュを維持するためにかなりの GPU メモリ フットプリントが必要になります。
量子化などの既存の KV キャッシュ圧縮方法は、コンテキストの長さが増加するにつれてメモリのボトルネックに直面しますが、エビクションなどの静的サイズのキャッシュは非効率なポリシーに悩まされます。
これらの制限により、単一の Nvidia 4090 GPU などのコンシューマー グレードのデバイスへの展開が制限されます。
これを克服するために、私たちは、KV キャッシュ ユニットの因果的重要性を評価する保持ヘッドを導入し、固定キャッシュ サイズ内でより正確な追い出しを可能にするロングコンテキスト LLM 推論のフレームワークである Locret を提案します。
Locret は、標準のロングコンテキスト SFT データセットからの最小限のデータを使用して、凍結されたバックボーン LLM 上で微調整されます。
推論中に、チャンク化されたプレフィル パターンとともに重要度の低いキャッシュ ユニットを削除し、ピーク時の GPU メモリ使用量を大幅に削減します。
私たちは、Locret を評価するために広範な実証研究を実施しています。その実験結果では、メモリ効率と生成されたコンテンツの品質の点で、Locret が InfLLM、Quantization、SirLLM、MInference などの最近の競合アプローチよりも優れていることが示されています。
Phi-3-mini-128K および Llama-3.1-8B-instruct の完全な KV キャッシュと比較して、20 倍および 8 倍の KV キャッシュ圧縮率。
さらに、Locret は、量子化やトークン マージなどの他の方法と組み合わせることができます。
私たちの知る限り、Locret は、Llama-3.1-8B または同様のモデルを単一の Nvidia 4090 GPU にデプロイできる最初のフレームワークであり、生成品質を損なうことなく 128K のロングコンテキスト推論を可能にし、追加のシステム最適化をほとんど必要としません。

要約(オリジナル)

Large language models (LLMs) have shown remarkable advances in supporting long-context comprehension and processing tasks. However, scaling the generation inference of LLMs to such long contexts incurs significant additional computation load, and demands a substantial GPU memory footprint to maintain the key-value (KV) cache of transformer-based LLMs. Existing KV cache compression methods, such as quantization, face memory bottlenecks as context length increases, while static-sized caches, such as eviction, suffer from inefficient policies. These limitations restrict deployment on consumer-grade devices like a single Nvidia 4090 GPU. To overcome this, we propose Locret, a framework for long-context LLM inference that introduces retaining heads to evaluate the causal importance of KV cache units, allowing for more accurate eviction within a fixed cache size. Locret is fine-tuned on top of the frozen backbone LLM using a minimal amount of data from standard long-context SFT datasets. During inference, we evict low-importance cache units along with a chunked prefill pattern, significantly reducing peak GPU memory usage. We conduct an extensive empirical study to evaluate Locret, where the experimental results show that Locret outperforms the recent competitive approaches, including InfLLM, Quantization, SirLLM, and MInference, in terms of memory efficiency and the quality of generated contents — Locret achieves over a 20x and 8x KV cache compression ratio compared to the full KV cache for Phi-3-mini-128K and Llama-3.1-8B-instruct. Additionally, Locret can be combined with other methods, such as quantization and token merging. To our knowledge, Locret is the first framework capable of deploying Llama-3.1-8B or similar models on a single Nvidia 4090 GPU, enabling 128K long-context inference without compromising generation quality, and requiring little additional system optimizations.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google