KV 快取策略

鍵值 (KV) 向量用於計算注意力分數。對於自迴歸模型，KV 分數是*每次*計算的，因為模型一次預測一個 token。每次預測都依賴於之前的 token，這意味著模型每次都會執行相同的計算。

KV *快取*儲存這些計算，以便可以重複使用它們而無需重新計算。高效快取對於最佳化模型效能至關重要，因為它可以減少計算時間並提高響應速度。有關快取工作原理的更詳細解釋，請參閱快取文件。

Transformers 提供了幾個實現不同快取機制的 Cache 類。其中一些 Cache 類經過最佳化以節省記憶體，而另一些則旨在最大限度地提高生成速度。請參閱下表以比較快取型別，並使用它來幫助您為您的用例選擇最佳快取。

快取型別	記憶體效率	支援 torch.compile()	建議初始化	延遲	長上下文生成
動態快取	否	否	否	中	否
靜態快取	否	是	是	高	否
解除安裝快取	是	否	否	低	是
解除安裝靜態快取	否	是	是	高	是
量化快取	是	否	否	低	是
滑動視窗快取	否	是	是	高	否

本指南向您介紹不同的 Cache 類，並展示如何將它們用於生成。

預設快取

DynamicCache 是大多數模型的預設快取類。它允許快取大小動態增長，以便在生成過程中儲存越來越多的鍵和值。

透過在 generate() 中配置 use_cache=False 來停用快取。

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).to("cuda:0")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

快取類也可以在呼叫並將其傳遞給模型的 past_key_values 引數之前進行初始化。此快取初始化策略僅推薦用於某些快取型別。

在大多數其他情況下，在 cache_implementation 引數中定義快取策略更容易。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

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).to("cuda:0")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache()
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

記憶體高效快取

KV 快取會佔用大量記憶體，並可能成為長上下文生成的瓶頸。記憶體高效快取側重於犧牲速度以減少記憶體使用。這對於大型語言模型 (LLM) 以及硬體記憶體受限的情況尤其重要。

解除安裝快取

OffloadedCache 透過將大部分模型層的 KV 快取移動到 CPU 來節省 GPU 記憶體。在模型對層進行 forward 迭代期間，只有當前層快取保留在 GPU 上。OffloadedCache 非同步預取下一層快取並將上一層快取傳送回 CPU。

此快取策略始終生成與 DynamicCache 相同的結果，並可作為即插即用替代品或備用方案。如果您有 GPU 並且出現記憶體不足 (OOM) 錯誤，您可能希望使用 OffloadedCache。

根據您的模型和生成選擇（上下文大小、生成的 token 數量、束的數量等），您可能會注意到與 DynamicCache 相比，生成吞吐量略有下降。

透過在 GenerationConfig 或 generate() 中配置 cache_implementation="offloaded" 來啟用 OffloadedCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

以下示例展示瞭如何在記憶體不足時回退到 OffloadedCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def resilient_generate(model, *args, **kwargs):
    oom = False
    try:
        return model.generate(*args, **kwargs)
    except torch.cuda.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch.cuda.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_beam_groups": 40, "num_return_sequences": 40, "diversity_penalty": 1.0, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

量化快取

QuantizedCache 透過將 KV 值量化為較低精度來減少記憶體需求。QuantizedCache 目前支援兩種量化後端。

HQQQuantizedCache 支援 int2、int4 和 int8 資料型別。
QuantoQuantizedCache 支援 int2 和 int4 資料型別。這是預設的量化後端。

如果上下文長度較短且有足夠的 GPU 記憶體可用於生成而無需啟用快取量化，則量化快取可能會損害延遲。嘗試在記憶體效率和延遲之間找到平衡點。

透過在 GenerationConfig 中配置 cache_implementation="quantized" 並指示 QuantizedCacheConfig 中的量化後端來啟用 QuantizedCache。任何其他量化相關引數也應作為字典或 QuantizedCacheConfig 例項傳遞。您應該使用這些額外引數的預設值，除非您記憶體不足。在這種情況下，請考慮減小殘差長度。

