AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
随着深度学习大语言模型的越来越火爆，大语言模型越做越大，使得其推理成本也水涨船高。模型量化，成为一个热门的研究课题。
近日，字节跳动语音团队推出一个全新的量化思路，抛弃传统的量化范式，从数学优化的角度来对量化任务建模。文章放在了 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 的显存利用提供更多空间。
7. 总结与讨论
在总结与讨论部分，作者也指出了 decoupleQ 目前存在的两个风险：
1. decoupleQ 旨在使用数学优化的手段，将量化前后的 L2 loss 最小化。然而，layer 层面或者 block 层面的 L2 loss 最小化，未必能够代表最终模型精度的最优；
2. 在 (5) 式和 (7) 式的优化过程中，在求解
和 (s,z) 的时候，只针对一小部分校准数据来求解，这使得 decoupleQ 容易对校准数据过拟合。
尽管如此，作者也指出，将模型参数解耦为整数部分和浮点部分的思路，是很有意义的。如果有标签数据集存在，我们可以量化完毕以后，固定整数部分，使用有标签数据集来专门训练 (s,z)，以进一步提升模型的精度。这样既保证了模型的泛化性能（源于固定整数部分
），又能在特定子任务上发挥能力（源于微调浮点部分）。在字节跳动的实际业务中，在上一版的模型量化完毕上线以后，下一版更新的时候，也可以只训练模型中的浮点部分。
参考文献：
【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.
机器之心将于2024年6月1日举办「CVPR 2024 线上论文分享会」，将设置 Keynote、 论文分享、企业招聘等环节，就业内关注的 CV 热门主题邀请顶级专家、论文作者与观众做学术交流。
© THE END 
转载请联系本公众号获得授权
投稿或寻求报道：[email protected]
继续阅读
阅读原文