Hugging Face KV 缓存量化

Hugging Face 在 Transformers 库中引入了 KV 缓存量化，使大语言模型（LLMs）能够通过降低键值（KV）缓存的内存占用来生成更长的文本序列。该特性允许用户在消费级 GPU 上通过牺牲少量生成速度来显著提升内存效率，从而扩展上下文长度。

KV 缓存在自回归生成中的作用

键值（KV）缓存对于优化自回归模型至关重要，这类模型逐 token 预测文本。要预测新 token，模型需要利用所有先前 token 的信息。如果没有缓存，模型必须在每一步对每个先前 token 重新计算矩阵乘法。KV 缓存充当记忆库，存储自注意力层对已处理 token 的键值对，使模型能够直接检索而非重新计算，从而显著加速文本生成。

然而，随着上下文长度或批量大小的增加，KV 缓存会成为内存瓶颈。以 7B Llama-2 模型、上下文长度为 10,000 token 为例，KV 缓存在 float16 精度下大约需要 5GB 内存，几乎占模型参数所需内存的三分之一。

KV 缓存量化的技术实现

受 KIVI 论文启发，Hugging Face 的实现采用仿射量化将 KV 缓存压缩为低精度格式。

量化方法

该实现对键和值进行逐 token 量化。为减轻在每个生成步骤中进行量化和反量化可能导致的速度下降，Hugging Face 使用了固定大小的残差缓存。该残差缓存以原始精度存储最近的键和值；当缓存达到最大容量时，存储的值会被量化并清空缓存。基准残差长度设为 128，以保持精度。

支持的后端和精度

该特性目前支持以下后端和精度：

  • Quanto：支持 int2 和 int4 精度。
  • HQQ：支持 int2、int4 和 int8 精度。

性能与质量的权衡

对 KV 缓存进行量化需要在内存节省、生成速度和模型质量之间进行权衡。

模型质量与准确性

在 PG-19 数据集上使用 Llama2-7b-chat 模型的测试表明，int4 缓存精度的表现几乎与 fp16 精度相同。但使用 int2 时质量会下降。在 LongBench 基准测试中，Via 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