通过键值缓存量化解锁更长文本生成

在 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)。高效神经网络推理的量化方法综述。