AI大模型量化技术:从QAT到BitNet
AI大模型量化技术:从QAT到BitNet
随着AI大模型的快速发展,模型量化技术成为降低计算成本、提高部署效率的关键手段。本文将深入探讨两种先进的量化方法:量化感知训练(QAT)和1位大型语言模型(BitNet),帮助读者理解这些技术如何在保持模型性能的同时,实现资源的优化利用。
一、量化感知训练
在前面的内容中,我们已经介绍了如何在训练之后量化一个模型。这种方法的一个缺点是,量化并不考虑实际的训练过程。
而量化感知训练(QAT)与训练后量化(PTQ)在模型训练完成之后进行量化不同,QAT旨在在训练期间学习量化过程。
QAT通常比PTQ更精确,因为量化过程已在训练中被考虑。其工作原理如下:
在训练过程中,引入所谓的“假”量化。这是一个首先将权重量化为例如INT4,然后再反量化回FP32的过程:
这个过程允许模型在训练、损失计算和权重更新过程中考虑量化过程。QAT试图探索损失中的“宽”极小值以最小化量化误差,因为“窄”极小值往往会导致较大的量化误差。
例如,假设我们在反向传播过程中没有考虑量化。根据梯度下降选择损失最小的权重。如果它处于“窄”极小值,那将引入更大的量化误差。
如果我们考虑量化,将在一个“宽”极小值中选择一个不同的更新权重,其量化误差将大大降低。
所以尽管PTQ在高精度(例如FP32)中有更低的损失,但QAT在低精度(例如INT4)中会获得更低的损失
二、1位大型语言模型的时代:BitNet
正如我们之前看到的,量化到4位已经相当小了,但如果我们进一步减少呢?
这就是BitNet的用武之地,它使用-1或1来表示模型权重的单一位。它通过将量化过程直接注入到Transformer 架构中实现这一点。
Transformer 架构是大多数LLM的基础,它由涉及线性层的计算组成:
这些线性层通常用更高的精度表示,如FP16,并且是大多数权重所在的地方。
而BitNet用它们称为BitLinear的东西替换了这些线性层:
BitLinear层的工作方式与普通线性层相同,根据权重乘以激活来计算输出。但是BitLinear层使用1位来表示模型的权重,并使用INT8来表示激活:
BitLinear层,如量化感知训练(QAT),在训练期间执行一种“假”量化形式,以分析权重和激活量化的效果:
这种方法显著减少了模型的存储和计算需求,使得在资源受限的环境中部署大型语言模型变得可行。同时,通过这种极端的量化方法,BitNet在维持性能的同时大幅降低了能耗和运行成本
在论文中,他们使用γ而不是α,但由于我们在这个示例中使用了a,所以我继续使用这个名词。另外,请β与我们在零点量化中使用的不同,是平均绝对值。
下面我们看看他是如何工作的
1. 权重量化
在训练过程中,权重存储在INT8中,然后使用一种称为符号函数的基本策略,将其量化为1位。
它将权重的分布移动到以0为中心,然后将0左边的所有值赋值为-1,右边的所有值赋值为1:
此外,它还跟踪一个值β(平均绝对值),因为稍后将用它进行去量化。
2. 激活量化
为了量化激活值,BitLinear使用_absmax量化_将激活值从FP16转换为INT8,因为在矩阵乘法(×)中它们需要更高的精度。
此外,它还跟踪了α(绝对值),因为稍后将用它进行去量化。
3. 去量化
上面跟踪了α(激活值的最大绝对值)和β(权重的平均绝对值),这些值将帮助我们将激活值反量化回FP16。
输出激活值使用 {α, γ} 重新缩放,以将其反量化到原始精度:
这个过程相对简单,并允许模型仅用两个值表示,要么是-1,要么是1。使用这种方法,作者观察到随着模型大小的增长,1位和FP16训练之间的性能差距变得越来越小。
并且作者发现,这仅适用于较大的模型(>30B 参数),而在较小的模型中,差距仍然相当大。