HQQQuantizedCache

Quanto

速度最佳化快取

預設的 DynamicCache 阻止您利用即時 (JIT) 最佳化，因為快取大小不是固定的。JIT 最佳化使您能夠以犧牲記憶體使用為代價來最大限度地提高延遲。以下所有快取型別都與 JIT 最佳化相容，例如 torch.compile 以加速生成。

靜態快取

StaticCache 為 KV 對預分配特定的最大快取大小。您可以在不修改快取的情況下生成最多達到最大快取大小的內容。

透過在 generate() 中配置 cache_implementation="static" 來啟用 StaticCache。

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="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

解除安裝靜態快取

OffloadedStaticCache 與 OffloadedCache 非常相似，只不過快取大小設定為最大快取大小。否則，OffloadedStaticCache 只在 GPU 上保留當前層快取，其餘的則移動到 CPU。

透過在 generate() 中配置 cache_implementation="offloaded_static" 來啟用 OffloadedStaticCache。

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="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="offloaded_static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

快取解除安裝需要 CUDA GPU。

滑動視窗快取

SlidingWindowCache 在之前的 kv 對上實現了一個滑動視窗，並且只保留最後的 sliding_window tokens。這種快取型別旨在僅與支援*滑動視窗注意力*的模型一起使用，例如 Mistral。較舊的 kv 狀態被丟棄並替換為新的 kv 狀態。

透過在 generate() 中配置 cache_implementation="sliding_window" 來啟用 SlidingWindowCache。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", torch_dtype=torch.float16).to("cuda:0")
inputs = tokenizer("Yesterday I was on a rock concert and.", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=30, cache_implementation="sliding_window")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]

模型快取

某些模型型別，例如編碼器-解碼器模型或 Gemma2 和 Mamba，具有專用的快取類。

編碼器-解碼器快取

EncoderDecoderCache 專為編碼器-解碼器模型設計。它管理自注意力快取和交叉注意力快取，以確保儲存和檢索以前的 kv 對。可以分別為編碼器和解碼器設定不同的快取型別。

此快取型別不需要任何設定。它可以在呼叫 generate() 或模型的 forward 方法時使用。

EncoderDecoderCache 目前僅支援 Whisper。

模型專用快取

有些模型儲存過去 kv 對或狀態的方式是獨有的，與任何其他快取類不相容。

Gemma2 需要 HybridCache，它在底層結合使用了 SlidingWindowCache 用於滑動視窗注意力，以及 StaticCache 用於全域性注意力。

Mamba 需要 MambaCache，因為該模型沒有注意力機制或 kv 狀態。

迭代生成

快取還可以在模型（聊天機器人）之間進行來回互動的迭代生成設定中工作。像常規生成一樣，帶有快取的迭代生成允許模型有效地處理正在進行的對話，而無需在每個步驟重新計算整個上下文。

對於帶快取的迭代生成，首先初始化一個空的快取類，然後可以輸入新的提示。使用聊天模板跟蹤對話歷史記錄。

以下示例演示了 Llama-2-7b-chat-hf。如果您正在使用不同的聊天風格模型，apply_chat_template() 可能會以不同的方式處理訊息。根據 Jinja 模板的編寫方式，它可能會剪下掉重要的 token。

例如，一些模型在推理過程中使用特殊的 <think> ... </think> tokens。這些 tokens 在重新編碼時可能會丟失，導致索引問題。您可能需要手動刪除或調整完成中的額外 tokens 以保持穩定。

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM
from transformers.cache_utils import (
    DynamicCache,
    StaticCache,
    SlidingWindowCache,
    QuantoQuantizedCache,
    QuantizedCacheConfig,
)

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache()

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

預填充快取

在某些情況下，您可能希望用特定字首提示的 kv 對填充 Cache，並重復使用它來生成不同的序列。

以下示例初始化一個 StaticCache，然後快取一個初始提示。現在，您可以從預填充的提示生成多個序列。

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_batch_size=1, max_cache_len=1024, device="cuda", dtype=torch.bfloat16)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)

< > 在 GitHub 上更新