透過鍵值快取量化解鎖更長文字生成

在 Hugging Face，我們很高興與您分享一項新功能，它將把您的語言模型提升到新的水平：鍵值快取量化 (KV Cache Quantization)。

總結：KV 快取量化在對質量影響最小的情況下，降低了大型語言模型（LLM）長上下文文字生成的記憶體使用，在記憶體效率和生成速度之間提供了可定製的權衡。

您是否曾嘗試使用您的語言模型生成一段冗長的文字，卻因惱人的記憶體限制而碰壁？隨著語言模型規模和能力的不斷增長，支援更長的生成可能會真正佔用記憶體。這是一種常見的挫折，尤其是在您處理有限資源時。這就是 KV 快取量化大顯身手的時候。

那麼，KV 快取量化到底是什麼？如果您不熟悉這個術語，請不要擔心！讓我們將其分解為兩部分：KV 快取和量化。

鍵值快取（或 KV 快取）是最佳化自迴歸模型生成所必需的，在自迴歸模型中，模型逐個預測文字標記。這個過程可能很慢，因為模型一次只能生成一個標記，並且每個新的預測都依賴於先前的上下文。這意味著，要在生成中預測第 1000 個標記，您需要來自前 999 個標記的資訊，這些資訊以這些標記表示的某些矩陣乘法的形式出現。但要預測第 1001 個標記，您還需要來自前 999 個標記的相同資訊，以及來自第 1000 個標記的額外資訊。這就是鍵值快取用於透過儲存先前的計算以在後續標記中重用，從而最佳化順序生成過程的地方，這樣它們就不需要再次計算。

更具體地說，鍵值快取充當自迴歸生成模型的記憶體庫，模型在其中儲存來自先前處理的標記的自注意力層派生的鍵值對。在 Transformer 架構中，自注意力層透過將查詢與鍵相乘來計算注意力分數，生成值向量的加權和作為輸出。透過儲存此資訊，模型可以避免冗餘計算，而是從快取中檢索先前標記的鍵和值。有關此概念的視覺解釋，請檢視下圖中鍵值快取的功能。當計算第 `K+1` 個標記的注意力分數時，我們不需要重新計算所有先前的鍵和值，而是從快取中獲取並連線到當前向量。這通常會帶來更快、更高效的文字生成。

接著說第二個術語，量化只是一個花哨的詞，指為了節省記憶體而降低數值精度。在量化過程中，每個數值都被四捨五入或截斷以適應降低精度的格式，這可能會導致資訊丟失。然而，仔細選擇量化引數和技術可以最大限度地減少這種損失，同時仍能實現令人滿意的效能。有不同的量化方法，如果您想了解更多資訊，請務必檢視我們之前的部落格文章，以深入瞭解量化世界。

即使 KV 快取加速了自迴歸生成，但對於長上下文或高批處理大小，它可能會成為記憶體瓶頸。讓我們估算一下，對於一個 7B Llama-2 模型，儲存序列長度為 10000 個標記的 KV 快取需要多少記憶體。儲存一個標記的 KV 快取所需的記憶體大約是 `2 * 2 * num_layers * num_key_value_heads * head_dim`，其中第一個 `2` 表示鍵和值，第二個 `2` 是我們需要的位元組數（假設模型以 `float16` 載入）。所以如果我們有一個長度為 10000 個標記的上下文，我們將需要

2 * 2 * 32 * 32 * 128 * 10000 ≈ 5GB

的記憶體僅用於儲存之前的鍵值快取，這幾乎是儲存半精度模型引數所需記憶體的三分之一。

因此，透過將 KV 快取壓縮成更緊湊的形式，我們可以節省大量記憶體並在消費級 GPU 上執行更長的上下文生成。在我們的實驗中，我們能夠透過將 KV 快取量化為較低精度格式，顯著減少記憶體佔用，而不會犧牲太多質量。藉助這項新的量化功能，我們現在可以支援更長的生成，而不會耗盡記憶體，這意味著您可以在不擔心遇到記憶體限制的情況下擴充套件模型的上下文長度。

實現細節

