AWQ: Activation-Aware Weight Quantization For On-Device LLM Compression AND Acceleration
AWQ: Activation-Aware Weight Quantization For On-Device LLM Compression AND Acceleration
研究方向: On-device LLM
两个主要的内容是
- Memory size: 内存容量(模型参数能不能放到设备里面)
- Memory bandwidth(内存带宽) : 即使模型装得下，能否快速的读取权重，进行计算
要理解这个问题，首先要理解到量化的一个核心 压缩率不等于加速比
除此之外，作者的一个核心观察是
Weights are not equally important
Protecting only 1% of salient weights can greatly reduce quatization error
作者把重要的权重称为 salient weights， 显著重要的权重
实验中，作者发现只保护一小部分这样的权重就能够大幅度优化量化后的模型表现。然而问题同时也来了，怎么判断那些权重重要?
当然，权重的重要程度是和 Activation相关的，而不能只关注weight本身的大小。
y = Wx ， 如果量化之后， Δy = (What - W) x
即使只量化了weight，也要了解activation的分布
作者声称，他们的方法不需要backpropagation，摘要下面写道，AWQ does not rely on any backpropagation or reconstruction
其中
- backpropagation是通过反向传播优化参数，例如 quantization-aware training (QAT）
- Reconstruction 类似于GPTQ，量化后某一层的输出尽量接近原模型输出。
AWQ 则希望避免昂贵的训练或者重建过程，主要通过校准数据观察 Activation,然后搜索缩放因子。
Introduction
别的都是套话，我们先不管
现有的方法QAT （ 训练成本高) PTQ(预训练后量化容易出现损失)
然后GPTQ和作者的这个很接近，GPTQ通过二阶信息补充量化误差，是作者认为最接近的工作，但是作者认为GPTQ的重建过程可能导致过拟合。
而常用的Mixed precision会产生不规则的混合精度表示，使得实际的kernel变得复杂。所以AWQ希望不需要保存单独的FP16副本，而是通过Per-channel Scaling， 在统一的量化格式内优先降低重要权重的误差。
Measures
他们提出了一个假设
activation Magnitude 是否真的比weight magnitude 更适合判断权重的重要性
他们然后提出了一个实验
FP16 原始模型的结果 是 10.86
全部INT3 RTN 量化 23.54
按照weight大小保留1% 22.37
按照activation 大小保留 11.39
按照activation大小保护1%，perplexity下降了很多，非常接近FP16
而这也带来了一个实际的问题，weight 相对于模型来说是固定的，而activation是由输入决定的。
所以activation如果是动态的，那怎么确定哪些weight重要呢？
AWQ并没有说要去每次都保护重要的权重，作者使用一小部分Calibration data，预先统计各个输入的平均activation magnitude
利用这个信息来在量化前确定哪些通道比较重要，但是这个东西建立在一个假设上: 虽然activation 随着输入变化，但是不同的通道的重要性存在统计规律。
所以思考到这一步，我们就知道了现在手上的两个情况
- 我们可以利用校准数据的activation 统计量，提前找出重要通道
- 直接把这些通道保留位fp16,会使得权重表示和硬件执行变得复杂。
然后我们来看一个东西，一个输出的计算是 权重w x 输入x
现在我们引入一个缩放系数 s， y = (ws) (x/s)
但是加入量化操作Q， y = Q(w)x 和 y = Q(ws) x/s 这两个结果通常是不一样的！
用具体数字看看:
W = [ 1.0 , 0.2 ]
量化步长是 Δ = 1/3 ， 那么正数就只有 0， 1/3 ， 2/3 ， 1.
现在假设第二个权重对应的activation 是 10, 真实输出贡献应该是 2
那么普通的INT3 的输出是3.333， 而 scaling s=2 的输出是 1.667，通过这样，我们就把误差减小了很多，原因是原来的权重近似是 0.333，而现在相当于是 0.333/2 来表示原始权重
但是如果我们把权重放缩过大，有可能会导致最大的权重不再是最大值，整个组的量化范围就会改变。 [1.0, 0.2] -> [1.0 , 2.0] 最大值就是2.0了，步长就是 2/3
所以，AWQ 通过调整每个输入通道的 Scaling，在固定比特预算下重新分配量化误差，最终追求整体输出的准确性。
我们刚刚已经知道了，可以通过 Y = (WD) (D^-1)x 来等价变换，但是呢量化会引出误差，优化目标就是最小化这个误差。
寻找一组scaling factors 让量化前后输出尽可能接近。问题是，如果由一个线性层有4096个输入通道，那么就有4096个待优化的sj
更麻烦的是量化包含round操作，不可能用梯度下降来优化的。所以作者只能说利用activation构造搜索空间
整个方法可以说是: 利用 Activation 统计量决定不同通道之间的相对保护强度，再由一个超参数控制保护程度。
那么搜索的时候，就可以是先收集每层输入X, 统计平均绝对值a， 然后对候选alpha 计算 a^α，然后算误差。
所以总结一下呢，就是量化前，利用校准数据统计activation, 搜索每层合适的α，计算每个输入通道对应的sj，缩放权重并进行分组低比特量化。最后保存量化权重和必要的量化元数据。
Tiny Chat
到这里，我们已经清楚AWQ怎么进行量化的了，但是从system角度来看，还有一个很大的问题:
权重压缩了四倍，不一定推理加速四倍
所以Tiny Chat主要做了三件事情来加快。
- On the fly Dequantization 在计算的kernel里面反量化，不把权重写回去到Dram
- SIMD-aware Packing 减少解包开销，离线重排低比特权重，使得运行的时候更方便用SIMD指令解包。
- Kernel fusion 减少中间结果和启动开销，融合QKV Normalization算子
Experiment
AWQ在INT3下面的优势更明显，INT4下面，GPTQ-Reorder的 PPL 是 5.63， AWQ 是5.60 ，两者差距只有0.03
但是在INT3下面 GPTQ 是 6.42， AWQ是6.24 ， 差距扩到到了0.18
模型越大，量化的误差未必越严重。 例如Llama-2-70B, AWQ的PPL和FP16只有0.09. 有可能是大模型在一些情况下的比特容忍度比较高。
然后作者在Mistral-7B 和 Mistral- 8x7B-instruct，其中后者是MoE.其中，AWQ也能够在MoE架构上面保持相对不错的INT3/INT4 量化质量
极低比特的情况也很值得关注，进一步尝试的INT2量化很有意思
FP16差不多是 14~10 量级的 ppl，而 RTN直接上万，GPTQ 四十多， 而AWQ+GPTQ 可以保持 13~35.
所以，这个地方其实证明AWQ 的 scaling 和 GPTQ 的误差补偿具有一定的互补性
然后INT2确实难度要大很多。
除此之外，AWQ在校准数据上面的需求也比GPTQ要低一些，在作者测试的两种分布情况下，calibration 和测试不一致的时候，AWQ也可以保留更好的结果。
Tinychat部分:
从 52 到 62 tokens/s，是 TinyChat 的 FP16 系统优化带来的收益，约为 1.19 倍。
从 62 到 194 tokens/s，才是在 TinyChat 自身 FP16 实现的基础上，使用 W4A16 量化 Kernel 后获得的约 3.13 倍加速。
在jetson Orin 64GB上，Tinychat达到了40tk/s的速度。
我的总结。
这篇文章的核心创新点可以归结到一个地方 ： 我们在量化的时候应该针对激活去量化，而不是对权重进行量化。
通过scaling，动态的选择量化的位，再辅佐加上system的优化，会带来更好的一个能力。
