【量化部署】AWQ in MLSys 2024

一、引言

论文： AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
作者： MIT
代码： AWQ
特点： 该方法是一个仅量化权重的方法（W4A16），AWQ认为只有极少部分（1%）模型权重是十分重要的，在量化时应该着重保护它们。AWQ根据校准数据输入模型后得到的激活值的量级来锁定重要权重，通过放大这些权重来实现对他们的保护，放大因子是对应激活量级的$\alpha$次方。

二、详情

2.1 量化与反量化

最常规的量化是将浮点数除以一个值并四舍五入取得指定范围内的整数，公式如下：
$$\text{quantify}(\text{w})=\text{Round}(\frac{\text{w}}{\Delta }),\Delta =\frac{\max (|\text{w}|)}{2^{N-1}}$$

其中，$N$是量化后的位数，例如从FP16量化到INT4，那么$N=4$。该量化过程的示意图如下：

我们可以举一个例子，假设FP16的待量化权重最大最小值分别为-120.56和2.5，要将其量化至INT4范围$[-8,7]$的整型上，则有：
$$\begin{gathered} \text{quantify}(-120.56)=\text{Round}(\frac{-120.56}{\frac{\max (|-120.56|,|2.5|)}{2^{N-1}}})=-8 \\ \text{quantify}(2.5)=\text{Round}(\frac{2.5}{\frac{\max (|-120.56|,|2.5|)}{2^{N-1}}})=0 \end{gathered}$$

由于AWQ仅对权重做量化，因此上层输出的FP16的激活值需要先经过量化才能与权重进行数学运算，运算得到的值可能已经超出INT4的表示范围，因此需要反量化至真实值，然后再做INT4量化得到输出，示意图如下：

反量化的公式如下：
$$\text{dequantify}(\text{w})=\Delta \cdot \text{Round}(\frac{\text{w}}{\Delta }),\Delta =\frac{\max (|\text{w}|)}{2^{N-1}}$$

其实就是记录量化时的缩放因子$\Delta$，在反量化时乘上它。
$$\begin{gathered} \text{dequantify}(-120.56)=\frac{\max (|\text{w}|)}{2^{N-1}}\times \text{quantify}(-120.56)=\frac{120.56}{8}\times -8=-120.56 \\ \text{dequantify}(2.5)=\frac{\max (|\text{w}|)}{2^{N-1}}\times \text{quantify}(2.5)=\frac{120.56}{8}\times 0=0 \end{gathered}$$

-2.5经量化和反量化后得到0，可见量化是有误差的，产生的主要原因在于$\text{Round}(\cdot )$操作不可逆。

为帮助理解，下面给出一个INT8量化时，矩阵的运算过程：

⚠️ 图中例子$\Delta$的计算公式为$\Delta =\frac{\max (|\text{w}|)}{2^{N-1}-1}$，与AWQ论文中的略有不同。

2.2 选取重要权重

了解了量化和反量化过程，接下来就要确定我们对什么进行量化。AWQ是对权重进行量化并保护重要权重，有些其他方法会同时对激活值和权重量化。

首先，AWQ说明了保留1%的重要权重不做量化能够减少量化误差。

如上图，PPL是一个衡量指标（越低越好），可见保持1%的重要权重不变（不做量化，仍为FP16），PPL更低，从而证明保留1%的重要权重能够减少量化误差。

其次，AWQ说明了这1%的重要权重可以根据激活值的量级来选取。

如上表，右侧三大列是分别保存0.1%、1%、3%时根据激活值量级、权重量级、随机选取时的PPL。与完全不做量化的FP16那一列相比，损失最低的是根据激活值量级选取的权重量化策略。

之所以可以根据激活值的量级选取，是因为有些重要的激活值的量级很大，它们对模型来说很重要，因此与之对应的权重也就很重要。

如上图，激活值在某些通道上会有极大的幅值（量级）

2.3 保护重要权重

