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

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

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

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

论文链接：
https://arxiv.org/abs/2404.12759

项目链接：
https://github.com/bytedance/decoupleQ

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

1. 背景

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

2. 传统量化

3. decoupleQ

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 的显存利用提供更多空间。

7. 总结与讨论

参考文献：

【1】Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Optq: Accurate quantization for generative pretrained transformers. In The Eleventh International Conference on Learning Representations, 2022.

【2】Wenqi Shao, Mengzhao Chen, Zhaoyang Zhang, Peng Xu, Lirui Zhao, Zhiqian Li, Kaipeng Zhang, Peng Gao, Yu Qiao, and Ping Luo. Omniquant: Omnidirectionally calibrated quantization for large language models. arXiv preprint arXiv:2308.13137, 2023

【3】Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023.