AWQ归纳总结
GPTQ简介
GPTQ 的思想最初来源于 Yann LeCun 在 1990 年提出的 OBD 算法,随后 OBS、OBC(OBQ) 等方法不断进行改进,而 GPTQ 是 OBQ 方法的加速版。在介绍 GPTQ 算法之前,需要先介绍 OBD -> OBS -> OBQ 的演进过程。
OBD: Optimal Brain Damage
OBD 实际上是一种剪枝方法,用于降低模型复杂度,提高泛化能力。
如果要在模型中去除一些参数(即剪枝),直觉上,我们希望去除对目标函数 ( E ) 影响小的参数。于是我们可以对目标函数 ( E ) 做泰勒展开:
其中,( g_i = \frac{\partial E}{\partial w_i} ) 为参数的梯度,( h_{ij} = \frac{\partial^2 E}{\partial w_i \partial w_j} ) 为海森矩阵的一个元素。
OBD 做了一些假设,对上述进行简化:
- 假设目标函数是二阶的,所以我们不考虑高阶项 ( O(\Delta w^3) )。
- 假设模型训练已充分收敛,因此所有参数的一阶偏导均为 0:( g_i = 0, \forall i )。
- 假设删除模型的任意一个参数后,其它参数对目标函数的影响不变。也就是说,每个参数对目标函数的影响是独立的。于是我们可以不考虑交叉项:( h_{ij} = 0, \forall i, j, i \neq j )。
于是,上式可以简化为:
因此,删除一个参数时,对目标函数的影响就是 ( \frac{1}{2} h_{ii} \Delta w_i^2 )。所以我们只要计算海森矩阵 ( h_{ii} ),就可以知道每个参数对目标的影响。然后就可以按照影响从小到大给参数排个序,这样就确定了参数剪枝的次序。
OBS: Optimal Brain Surgeon
OBS认为,参数之间的独立性不成立,还是要考虑交叉项,因此上式变成了
用向量/矩阵形式表达会更加简明:
删除一个权重 ,那么 的第 维固定为 ,但其他维度的值可变,可以用于减少删除该权重带来的目标偏离。
的第 维固定为 ,这是一个约束条件,我们可以表示为一个等式:
[
\mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}=0
]
其中 是一个 one - hot 向量,第 个位置是 1,其余位置是 0。
我们希望找到最合适的权重 ,使得删除它对目标的影响最小。这可以表示为一个最优化问题:
[
\min {\Delta \mathbf{w}, w{q}} \frac{1}{2} \Delta \mathbf{w}^{\mathrm{T}} \mathbf{H} \Delta \mathbf{w} \quad \text { s.t. } \mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}=0
]
用 Lagrange (拉格朗日)乘数法求解:
[
L=\frac{1}{2} \Delta \mathbf{w}^{\mathrm{T}} \mathbf{H} \Delta \mathbf{w}+\lambda\left(\mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}\right)
]
可以得到:
[
\Delta \mathbf{w}=-\frac{w_{q}}{\left[\mathbf{H}^{-1}\right]{q q}} \mathbf{H}^{-1} \cdot \mathbf{e}{\mathbf{q}} \quad \text { and } \quad L=\frac{1}{2} \frac{w_{q}^{2}}{\left[\mathbf{H}^{-1}\right]_{q q}}
]
于是,我们也只需要求解海森矩阵的逆,就可以计算每个参数 对目标的影响 ,然后就可以按照影响从小到大给参数排个序,这样就确定了参数剪枝的次序。同时,每次剪枝一个参数,其他的参数也按照 更新一次。
这里的思想一直沿用到了 GPTQ 算法:也就是对某个 block 内的所有参数逐个量化,每个参数量化后,需要适当调整这个 block 内其他未量化的参数,以弥补量化造成的精度损失。
OBC
OBD 和 OBS 都存在一个缺点,就是剪枝需要计算全参数的海森矩阵(或者它的逆矩阵)。但在动辄上亿参数的神经网络下求解海森矩阵显然不可能。于是,我们可以假设参数矩阵的同一行参数互相之间是相关的,而不同行之间的参数互不相关,这样,海森矩阵就只需要在每一行内单独计算就行啦。
为了求解海森矩阵,我们需要确定目标函数的具体形式。
令参数矩阵的维度为 ,OBC 论文的目标函数定义为:
[
E=\sum_{i = 1}^{d_{row}} |\mathbf{W}{\mathbf{i},:} \mathbf{X}-\hat{\mathbf{W}}{\mathbf{i},:} \mathbf{X}|_{2}^{2}
]
直观来看,就是参数量化前后,给同样的输入(具有代表性的校准数据),输出结果的差异要尽可能小。
由于每一行的参数独立,所以我们只需要对每一行的量化后参数 求海森矩阵:
如果我们对一行共 个参数删除 k 个参数,那么 OBC 的算法流程如下:
从 for 循环开始逐行分析:
Line 1:找到对目标函数影响最小的参数 p
Line 2:对参数 p 剪枝,并更新其他参数
Line 3:删除海森矩阵的 p 行 p 列,再求逆(这里用了数学的等价表达,降低了计算复杂度)
该算法的时间复杂度为。
OBQ
OBQ(和OBC是同一篇文章)指出,剪枝是一种特殊的量化(即剪枝的参数等价于量化到0点),我们只需要修改一下OBC的约束条件即可:
[
\mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}=\operatorname{quant}(w_{q})
]
也就是说,OBC推导结果中的 替换成 ,就能得到一般量化情况下的权重更新公式:
OBQ对一行做量化的时间复杂度为 ,因此对整个参数矩阵做量化的时间复杂度为 。
至此,GPTQ的前置介绍就基本结束了。
GPTQ
OBQ 的算法复杂度还是太高了,GPTQ 对 OBQ 做了一些算法和性能上的优化,因而可以实现大模型的高效量化。
GPTQ 的创新点有:
- OBS 采用贪心策略,先量化对目标影响最小的参数;但 GPTQ 发现直接按顺序做参数量化,对精度影响也不大。
- 这项改进使得参数矩阵每一行的量化可以做并行的矩阵计算,同时将时间复杂度从 降低到 ,在大模型环境下,这项改进使得量化速度快了一个数量级;
- Lazy Batch - Updates,延迟一部分参数的更新,它能够缓解 bandwidth 的压力;
- Cholesky Reformulation,用 Cholesky 分解求海森矩阵的逆(参数更新的一部分 ),在增强数值稳定性的同时,不再需要对海森矩阵做更新计算,进一步减少了计算量。
我们主要关注第二个创新点,也就是 Lazy Batch - Updates 是如何缓解 bandwidth 的压力的。
问题:虽然 GPTQ 降低了时间复杂度,但这个算法的计算/通信比太低,通信带宽成为了瓶颈。
例如在量化某个参数矩阵的情况下,每次量化一个参数,其他所有未量化的参数都要按公式全都要更新一遍:
如果每行的量化并行计算,那么每次更新过程就需要 read + write 一次参数矩阵。如果参数矩阵的维度为 ,那么量化这个参数矩阵就需要读写 k 次参数,总共的 IO 量为 个元素。当 k 比较大时 (),需要读写的元素就非常多了,运行时间大都被 IO 占据。
思路:由于参数量化是一列一列按次序进行的,第 i 列的参数的量化结果受到前 i-1 列量化的影响,但第 i 列的量化结果不影响前面列的量化。因此我们不需要每次量化前面的列,就更新一遍第 i 列的参数,而是可以先记录第 i 列的更新量,在量化到第 i 列时,再一次性更新参数,这样就可以减少 IO 的次数。
具体实现:将参数矩阵按每 128 列划分为一个个 group,量化某一列时,group 内的参数立即更新,而 group 后面的列只记录更新量,延迟更新。当一个 group 的参数全部量化完成,再统一对后面的所有参数做一次更新。这就是 Lazy Batch-Updates。
Lazy Batch-Updates 不减少实际的计算量,但它能有效解决吞吐的瓶颈问题。