开源模型

字节开源大模型量化新思路，2-bit量化模型精度齐平fp16

2024 年 5 月 14 日

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：；

随着深度学习大语言模型的越来越火爆，大语言模型越做越大，使得其推理成本也水涨船高。模型量化，成为一个热门的研究课题。

近日，字节跳动语音团队推出一个全新的量化思路，抛弃传统的量化范式，从数学优化的角度来对量化任务建模。文章放在了 arXiv，代码已经开源，可以一键复现文中的所有结果：

论文链接：https://arxiv.org/abs/2404.12759
项目链接：https://github.com/bytedance/decoupleQ
W2 算子：https://github.com/NVIDIA/TensorRT-LLM/pull/1568

1. 背景

大模型的迅速发展，使得推理成本越来越高。模型量化，作为一个降低推理成本的技术方案，得到了越来越多的关注与研究。然而，在传统的量化范式下，模型的精度在极低比特下会迅速下降。基于此，作者们提出了一种新的量化思路，将模型参数解耦为整数部分和浮点部分，从数学优化的角度来对量化任务建模，使得在极低比特下，模型依然能保持较高的精度。这样做的优势是明显的，我们不再需要关注量化特有的问题，比如如何处理敏感通道，如何处理 outlier 等等，而是只需要将量化问题进行数学建模，找到一个合适的优化目标函数，然后去求解该函数。

2. 传统量化

传统上，我们对一个模型的量化思路是：

其中，是量化前模型的浮点型 weights；s 和 z 是一个线性变换系数，表示 scale 和 zero point；α 和 β 是整数表示范围的上下界，例如对于 int4 量化，可以取 α = -8, β = 7；表示取整函数，一般是四舍五入到临近的整数。

关于 s 和 z 的取值，一般来说，对于非对称量化，我们可以取：

这样就将一个分布在的浮点的 weights，线性映射到的区间范围。

在反量化的时候，一般使用如下公式：

在这种传统的量化方案中，我们需要关注很多量化特有的细枝末节问题，比如针对敏感通道，我们有敏感通道的处理方式；针对 outlier，我们有 outlier 的处理方式。这种头痛医头脚痛医脚的处理范式，难以应付复杂多变的业务场景。字节跳动的研究者们试图把这些问题都抽象出来，从一个宏观的视角来看待量化问题。我们只需要建立一个抽象的优化目标函数，再去求解这个目标函数即可。

3. decoupleQ

观察 (1)~(3) 式在量化中的作用，如果我们换一种思路，就会发现其实我们不需要知道 (1) 式和 (2) 式。我们将一个大模型量化完毕以后，交付给下游引擎同学的时候，只需要知道 (3) 式中的和 (s,z) 即可。也就是说，(3) 式中的 (s,z) 完全可以视作一个普通的仿射变换的系数，不需要保留其在 (2) 式中的含义。该仿射变换系数，可以通过数学优化的方法求出来。

进一步挖掘 (3) 式中，我们可以将一个大模型的参数，解耦 (decouple) 成整数部分和浮点部分 (s,z)。这样解耦以后，模型量化的过程，就可以视作一个求解模型整数部分和浮点部分 (s,z) 的过程。我们可以交替优化求解。为此，要确定优化目标函数和其约束条件。

对于一个 linear 层，我们可以构建如下优化目标函数：

其中，是该层的输入，是一个对称矩阵（若 X 的列都不全为零，则 H 是正定对称矩阵）。

一般来说，为了提高量化精度，我们可以对模型的权重用 per-channel 量化。在 per-channel 量化中，在优化 (4) 式的时候，的各个列是独立进行优化的。所以我们只需要关注其中一列即可。

至此，优化目标可以写作如下：（为了符号简洁，文章中重新定义符号）：

其中优化目标函数为

其中，w 是的某一列，b 是中对应的列。其他符号的定义与前文相同。

事实上，优化目标函数 (6) 与 (4) 是完全一致的，就是反量化过程。

将一个量化问题转化为形如(5)式的数学优化问题，正是 decoupleQ 区别于传统量化论文的关键所在。这种转化，使得我们只需要专注于求解 (5) 式即可，不再需要处理关于量化本身的细枝末节的问题，比如 outlier，等等。

