字节开源大模型量化新思路,2

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

字节开源大模型量化新思路,2-bit量化模型精度齐平fp16
AIxiv专栏是本站发布学术、技术内容的栏目。过去数年,本站AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。如果您有优秀的工作想要分享,欢迎投稿或者联系报道。投稿邮箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

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

最近,字节跳动推出了一个全新的量化思路,抛弃传统的量化范式,从数学优化的角度来对量化任务建模。文章放在了 arXiv,代码已经开源,可以一键复现文中的所有结果。该量化思路基于数学优化,从数学优化的角度来对量化任务建模,通过最大化目标函数或最小化损失函数来寻找最优解。这种思路在实验中获得了很好的效果,取得了令人满意的结果。

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

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

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

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

1. 背景

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

2. 传统量化

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

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

其中,字节开源大模型量化新思路,2-bit量化模型精度齐平fp16是量化前模型的浮点型 weights;s 和 z 是一个线性变换系数,表示 scale 和 zero point;α 和 β 是整数表示范围的上下界,例如对于 int4 量化,可以取 α = -8, β = 7;字节开源大模型量化新思路,2-bit量化模型精度齐平fp16表示取整函数,一般是四舍五入到临近的整数。

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

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

这样就将一个分布在字节开源大模型量化新思路,2-bit量化模型精度齐平fp16的浮点的 weights,线性映射到字节开源大模型量化新思路,2-bit量化模型精度齐平fp16的区间范围。

在反量化的时候,一般使用如下公式:

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

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

3. decoupleQ

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

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

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

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

其中,字节开源大模型量化新思路,2-bit量化模型精度齐平fp16是该层的输入,字节开源大模型量化新思路,2-bit量化模型精度齐平fp16是一个对称矩阵(若 X 的列都不全为零,则 H 是正定对称矩阵)。

一般来说,为了提高量化精度,我们可以对模型的权重用 per-channel 量化。在 per-channel 量化中,在优化 (4) 式的时候,字节开源大模型量化新思路,2-bit量化模型精度齐平fp16的各个列是独立进行优化的。所以我们只需要关注其中一列即可。

至此,优化目标可以写作如下:(为了符号简洁,文章中重新定义符号):

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

其中优化目标函数为

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

其中,w 是图片的某一列,b 是字节开源大模型量化新思路,2-bit量化模型精度齐平fp16中对应的列。其他符号的定义与前文相同。

事实上,优化目标函数 (6) 与 (4) 是完全一致的,字节开源大模型量化新思路,2-bit量化模型精度齐平fp16就是反量化过程。

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

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

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

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

在这一步,作者在一个 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

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

                                  w2 cuda kernel的实现原理

5. 实验

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

其中内部实验结果是:

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

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

开源对比实验结果是:

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

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

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

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

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.

以上就是字节开源大模型量化新思路,2-bit量化模型精度齐平fp16的详细内容,更多请关注叮当号网其它相关文章!

文章来自互联网,只做分享使用。发布者:周斌,转转请注明出处:https://www.dingdanghao.com/article/480420.html

(0)
上一篇 2024-05-13 20:45
下一篇 2024-05-13 21:03

相关推荐

联系我们

在线咨询: QQ交谈

邮件:442814395@qq.com

工作时间:周一至周五,9:30-18:30,节假日休息

关注微信公众号