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.