Hugging Face KV 快取量化

Hugging Face 在 Transformers 函式庫中引入了 KV 快取量化，使大型語言模型（LLM）能透過減少鍵值（KV）快取的記憶體佔用，生成更長的文字序列。此功能允許使用者在消費級 GPU 上透過犧牲少量生成速度，換取顯著的記憶體效能，從而擴展上下文長度。

KV 快取在自回歸生成中的角色

鍵值（KV）快取對於優化自回歸模型至關重要，這類模型逐字元（token）預測文字。要預測新 token，模型需要所有先前 token 的資訊。若沒有快取，模型必須在每一步重新計算所有先前 token 的矩陣乘法。KV 快取充當記憶庫，儲存自注意力層對先前處理過的 token 所產生的鍵值對，讓模型直接檢索而非重新計算，從而大幅加速文字生成。

然而，隨著上下文長度或批次大小的增加，KV 快取會成為記憶體瓶頸。以 7B Llama-2 模型、10,000 token 的上下文長度為例，KV 快取在 float16 精度下大約需要 5GB 記憶體，佔模型參數所需記憶體的近三分之一。

KV 快取量化的技術實作

受 KIVI 論文啟發，Hugging Face 的實作使用仿射量化（affine quantization）將 KV 快取壓縮為較低精度的格式。

量化方法

此實作對鍵和值皆執行逐 token 量化。為減少在每一次生成步驟中量化與反量化可能帶來的速度下降，Hugging Face 採用了固定大小的殘差快取。該殘差快取以原始精度儲存最近的鍵和值；當快取達到最大容量時，儲存的值會被量化，快取隨即被清空。以 128 的殘差長度作為基線，以維持精度。

支援的後端與精度

此功能目前支援以下後端與精度：

  • Quanto：支援 int2 與 int4 精度。
  • HQQ：支援 int2、int4 與 int8 精度。

效能與品質的取捨

對 KV 快取進行量化需要在記憶體節省、生成速度與模型品質之間取得平衡。

模型品質與準確度

在 PG-19 資料集上使用 Llama2-7b-chat 模型的測試顯示，int4 快取精度的表現與 fp16 精度幾乎相同。然而，使用 int2 時品質會下降。在 LongBench 基準測試中，透過 Quanto 後端的 int4 精度在多個資料集（包括 TREC、SAMSum 與 TriviaQA）上與 fp16 相當，甚至略有超越。

記憶體效率與延遲

將快取量化為 int4 可帶來約 2.5 倍的記憶體節省。雖然減輕了記憶體壓力，但可能導致生成速度下降，尤其在批次大小增大時更為明顯。此外，將 KV 快取量化與權重量化結合使用，生成速度可能下降至原來的三分之一。

上下文長度容量

在 80GB A100 GPU 上結合 Flash Attention 時，KV 快取量化使模型能支援最高 128k token，相較於使用半精度快取時最高僅 40k token。

與 Transformers 的整合

KV 快取量化與裝置無關，可在 CPU、GPU 與 MPS（Apple Silicon）上運作。要在 🤗 Transformers 中使用此功能，使用者必須安裝 quanto 套件，並在 generate 呼叫時指定 cache_implementation="quantized" 參數以及 cache_config 字典。

# Example usage
out = model.generate(
    **inputs, 
    do_sample=False, 
    max_new_tokens=20, 
    cache_implementation="quantized", 
    cache_config={"backend": "quanto", "nbits": 4}
)

Sources