Transformers 中的鍵值快取量化主要受到 KIVI: A Tuning-Free Asymmetric 2bit Quantization for kv Cache 論文的啟發。該論文引入了一種用於大型語言模型的 2 位非對稱量化，且不降低質量。KIVI 對鍵快取進行每通道量化，對值快取進行每標記量化，因為他們表明對於 LLM，鍵在某些通道中具有更高的異常值，而值不顯示這種模式。因此，當鍵進行每通道量化而值進行每標記量化時，量化精度和原始精度之間的相對誤差要小得多。

在我們整合到 Transformers 中的方法中，鍵和值都進行每標記量化。每標記量化的主要瓶頸是每次新增新標記時（即每個生成步驟）都需要對鍵和值進行量化和反量化。這可能會導致生成速度變慢。為了克服這個問題，我們決定保留一個固定大小的殘差快取，以原始精度儲存鍵和值。當殘差快取達到其最大容量時，儲存的鍵和值將被量化並丟棄快取內容。這個小技巧還可以保持準確性，因為最新的一些鍵和值始終以其原始精度儲存。主要考慮因素是設定殘差快取長度時的記憶體效率權衡。雖然殘差快取以其原始精度儲存鍵和值，但這可能會導致總體記憶體使用量增加。我們發現使用 128 的殘差長度作為基線效果很好。

因此，給定形狀為 `batch size, num of heads, num of tokens, head dim` 的鍵或值，我們將其分組為 `num of groups, group size` 並執行仿射量化如下：

X_Q = round(X / S) - Z

其中，

• X_Q 是量化張量
• S 是比例因子，計算公式為 `(maxX - minX) / (max_val_for_precision - min_val_for_precision)`
• Z 是零點，計算公式為 `round(-minX / S)`

目前，kv 量化適用於 `quanto` 後端（支援 `int2` 和 `int4` 精度）和 `HQQ` 後端（支援 `int2`、`int4` 和 `int8` 精度）。有關 `quanto` 的更多資訊，請參閱之前的部落格文章。雖然我們目前不支援更多的量化後端，但我們歡迎社群貢獻，以幫助整合它們。具體來說，不需要校準資料並可以動態計算低位張量的量化方法可以很容易地整合。此外，您可以在配置中指定最常見的量化引數，從而可以自由調整量化過程，例如根據您的用例決定是執行每通道量化還是每標記量化。

對比 fp16 和量化快取的效能

我們知道圖表比文字更具說服力，所以我們準備了一些對比圖，讓您一目瞭然地瞭解量化與 FP16 精度相比的表現。這些圖表能讓您快速瞭解，當我們調整 KV 快取的精度設定時，模型生成的質量表現如何。我們使用以下量化引數在 `PG-19` 資料集上計算了 Llama2-7b-chat 模型的困惑度：`nbits=4, group_size=64, resildual_length=128, per_token=True`

我們可以看到，對於兩個後端，`int4` 快取的效能幾乎與原始 `fp16` 精度相同，而使用 `int2` 時質量會下降。重現結果的指令碼可在此處獲取。

與 KIVI 論文的結果進行比較，在 LongBench 基準測試中計算效能時，也得出相同的結論。`Int4 quanto` 精度與 `fp16` 相當，甚至在下表中的所有資料集中都略優於 `fp16`（越高越好）。

現在，我們來談談記憶體節省和速度之間的權衡。當我們對模型中的 KV 快取進行量化時，會使其記憶體佔用更小，但有時這會以犧牲一點生成速度為代價。雖然將快取量化為 `int4` 可以節省大約 2.5 倍的記憶體，但隨著批處理大小的增加，生成速度開始下降。人們必須決定是否值得為了記憶體效率的顯著提高而使用量化 KV 快取並可能犧牲一點速度。這完全是為了找到最適合您特定用例和優先順序的方法。

以下是 KV 快取原始精度和量化格式的效能指標。獲取以下資料的指令碼可在此處獲取。

想知道當我們加入權重量化時會發生什麼嗎？當然，結合這些技術可以進一步縮小模型的記憶體佔用，但也有一個缺點——它可能會進一步降低速度。事實上，我們的實驗表明，權重量化與 KV 快取量化結合使用會導致速度降低三倍。但我們一直在努力尋找使這種組合無縫工作的方法。雖然我們目前在 `quanto` 庫中沒有最佳化的核心，但我們歡迎社群貢獻，以幫助提高計算效率。我們的目標是確保您的模型執行流暢，同時保持高延遲和準確性。

