AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

目录

• 概
• AWQ
• 代码

Lin J., Tang J., Tang H., Yang S., Chen W., Wang W., Xiao G., Dang X., Gan C. and Han S. AWQ: Activation-aware weight quantization for llm compression and acceleration. MLSys, 2024.

概

随着模型的参数量的增加, 推理成本也在显著增加, 本文提出一种量化方法: AWQ 量化, 以缓解这一问题. 其主要贡献在于对于"重要"权重的特殊处理, 以及 per-channel 的 scaling.

AWQ

• 作者首先发现, 权重中的元素并不是同等重要的, 大约有 1% 的权重, 如果把他们以更高精度的方式保存 (如, FP16), 就能取得显著的效果提升 (上图 (a) -> (b)).

• 但是这种方式有一个显著的缺点, 这种混合精度会使得实际的实现变得异常麻烦, 所以需要另外的手段取解决.

• 一般的对称量化形如:

  $$y = \mathbf{w} \mathbf{x} \longrightarrow y = Q(\mathbf{w}) \mathbf{x},$$

  其中

  $$Q(\mathbf{w}) = \Delta \cdot \text{Round}(\frac{\mathbf{w}}{\Delta}), \quad \Delta = \frac{\max(|\mathbf{w}|)}{2^{N-1}}.$$

  这里 $N$ 是量化 Bits.

• 我们可以采用另外一种更加灵活的方式, 考虑只改变其中的一个元素 $w$ 的量化方式:

  $$Q(w \cdot s) \cdot \frac{x}{s} = \Delta' \cdot \text{Round}(\frac{ws}{\Delta'}) \cdot x \cdot \frac{1}{s}.$$

  这里 $\Delta' = \max(|w_1|, |w|_2, \ldots, w \cdot s, \cdots) / 2^{N-1}$. 注意到:

  1. $\text{Round}(\cdot)$ 所带来的误差是差不多的;
  2. $\Delta'$ 由于只改变一个元素, 通常 $\Delta' \approx \Delta$.

• 当 $s > 1$ 的时候, 由于 $w \cdot s$ 的分布更加均匀地分布在量化区间内, 所以 $w \cdot s$ 相较于 $w$ 通常能够被更加精准地量化.

• 于是, 作者最终的 AWQ 为:

  $$\mathbf{s^*} = \text{argmin}_{\mathbf{s}} \mathcal{L}(\mathbf{s}) \\ \mathcal{L}(\mathbf{s}) = \| Q(\mathbf{W} \text{diag}(\mathbf{s})) (\text{diag}(\mathbf{s})^{-1} \cdot \mathbf{X}) - \mathbf{W} \mathbf{X} \|.$$

  其中 $\mathbf{X}$ 是根据一个比较小的 calibration set 得到的.

代码

[official-code]