KIVI：用于KV缓存的无需调优的非对称2位量化

《KIVI：用于KV缓存的无需调优的非对称2位量化》的实现

更新

[2024.06.07]：🎉 KIVI 在很大程度上启发了 HuggingFace Transformers 的 KV 缓存量化
[2024.06.06]：（测试版）我们在 develop 分支中对代码库进行了广泛的优化，以减少 KIVI 的延迟。请注意，您需要重新安装我们在 quant 文件夹下的 CUDA 实现。我们将很快发布一篇博客，详细介绍优化内容。
[2024.05.01]：🎉 KIVI 已被 ICML 2024 接收！维也纳见！
[2024.04.12]：我们添加了对 Mistral 模型系列的支持。LongChat-7b-v1.5-32K 和 Mistral-7B-Instruct-v0.2 在 LongBench 的 15 个任务上的性能可以在 long_bench.md 中找到。
[2024.04.05]：我们发布了使用 LM-Eval 复现 CoQA/TruthfulQA/GSM8K 结果的代码。请查看 lmeval 分支的 README。
[2024.04.04]：🔥🔥我们添加了一个新的 5 位数密钥示例，上下文长度为 12k，以展示 2 位 KIVI 在长上下文场景下的性能。
[2024.04.04]：（测试版）我们在预填充阶段为 KIVI 添加了 flash-attention 支持。
[2024.04.03]：我们添加了一个新的 5-shot GSM8K 示例，展示了带有 32 个全精度 token 的 2/4 位 KIVI 的性能。
[2024.02.05]：KIVI 第 2 版发布在 arXiv 上。
[2024.02.03]：KIVI 代码发布。
[2023.12.29]：KIVI 第 1 版发布在 researchgate 上。

概述

KIVI 是一种新的即插即用的 2 位 KV 缓存量化算法，无需任何微调。该算法通过对每个通道的键缓存和每个 token 的值缓存进行 2 位量化来优化内存使用。KIVI 的硬件友好设计允许 Llama-2、Falcon 和 Mistral 等 LLM 在保持可比质量水平的同时，将峰值内存使用减少 2.6 倍。这使得批量大小可以增加到 4 倍，并在实际 LLM 推理工作负载中显著提高吞吐量 2.35 到 3.47 倍，有效解决了速度和内存使用的瓶颈问题。

KIVI 量化方案示意图：每通道键缓存和每 token 值缓存。

KIVI 算法在推理预填充和解码阶段的示意图：

如何使用 KIVI

设置

安装所需包：

conda create -n kivi python=3.10
conda activate kivi
pip install --upgrade pip  # 启用 PEP 660 支持
pip install -e .

然后安装我们的 CUDA 实现：

cd quant && pip install -e .

示例

使用 KIVI 加载模型：（例如，Llama-2-7b）

# 使用 KIVI 的 LLaMA 模型
import torch
import os
from models.llama_kivi import LlamaForCausalLM_KIVI
from transformers import LlamaConfig, AutoTokenizer
config = LlamaConfig.from_pretrained("meta-llama/Llama-2-7b-hf")

config.k_bits = K_BITS # 当前支持 KV 缓存的 2/4 位
config.v_bits = V_BITS # 当前支持 KV 缓存的 2/4 位
config.group_size = GROUP_SIZE
config.residual_length = RESIDUAL_LENGTH # 最近的 fp16 token 数量
CACHE_DIR = PATH_TO_YOUR_SAVE_DIR

model = LlamaForCausalLM_KIVI.from_pretrained(
    pretrained_model_name_or_path='meta-llama/Llama-2-7b-hf',
    config=config,
    cache_dir=CACHE_DIR,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained(
    'meta-llama/Llama-2-7b-hf', 
    use_fast=False, 
    trust_remote_code=True, 
    tokenizer_type='llama')

# 推理
# 例如，model.generate(...)

GSM8K 示例

我们使用 GSM8K 作为示例来展示如何使用 KIVI。您可以查看 example.py：

python example.py

密钥检索示例

使用 KIVI 进行密钥检索。您可以查看 long_context_example.py：

python long_context_example.py

在 LongBench 上评估 KIVI

我们目前支持 Llama 和 Mistral 系列模型。我们最近在 Mistral-7B-Instruct-v0.2 和 Longchat-7b-v1.5-32k 上测试了 KIVI。更多详情请查看 long_bench.md。

bash scripts/long_test.sh {GPU_ID} {K_BITS} {V_BITS} {GROUP_LENGTH} {RESIDUAL_LENGTH} {MODEL_NAME}
python eval_long_bench.py --model {MODEL} # MODEL 是 pred/ 下的目录名。目前支持 Llama 系列模型和 Mistral 模型。

引用

如果您觉得我们的方法有用，请引用我们的论文。

@article{liu2024kivi,
  title={KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache},
  author={Liu, Zirui and Yuan, Jiayi and Jin, Hongye and Zhong, Shaochen and Xu, Zhaozhuo and Braverman, Vladimir and Chen, Beidi and Hu, Xia},
  journal={arXiv preprint arXiv:2402.02750},
  year={2024}
}