KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

要約

LLMは、大きなコンテキストウィンドウを必要とする文書解析や要約などのアプリケーションで使用されるようになってきており、このような大きなコンテキストウィンドウでは、KVキャッシュの活性化が推論時のメモリ消費に大きく影響します。量子化はKVキャッシュの活性度を圧縮するための有望なアプローチであるが、既存のソリューションでは、4ビット以下のような超低精度で活性度を正確に表現することができない。KVQuantは、キャッシュされたKVアクティベーションを量子化するための新しい手法を取り入れることで、この問題に対処している:(i)チャネルごとのキー量子化。キーアクティベーションを量子化する次元を調整することで、より分布にマッチするようにする。(ii)プリRoPEキー量子化。ロータリ位置埋め込み前にキーアクティベーションを量子化することで、量子化への影響を緩和する;(iv)ベクトル毎のDense-and-Sparse量子化。各ベクトル毎に外れ値を分離し、量子化範囲のスキューを最小化する。(v)Q-Norm。分布のシフトを緩和するために量子化セントロイドを正規化し、2ビット量子化にさらなる利点をもたらす。本手法をLLaMA、LLaMA-2、Mistralモデルに適用することで、Wikitext-2とC4の両方において、3ビット量子化で$<0.1$の当惑度劣化を達成し、既存のアプローチを凌駕した。我々の手法により、単一のA100-80GB GPUで最大100万、8GPUシステムで最大1000万のコンテキスト長を持つLLaMA-7Bモデルを扱うことができる。

要約(オリジナル)

LLMs are seeing growing use for applications such as document analysis and summarization which require large context windows, and with these large context windows KV cache activations surface as the dominant contributor to memory consumption during inference. Quantization is a promising approach for compressing KV cache activations; however, existing solutions fail to represent activations accurately in ultra-low precisions, such as sub-4-bit. In this work, we present KVQuant, which addresses this problem by incorporating novel methods for quantizing cached KV activations, including: (i) Per-Channel Key Quantization, where we adjust the dimension along which we quantize the Key activations to better match the distribution; (ii) Pre-RoPE Key Quantization, where we quantize Key activations before the rotary positional embedding to mitigate its impact on quantization; (iii) Non-Uniform KV Cache Quantization, where we derive per-layer sensitivity-weighted non-uniform datatypes that better represent the distributions; (iv) Per-Vector Dense-and-Sparse Quantization, where we isolate outliers separately for each vector to minimize skews in quantization ranges; and (v) Q-Norm, where we normalize quantization centroids in order to mitigate distribution shift, providing additional benefits for 2-bit quantization. By applying our method to the LLaMA, LLaMA-2, and Mistral models, we achieve $<0.1$ perplexity degradation with 3-bit quantization on both Wikitext-2 and C4, outperforming existing approaches. Our method enables serving the LLaMA-7B model with a context length of up to 1 million on a single A100-80GB GPU and up to 10 million on an 8-GPU system.

arxiv情報

著者 Coleman Hooper,Sehoon Kim,Hiva Mohammadzadeh,Michael W. Mahoney,Yakun Sophia Shao,Kurt Keutzer,Amir Gholami
発行日 2024-04-04 17:45:34+00:00
arxivサイト arxiv_id(pdf)

提供元, 利用サービス

arxiv.jp, DeepL

カテゴリー: cs.LG パーマリンク