LLM 进 KV 缓存的量化

一、推理需要的内存

推理需要的内存主要有三部分：加载模型、激活和KV-cache

量化通过降低大型语言模型 (LLM) 的参数精度（例如从 16 位到 4 位）来减小其大小。

然而，量化模型只会减少模型参数的内存消耗。虽然可以在较小的 GPU 上加载 LLM，但推理过程仍然需要额外的内存来主要存储：

• 激活，即在前向传递过程中创建的张量

• KV缓存

  
  

正常情况下，推理中的激活由于是一层层串行传递，因此通常现存占用不是很大。

KV-Cache方面，随着上下文变长和模型变深，KV 缓存的大小会迅速增长。对于处理长上下文如 RAG 系统来说，情况更为突出。

二、Transformer 模型的 KV 缓存是什么

在 Transformer 模型中，文本生成是逐个标记进行的，其中每个新预测都取决于所有先前生成的标记的上下文。

这种顺序依赖性会使过程变慢，因为预测下一个 token 需要重新处理迄今为止生成的整个序列。例如，要预测第 100 个 token，模型需要整合前 99 个 token 的信息，这涉及对它们的表示进行复杂的矩阵乘法。预测第 101 个 token 需要再次对前 99 个 token 执行相同的计算，再加上对第 100 个 token 执行的额外计算。

键值 (KV) 缓存通过存储这些计算的结果来优化此过程，允许模型将它们重新用于后续标记而无需重新计算。这意味着，为了生成第 101 个标记，不必重新计算前 99 个标记的信息，而是可以从 KV 缓存中检索它们，然后只计算第 100 个标记的缺失信息。

具体来说，KV 缓存存储了从先前处理的 token 的自注意力层派生的键值对。在 Transformer 架构中，自注意力层通过将查询与键相乘并产生值向量的加权和作为输出来生成注意力分数。通过将这些键和值存储在缓存中，模型可以有效地检索它们以加快生成过程。

在推理过程中利用 KV 缓存现在已成为标准做法。生成的序列越长，加速效果就越显著。然而，这也意味着 KV 缓存正在增长。必须将所有这些张量存储在 GPU 上，以便快速检索和利用它们。对于深度模型、长序列和大批量，KV 很容易占用数十 GB 的 GPU RAM。

二、估算 KV 缓存的内存消耗

KV 缓存通常以 16 位存储张量，使用 float16 或 bfloat16 数据类型。

对于一个 token，KV 缓存会为每一层和注意力头存储一对张量（键值对）。这些张量的大小由注意力头的维度决定。这对张量的总内存消耗（以字节为单位）为：

number of layers * number of KV attention heads * dimension of the attention head * (bit width / 8) * 2

公式中：

• hidden dimension是隐藏层的维度。

• dimension of the attention head= hidden dimension/number of attention heads. 

• 最后一个“2”是因为有两个张量，即键和值。位宽大多数情况下为 16。由于 8 位为 1 字节，将位宽除以 8，这样 KV 缓存中每个 16 位参数就有 2 个字节。

如果我采用 Llama 3 8B，则该等式变为：

32 * 8 * 128 * 2 * 2 = 131,072

注意：Llama 3 8B 有 32 个注意力头。但是，得益于分组查询注意力 (GQA) ，只有 8 个注意力头用于键和值。

一个 token 的 KV 缓存占用 131,072 字节，即 0.1 MB。它看起来很小，但对于许多不同类型的应用程序来说，LLM 需要生成数千个 token。例如，如果想利用 Llama 3 8B 的完整上下文大小（即 8192），KV 缓存将存储 8191 个 token 的 KV 张量。这是 1.1 GB。对于具有 24 GB RAM 的消费级 GPU，KV 缓存将占用其总内存的 4.5%。

对于较大的模型，KV 缓存增长得更快。例如，对于具有 80 层的 Llama 3 70B，公式变为：

80 * 8 * 128 * 2 * 2 = 327,680

对于 8191 个令牌，Llama 3 70B 的 KV 缓存将占用 2.7 GB。

还要注意，这是一个seq的内存消耗。如果进行批量解码，必须将此值乘以批量大小。

例如，使用 Llama 3 8B 的批处理大小 32 需要 35.2 GB 的 GPU RAM。这无法通过一个消费级 GPU 实现。

三、Llama 3 的 KV 缓存量化

量化会降低参数的精度，例如从 16 位降低到 4 位。这意味着 16 位张量的总大小可以除以 4，或者使用 2 位量化除以 8。

如果拿上面最后一个例子来说明内存消耗，理论上 4 位量化可以将 KV 缓存的大小从 35.2 GB 减少到 8.8 GB，2 位量化可以减少到 4.4 GB。但实际操作中，量化算法的效果并不好。内存消耗的减少取决于量化方法及其超参数，尤其是组大小。

