QUIK：LLM上的一种4-bit推理方法

论文概述

 论文提出了一种叫做 QUIK 的混合量化策略，QUIK 把大部分的激活和权重量化为 4-bit，通过校准数据集进行激活进行统计，把一些 outliers 激活和对应的权重用 FP16 表示。该机制的特点是计算高效，同时能保持很好的精度。通过 QUIK，端对端吞吐量能达到 FP16 的 3.4 倍。论文研究了 OPT 模型、LLaMA-2 模型和 Falcon 模型，在推理时，除了量化，还进行了 2:4 稀疏的尝试。

   同时，论文还指出对 LLaMA 类模型，如 LLaMA-2 模型，down-project 对量化更加敏感，对激活权重不进行 4-bit 量化，而改用 8-bit 量化，对保持模型的精度非常重要。

   下图是 LLaMA-2 模型使用 QUIK 量化后的效果，对 70B 模型，精度降低在 0.5 PPL 以内，70% 的线性层使用 INT4 进行计算，带来了端对端 3.4 倍的加速。

论文细节

动机

论文关注于计算受限的场景。

NVIDIA GPU 本来就支持 INT4，INT8 计算，下图是不同矩阵大小下不同数据类型矩阵乘法的吞吐量。从图可以看出，INT8 的吞吐量是 FP16 的 2 倍，而 INT4 是 INT8 吞吐量的 2 倍。要利用这种硬件特性，激活和权重必须量化为相同的数据类型。但是目前，把权重和激活都量化成 INT4，会带来较大的精度损失，基于此，QUIK 致力于提出一种能够把权重和激活都压缩成 INT4 数据类型的高精度的 PTQ 量化方法。

QUIK量化

1. QUIK量化工作流

线性变换可以表示为：$X W^T$

QUIK 的工作流如下图，由于 outliers 总是出现在激活的固定的列，step1：可以通过校准数据集统计激活，step2：把激活的 outliers 放到最后，对应调整权重矩阵的列的分布。step3：对权重$W$进行 GPTQ INT4 量化，GPTQ 更新权重的范围是整个$W$，但是量化的范围是 outliers 之前的列，groupsize = -1。这样量化误差就补偿到了 outliers 里。

后面的步骤下图就没有显示了。后面首先对激活的正常列进行 per token 量化。量化完成后，激活权重的正常列进行 INT4 矩阵乘，再反量化为 FP16 值，outliers 进行 FP16 的乘法。最后把正常部分的结果与 outliers 的结果相加作为最后的结果矩阵。

总体而言，激活是非对称 per token 动态量化，权重是对称静态 GPTQ 量化。

通过这样的工作流，GPTQ 量化把量化误差累积到那些用 FP16 计算的列，并且避免了 INT4 量化那些量化困难的列。

2. 权重裁剪

在进行 GPTQ 量化，计算量化参数前，使用 linear search 的方式来确定 weight clipping 的 threshold，这样可改善 PPL。

3. 基于敏感度的部分量化

通过列的无穷范数（元素绝对值的最大值）来选择 outliers。对 outliers 保持 FP16 不量化。outliers=256。

上面的做法对 OPT 模型来说，已经能保持足够的精度了，但是对 LLaMA2-70B 模型模型是不够的。这主要由于在 FeedForward 层，有 3 个线性转换，其间包含了按元素相乘和 SiLU 激活函数。下图是不同层激活的方差分析，可见 Down_proj 层对量化更敏感，为了恢复量化精度，Down_proj 使用 INT8 量化。

4. 高效推理实现

推理实现的示意图见下图，该图展示了 QUIK-4B 前向推理时线性层的实现方式。

具体算法见下图。Line 6 的矩阵乘法是利用 CUTLASS 进行的，能高效的利用硬件的 INT8/INT4 tensor-cores 进行快速低精度计算，结果用 INT32 进行累加。

实验结果

1. 4-bit PPL 精度 与 zero-shot 任务  

下表是 4 个 OPT 模型使用不同的量化方法在 WikiText2 任务上的效果。在 QUIK 中，每一层线性层都使用了统一的 256 个异常点，除了异常点，其他权重和对应激活都量化成 INT4。可以看出，QUIK 在精度上取得了比之前的量化方法都好的效果。

