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

发布于 2024-5-13 13:12
浏览
0收藏

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


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


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


  • 论文链接: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-AI.x社区


其中,

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

是量化前模型的浮点型 weights;s 和 z 是一个线性变换系数,表示 scale 和 zero point;α 和 β 是整数表示范围的上下界,例如对于 int4 量化,可以取 α = -8, β = 7;

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

表示取整函数,一般是四舍五入到临近的整数。


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


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


这样就将一个分布在

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

的浮点的 weights,线性映射到

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

的区间范围。


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


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


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


3. decoupleQ


观察 (1)~(3) 式在量化中的作用,如果我们换一种思路,就会发现其实我们不需要知道 (1) 式和 (2) 式。我们将一个大模型量化完毕以后,交付给下游引擎同学的时候,只需要知道 (3) 式中的

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

和 (s,z) 即可。也就是说,(3) 式中的 (s,z) 完全可以视作一个普通的仿射变换的系数,不需要保留其在 (2) 式中的含义。该仿射变换系数,可以通过数学优化的方法求出来。


进一步挖掘 (3) 式中,我们可以将一个大模型的参数,解耦 (decouple) 成整数部分

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

和浮点部分 (s,z)。这样解耦以后,模型量化的过程,就可以视作一个求解模型整数部分

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

和浮点部分 (s,z) 的过程。我们可以交替优化求解。为此,要确定优化目标函数和其约束条件。


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


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


其中,

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

是该层的输入,

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

是一个对称矩阵(若 X 的列都不全为零,则 H 是正定对称矩阵)。


一般来说,为了提高量化精度,我们可以对模型的权重用 per-channel 量化。在 per-channel 量化中,在优化 (4) 式的时候,

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

的各个列是独立进行优化的。所以我们只需要关注其中一列即可。


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


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


其中优化目标函数为


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


其中,w 是

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

的某一列,b 是

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

中对应的列。其他符号的定义与前文相同。


事实上,优化目标函数 (6) 与 (4) 是完全一致的,

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

就是反量化过程。


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


式子 (5) 的求解并不容易,因为有关于

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

的约束的存在,特别是

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

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


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


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


在这一步,作者在一个 transformer block 层面,将其中的各个 linear 层量化完毕以后,固定他们的整数部分

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

,而去微调浮点部分 (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-AI.x社区

w2 cuda kernel的实现原理


5. 实验


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


其中内部实验结果是:


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


该表格中,作者用 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-AI.x社区


该表格是 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-AI.x社区


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


7. 总结与讨论


在总结与讨论部分,作者也指出了 decoupleQ 目前存在的两个风险:


1. decoupleQ 旨在使用数学优化的手段,将量化前后的 L2 loss 最小化。然而,layer 层面或者 block 层面的 L2 loss 最小化,未必能够代表最终模型精度的最优;

2. 在 (5) 式和 (7) 式的优化过程中,在求解

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

和 (s,z) 的时候,只针对一小部分校准数据来求解,这使得 decoupleQ 容易对校准数据过拟合。


尽管如此,作者也指出,将模型参数解耦为整数部分和浮点部分的思路,是很有意义的。如果有标签数据集存在,我们可以量化完毕以后,固定整数部分,使用有标签数据集来专门训练 (s,z),以进一步提升模型的精度。这样既保证了模型的泛化性能(源于固定整数部分

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

),又能在特定子任务上发挥能力(源于微调浮点部分)。在字节跳动的实际业务中,在上一版的模型量化完毕上线以后,下一版更新的时候,也可以只训练模型中的浮点部分。


本文转自机器之心 ,作者:机器之心


原文链接:​​https://mp.weixin.qq.com/s/tv4H0LA-e0qzX33NIWCdAg​

收藏
回复
举报
回复
相关推荐