使用HQQ 量化，可以预期 4 位量化的内存消耗将减少 3 倍至 2 倍。使用块量化等更简单的方法，可以预期减少接近 3 倍。

Hugging Face Transformers 支持 KV 缓存量化。

目前，支持它的版本不能直接通过 pip 使用。必须从源代码安装 Transformers：

pip install git+https://github.com/huggingface/transformers

它支持 HQQ 量化和 Quanto（块式）。HQQ 比 Quanto 更准确，但效率不如 Quanto。据 Hugging Face 报告，HQQ 也比 Quanto 慢得多：

如果应用程序不太关心解码速度，建议使用 HQQ。否则，请使用 Quanto。

如果想使用 HQQ，请运行：

pip install HQQ

pip install quanto==0.1.0

Hugging Face 目前建议使用 0.1.0 版本。

下面是用来量化Llama 3 8B Instruct的KV缓存的代码，使用bnb动态量化。

import torchfrom transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, BitsAndBytesConfig
set_seed(1234)  # For reproducibility
prompt = "The best tomato sauce is"
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
bnb_config = BitsAndBytesConfig(        load_in_4bit=True,        bnb_4bit_quant_type="nf4",        bnb_4bit_compute_dtype=torch.float16,        bnb_4bit_use_double_quant=True,)
tokenizer = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, attn_implementation="flash_attention_2", torch_dtype=torch.float16, device_map="cuda:0")
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)outputs = model.generate(**inputs, max_new_tokens=150, cache_implementation="quantized", cache_config={"backend": "HQQ", "nbits": 4, "q_group_size": 128, "residual_length": 64, "device":model.device})result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)

其中outputs 部分为介绍KV缓存量化的部分。cache_implementation表示在推理过程中会量化KV缓存。cache_config表示量化用到的所有超参数：

1. backend

• 描述：选择量化算法的类型。

• 选项：

• HQQ：适用于对精度要求较高的场景，但推理吞吐量较低。

• Quanto：适用于对推理速度要求较高的场景，但精度可能会有所损失。

• 建议：如果更关注模型的精度，选择 HQQ；如果更关注推理速度，选择 Quanto。
  

2. nbits

• 描述：量化的精度，即每个参数使用的位数。

• 选项：

• 4：表示4位量化。这是一个常见的选择，可以显著减少内存消耗，同时保持较好的模型精度。

• 2：表示2位量化。尽管可以进一步减少内存消耗，但可能会严重降低模型的准确性。

• 建议：一般情况下选择 4 位量化。只有在内存极其紧张的情况下才考虑使用 2 位量化。
  

3. q_group_size

• 描述：量化时的组大小，决定了每组参数的量化粒度。

• 默认值： 128

• 调整：

• 减小值：可以提高量化的精度，但内存消耗的减少效果会降低。

• 增大值：可以提高内存节省效果，但可能会降低量化精度。

• 建议：通常保持默认值 128。只有在需要更高精度时才考虑减小该值。
  

4. residual_length

• 描述：表示有多少个令牌的KV缓存不进行量化。这对于保持模型的准确性至关重要。

• 默认值： 64

• 调整：

• 增大值：可以提高模型的准确性，但会增加内存消耗。

• 减小值：可以减少内存消耗，但可能会降低模型的准确性。

• 建议：设置为 64 是一个较好的平衡。如果模型的准确性不够，可以尝试增大该值；如果内存消耗仍然过高，可以尝试减小该值。
  

5. device

• 描述：指定KV缓存量化运行的设备，通常应该与模型加载的设备相同。

• 建议：明确设置为与模型加载设备一致。例如，如果模型加载在GPU上（如 cuda:0），那么 device 应该也设置为 cuda:0。如果未设置，默认情况下量化可能会在CPU上进行，这可能会导致性能问题。

如果所有这些超参数都设置正确，则使用和不使用 KV 缓存量化生成的输出不会有太大差异。

至于内存消耗的差异，只有在处理大批量和非常长的序列时才会注意到显著的差异（同时处理来自多个用户的查询时，批处理可以让GPU并行处理，从而提高效率）。

KV 缓存量化是进一步减少 LLM 推理内存消耗的另一种方法。它可以将 KV 缓存的大小减少 2 倍或 3 倍，具体取决于后端和量化超参数，从而可能节省数十GB。

然而，这种量化并非没有缺点。它会降低解码速度，并可能显著降低 LLM 的准确性。通过调整量化超参数可以提高速度和准确性。

参考：https://kaitchup.substack.com/p/kv-cache-quantization-for-memory