下表是把 QUIK 用在了 LLaMA-2 和 Falcon 模型上，同样是 WikiText2 任务。和 OPT 模型一样， 每个线性层都使用了 256 个异常点。不同在于，对 LLaMA-2 的 down-projection 层和 Falcon 的 FC2 层使用了 INT8 量化，并且按这些层的输入特征数成比例的增大了 outliers 数量。可以看出，QUIK 也取得了较好的效果。（ QUIK-4B 表示 4-bit QUIK 量化）

下表是 QUIK 在一些 zero-shot 任务上的效果。

2. 8-bit 量化 （QUIK-8B）

下表对比了 QUIK-8B 和 SmoothQuant 在 OPT， LLaMA-2， Falcon上的效果。  虽然两种方法都接近无损了，但是 QUIK 的效果更好点。

性能分析

1.  峰值显存使用

对 OPT-66B 模型，QUIK-8B 和 QUIK-4B 的模型分别降低了 47%（理想情况是 50%）和 74%（理想情况是 75%）的峰值显存使用。

对 LLaMA2-70B 模型，QUIK-8B 和 QUIK-4B 的模型分别降低了 32% 和 67% 的峰值显存使用，因为在 LLaMA2-70B 中，down-project 使用了 8-bits 并使用了额外的 outliers。具体见下表。

2. 层加速

理想的层加速见**动机**章节，没有考虑量化/反量化开销。

QUIK 的实际层加速见下图（QUIK-4B 每层使用 256 个 outliers, 在单个 RTX3090 上进行测试）。可见 QUIK-4B 在大的层上能达到高于 4 倍的层加速，在小的层上能有 2 倍多的层加速。

3. 端到端加速

QUIK-4B 使用 256 个 outliers 除了down_project 层（8-bit 量化，使用超过 600 个 outliers）。

从下图可见，OPT-66B 模型除了能减少接近 4 倍的内存，吞吐量是 FP16 的3.05倍。LLaMA2-70B 模型吞吐量是 FP16 的 3.4倍。

这些加速都仅仅是 QUIK 对线性层的加速带来的。从下图的右图可见，当大部分计算都是 4-bit 精度时，attention，softmax，layernorm 的开销所占的比例就比较大了（19%）。

上图左图 Ideal 8-bits 和 Ideal 4-bits 是指直接使用 8-bit 核和 4-bit 核，而不考虑保持精度（即没有 outliers）

从上图可见 8-bit 实现提供了接近理想的加速。部署时，GPU 从 FP16 的 7 个降低为 5 个。

QUIK-4B 比 QUIK-8B 吞吐量提高了 15%，部署时，GPU 数降为了 3 个，显存只需要不到 50GB。为了保持精度，使用 outliers 和 8-bit 量化 down-project，使得 QUIK-4B 的吞吐量和理想的 4-bit 相比，下降在 15% 以内。目前还不知道权重激活全部 4-bit 量化，能够恢复精度的方法。

论文存在的问题

论文源码使用模拟的方式测出的 QUIK-4B 的 PPL 不准确

笔者使用论文源码进行了一些测试。发现使用源码提供的模拟方式（--sim_eval）测出的 PPL 与使用核（--kernels_eval）测出的 PPL 不一致的，特别对 CodeLLaMA 模型在 4-bit 量化时。笔者按照论文中的描述方式自己实现了模拟量化，与 QUIK 核测试结果一致。

论文源码不能复现论文中的层加速效果

笔者使用源码提供的核在 layer_benchmark.py 上进行测试（单卡 NVIDIA A100 ），并不能复现层加速效果。具体测试结果见下表：例如 QUIK-4B 最高只提供了 2.2 倍的加速，并不能达到论文中的 4.7 倍。

总结

虽然论文的源码存在一些问题，但是论文提出的这种通过对激活进行统计重排，而后利用矩阵的二阶信息进行量化误差补偿的思想，是值得我们学习借鉴的。

参考文献

[1]Ashkboos S, Markov I, Frantar E, et al. Towards end-to-end 4-bit inference on generative large language models[J]. arXiv preprint arXiv:2310.09259, 2023.