還值得注意的是，輸入提示的初始處理（又稱預填充階段）仍然需要一次性為整個輸入計算完整的鍵值矩陣，這對於長上下文來說可能是另一個記憶體瓶頸。這就是為什麼生成第一個標記的延遲往往高於後續標記的原因。還有其他不同的策略可以透過最佳化注意力計算階段來減少預填充階段的記憶體負擔，例如區域性視窗注意力（Local Windowed Attention）或 Flash-Attention。如果您在預填充階段記憶體不足，您可以在 🤗 Transformers 中使用 `FlashAttention` 以及 KV 快取量化，以進一步減少長輸入提示的記憶體使用。有關更多資訊，請參閱文件。

如果您對上下文能容納多少標記感興趣，在記憶體使用達到極限的情況下，啟用 Flash Attention 的量化 KV 快取可以在 80GB A100 中支援多達 128k 個標記。對於半精度快取，最大容量為 40k 個標記。

如何在 🤗 Transformers 中使用量化 KV 快取？

要在 🤗 Transformers 中使用 KV 快取量化，我們首先必須透過執行 `pip install quanto` 來安裝外部依賴項。要啟用 KV 快取上的量化，我們必須傳入 `cache_implementation="quantized"` 並在字典格式的快取配置中指示量化引數。這就是我們開始使用 KV 快取量化所需的全部內容。此外，由於 quanto 與裝置無關，無論您是在 CPU/GPU/MPS（Apple Silicon）上，您都可以量化和執行您的模型。

您可以在此處的 Colab 筆記本中找到簡短的使用示例。

>>> import torch
>>> from transformers import AutoTokenizer, AutoModelForCausalLM

>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
>>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.float16, device_map="cuda:0")
>>> inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

>>> out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "quanto", "nbits": 4})
>>> print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

>>> out = model.generate(**inputs, do_sample=False, max_new_tokens=20)
>>> print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. I like to listen to it when I'm feeling

結論

有許多不同的方法可以減少鍵值快取的記憶體使用，包括 MultiQueryAttention、GroupedQueryAttention 或最近的 kv 快取檢索方法。雖然其中一些方法受限於模型架構選擇，但另一些方法可以在訓練後應用。量化就是其中一種訓練後最佳化技術，我們可以從這篇簡短的部落格文章中得出以下結論：

1. 記憶體與速度的權衡：透過將 KV 快取量化為較低精度格式，記憶體使用顯著減少，從而可以在不遇到記憶體限制的情況下生成更長的文字。但使用者必須決定犧牲一點生成速度是否適合他們的用例。

2. 保持準確性：儘管精度降低，但 `int4` 中的 KV 快取量化在令人滿意的程度上保持了模型準確性，確保生成的文字保持上下文相關性和連貫性。

3. 靈活性：使用者可以根據自己的特定要求靈活選擇不同的精度格式，從而可以進行定製以適應不同的用例和優先順序。

4. 進一步最佳化的潛力：雖然 KV 快取量化本身提供了顯著的優勢，但它也可以與其他最佳化技術（例如權重量化）結合使用，以進一步提高記憶體效率和計算速度。

致謝

特別感謝 Younes 和 Marc 在量化技術方面提供的幫助和建議。他們的專業知識為該功能的開發做出了巨大貢獻。

此外，我還要感謝 Joao 提供的寶貴支援。

附加資源

1. Zirui Liu、Jiayi Yuan、Hongye Jin、Shaochen Zhong、Zhaozhuo Xu、Braverman, V.、Beidi Chen 和 Hu, X. (2023)。KIVI : 具有流式非對稱量化的即插即用 2 位 KV 快取量化。
2. Databricks 關於 LLM 推理效能工程：最佳實踐的部落格文章
3. Coleman Hooper、Sehoon Kim、Hiva Mohammadzadeh、Michael W. Mahoney、Yakun Sophia Shao、Kurt Keutzer 和 Amir Gholami。（2024）。KVQuant：實現 1000 萬上下文長度 LLM 推理與 KV 快取量化。
4. T. Dettmers、M. Lewis、Y. Belkada 和 L. Zettlemoyer，(2022)。LLM.int8()：大規模 Transformer 的 8 位矩陣乘法。
5. A. Gholami、S. Kim、Z. Dong、Z. Yao、M. W. Mahoney 和 K. Keutzer，(2021)。高效神經網路推理的量化方法綜述。