但是，如果我们保持重要权重为FP16，其他权重被量化为INT4，这就导致系统在执行矩阵运算时需要单独执行两个部分的运算，即混合精度，这会降低运算效率。

于是，AWQ提出放大重要权重然后统一执行INT4量化，从而起到保护重要权重又不影响运算效率的作用。

AWQ对放大能够降低量化误差做了简单的证明，假设不放大时的量化+反量化公式如下：

这里的$\text{w}$是指多个通道中的权重最大值，例如128个通道的权重（原文中称其为一组）。

则对于一个通道的权重来说，量化误差为：

其中$\Delta$和$x$都是FP16的，因此需要反量化后才能进行运算。

最原始的计算公式应该是$w\cdot x$，如果我们对$w$进行放大，并对$x$进行缩小，能够得到相同的结果，即$w\cdot x=(w\cdot s)\cdot (x\cdot \frac{1}{s})$。于是，我们可以得到将$w\cdot s$视为$w$，将$x\cdot \frac{1}{s}$视为$x$时的量化误差：

其中，${\Delta }^{\prime }$是对某通道进行放大后的一组通道的最大值（可能仍然等于$\Delta$）。

由于$\text{RoundErr}(\cdot )$所产生的误差始终在$[0,0.5]$之间，因为四舍五入只会产生这么多误差，所以其平均误差为0.25。无论是$\text{RoundErr}(\frac{w}{\Delta })$还是$\text{RoundErr}(\frac{w\cdot s}{{\Delta }^{\prime }})$产生的平均误差都一样。所以放大与不放大的误差比为：

${\Delta }^{\prime }$是对某通道进行放大后的一组通道的最大值，通常还是等于$\Delta$。例如，我们有两个通道的权重值为$[[0.1,0.2,0.3],[0.5,0.6,0.7]]$，我们对第一个通道进行2倍的放大，于是有$[[0.2,0.4,0.6],[0.5,0.6,0.7]]$，此时放大前后的最大值都还是0.7，所以有${\Delta }^{\prime }\approx \Delta$。

因为，${\Delta }^{\prime }\approx \Delta$且$s>1$（放大$w$），所以放大操作能使量化误差降低。

2.4 搜索放大因子$s$

理论上，我们只需要选取一个大于1的$s$即可做到对重要权重的保护，但其实并非所有大于1的$s$值都有利于降低量化误差。

如上表，取不同的$s$可能产生不同的性能表现。

于是，AWQ提出引入校准集帮助搜索放大因子，这样在推理时就可以直接使用事先确定的放大因子$s$进行重要权重的保护。

AWQ以量化+反量化权重与不量化权重乘上校准集$\text{X}$之间的误差为目标，进行$s$的网格搜索。

但是，这种搜索需要事先指定一个$s$的集合，这对人工的经验要求过高，于是AWQ提出根据校准集激活值各通道的平均量级来确定$s$，于是有：

其中，${\text{s}}_{\text{X}}$就是校准集激活值各通道的平均量级，即校准集$\text{X}$进入模型后得到激活值，然后求激活值的各个通道的幅值的平均值。

但是这个平均值${\text{s}}_{\text{X}}$并不一定大于1，所以AWQ给${\text{s}}_{\text{X}}$加了一个$\alpha$的指数项（$\alpha \in [0,1]$），最终将${{\text{s}}_{\text{X}}}^{\alpha }$确定为放大因子。

于是，我们可以在$[0,1]$的范围内对$\alpha$进行最小误差的搜索，确定之后即可得到最终的放大因子$s$。

• 校准集对放大因子的质量有直接影响，因此一般所使用的校准集需要具有较高的多样性。
• 由于AWQ并不需要进行梯度传递，只是使用校准集进行$s$的搜索，所以它不会在校准集上过拟合，具有更高的通用性。

致谢：

本博客仅做记录使用，无任何商业用途，参考内容如下：
模型量化一：量化基础 对称量化 非对称量化 极大值量化 零点量化
大模型量化一网打尽