式子 (5) 的求解并不容易，因为有关于的约束的存在，特别是这个非凸约束。在论文中，作者给出一种交替求解的思路，即获得关于 (s,z) 和 w 的良好的初始化以后，交替迭代求解 (s,z) 和 w。在求解 (s,z) 的时候，注意到式子 (5) 关于 (s,z) 是一个无约束的二次型，可以直接将目标函数求导，令导数为零，从而获得解析解。在求解 w 的时候，作者采取了两个层面的近似，第一层近似收敛性更高，但是求解慢；第二层近似采样 GPTQ【1】的思想，收敛性稍差，但是求解更快。

为了进一步提升量化模型精度，作者指出，除了在 layer 层面做 mse 最小化以外，还可以在 block 层面做 mse 最小化，即：

在这一步，作者在一个 transformer block 层面，将其中的各个 linear 层量化完毕以后，固定他们的整数部分，而去微调浮点部分 (s,z) 以及 layer norm 的相关参数。实验表面，这一步微调，能够进一步提升模型精度。

4. W2 算子实现

要对量化后的模型进行推理，需要量化算子的支持，在业界没有现成的 w2a16 的算子可用，作者们基于 Tensorrt-LLM 中的 w4 算子开发了 w2 的 Gemm cuda kernel, 实现了 w2a16 模型的高效推理。

量化模型本身是以 2bit weight 的形式加载和存储在显存中，因此会占用比较小的显存。我们的 cuda kernel 通过在运行时将 2bit 的 weight 加载到寄存器中，再利用硬件指令高效转换成 bf16 的形式与 activation 进行 gemm 运算。因为我们的场景受限于 latency, generation 阶段的 batchsize 比较小，此时矩阵乘受限于 weight 的访存，这种实现会大大减少访存量，提升模型的性能。在实现过程中，结合了算法搜索以及 SpiltK Parallel Reduce，进一步能提升模型的性能，实测在 batchsize=1 的情况下，在 L 卡上 w2a16 Gemm 性能相比 w4a16 能提升 1.4x-1.7x 不等。

算子链接：https://github.com/NVIDIA/TensorRT-LLM/pull/1568

w2 cuda kernel的实现原理

5. 实验

作者在文章给出了字节跳动内部的 ASR 实验结果，和开源的实验对比结果：

其中内部实验结果是：

该表格中，作者用 word err rate (WER) 来衡量 ASR 的准确率。作者尝试使用不同的方法将模型量化为 W2A16g64。量化前的浮点模型的 wer 是 6.68%，使用 GPTQ【1】量化以后是 6.83%，带有 block 最小化的 decoupleQ 量化以后的 wer 是 6.70%，该结果与量化前的浮点模型的 wer 很接近。同时也 report 了量化所需要的耗时。量化高精度的代价，是量化耗时较长。在实际业务中，在使用 decoupleQ 对模型量化完毕以后，固定整数部分，使用有标签数据集微调 scale 和 zero，模模型精度有进一步的提升。

开源对比实验结果是：

该表格是 decoupleQ 和其他方法在 Llama-1/2 上的量化结果比较。以 perplexity (PPL) 作为评价指标。可以看出，在同样的量化配置下，deoucpleQ 的 PPL 在绝大多数时候会低于其他方法。

6. 业务收益

decoupleQ 量化技术在字节跳动语音部门现在被广泛使用。已经上线于语音生成模型（Text-to-Speech），语音识别模型（automic speech recognition）等等，落地于豆包、飞书、抖音等产品中。大量上线业务表明，基于 decoupleQ 的量化，W4A16 的推理精度已经完全能和 fp16/bf16 推理持平；W2A16 的精度只略差于 fp16/bf16 精度（对浮点部分 sft 以后，精度能和 fp16/bf16 持平）。尽管论文中只介绍了 weight-only 的量化，但是在实际业务中，在 weight 获得良好的量化以后，对 activation 的量化也便能简单许多。

在硬件加速上相比 fp16、w8fp16、w4fp16 获得了不错的加速效果，在小 batch 下 w2 矩阵乘的性能相比 fp16 提升 5-6 倍，相比 w4 提升 1.5-1.7 倍。在内部业务模型上，w2fp16 相比 fp16 性能有 3-5 倍的提升，相比 w4fp16 性能有 1.25-1.4 倍的性能提升，同时也会使得模型 weight 占用显存大幅下降，为 runtime 的显存利用提供更多空间。