2024软件学院创新项目实训--KV Cache 量化

上篇种介绍了模型量化的相关知识，针对我们使用的InternLM大模型，采用KV Cache 量化。

1.量化步骤

KV Cache 量化是将已经生成序列的 KV 变成 Int8，使用过程一共包括三步：

第一步：计算 minmax

主要思路是通过计算给定输入样本在每一层不同位置处计算结果的统计情况。

• 对于 Attention 的 K 和 V：取每个 Head 各自维度在所有Token的最大、最小和绝对值最大值。对每一层来说，上面三组值都是 (num_heads, head_dim) 的矩阵。这里的统计结果将用于本小节的 KV Cache。
• 对于模型每层的输入：取对应维度的最大、最小、均值、绝对值最大和绝对值均值。每一层每个位置的输入都有对应的统计值，它们大多是 (hidden_dim, ) 的一维向量，当然在 FFN 层由于结构是先变宽后恢复，因此恢复的位置维度并不相同。这里的统计结果用于下个小节的模型参数量化，主要用在缩放环节（回顾PPT内容）。

第一步执行命令如下：

# 计算 minmax
lmdeploy lite calibrate \
  --model  /root/share/temp/model_repos/internlm-chat-7b/ \
  --calib_dataset "c4" \
  --calib_samples 128 \
  --calib_seqlen 2048 \
  --work_dir ./quant_output

在这个命令行中，会选择 128 条输入样本，每条样本长度为 2048，数据集选择 C4，输入模型后就会得到上面的各种统计值。值得说明的是，如果显存不足，可以适当调小 samples 的数量或 sample 的长度。

这一步由于默认需要从 Huggingface 下载数据集，国内经常不成功。所以我们导出了需要的数据，大家需要对读取数据集的代码文件做一下替换。共包括两步：

• 第一步：复制 calib_dataloader.py 到安装目录替换该文件：cp /root/share/temp/datasets/c4/calib_dataloader.py /root/.conda/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/lite/utils/
• 第二步：将用到的数据集（c4）复制到下面的目录：cp -r /root/share/temp/datasets/c4/ /root/.cache/huggingface/datasets/

第二步：通过 minmax 获取量化参数

主要就是利用下面这个公式，获取每一层的 K V 中心值（zp）和缩放值（scale）。

zp = (min+max) / 2
scale = (max-min) / 255
quant: q = round( (f-zp) / scale)
dequant: f = q * scale + zp

有这两个值就可以进行量化和解量化操作了。具体来说，就是对历史的 K 和 V 存储 quant 后的值，使用时在 dequant。

第二步的执行命令如下：

# 通过 minmax 获取量化参数
lmdeploy lite kv_qparams \
  --work_dir ./quant_output  \
  --turbomind_dir workspace/triton_models/weights/ \
  --kv_sym False \
  --num_tp 1

在这个命令中，num_tp 的含义前面介绍过，表示 Tensor 的并行数。每一层的中心值和缩放值会存储到 workspace 的参数目录中以便后续使用。kv_sym 为 True 时会使用另一种（对称）量化方法，它用到了第一步存储的绝对值最大值，而不是最大值和最小值。

第三步：修改配置

也就是修改 weights/config.ini 文件，这个我们在《2.6.2 模型配置实践》中已经提到过了（KV int8 开关），只需要把 quant_policy 改为 4 即可。

这一步需要额外说明的是，如果用的是 TurboMind1.0，还需要修改参数 use_context_fmha，将其改为 0。

接下来就可以正常运行前面的各种服务了，只不过咱们现在可是用上了 KV Cache 量化，能更省（运行时）显存了。

2.量化结果

首先，我们先来看一个官方给出的量化结果：

接下来我们对自己的大模型进行量化处理：

量化前后对比

未进行量化-运行：

lmdeploy chat turbomind ./workspace --meta-instruction "你是考研政治题库，内在是InternLM-7B大模型。你将对考研政治单选题，多选题以及综合题做详细、耐心、充分的解答，并给出解析。"

运行结果-占用缓存如下：

可见，占用显存23285MB 。

使用KV cache 0.5：

lmdeploy chat workspace --meta-instruction "你是考研政治题库，内在是InternLM-7B大模型。你将对考研政治单选题，多选题以及综合题做详细、耐心、充分的解答，并给出解析。" --cache-max-entry-count 0.5 

运行结果：

在使用KV cache 0.5量化之后，占用显存20725MB，有明显的下降。

使用KV cache 0.01：

lmdeploy chat workspace --meta-instruction "你是考研政治题库，内在是InternLM-7B大模型。你将对考研政治单选题，多选题以及综合题做详细、耐心、充分的解答，并给出解析。" --cache-max-entry-count 0.01 

运行结果：

显存占用下降更多，仅为16501MB

总结

可见，使用KV Cache量化后，占用显存明显有了下降，关于识别精度将通过后续测试得出，后续还会进行性能的相关测试。