GPTQ简介

GPTQ 的思想最初来源于 Yann LeCun 在 1990 年提出的 OBD 算法,随后 OBS、OBC(OBQ) 等方法不断进行改进,而 GPTQ 是 OBQ 方法的加速版。在介绍 GPTQ 算法之前,需要先介绍 OBD -> OBS -> OBQ 的演进过程。

OBD: Optimal Brain Damage

OBD 实际上是一种剪枝方法,用于降低模型复杂度,提高泛化能力。

如果要在模型中去除一些参数(即剪枝),直觉上,我们希望去除对目标函数 ( E ) 影响小的参数。于是我们可以对目标函数 ( E ) 做泰勒展开:

ΔE=igiΔwi+12ihiiΔwi2+12ijhijΔwiΔwj+O(Δw3)\Delta E = \sum_i g_i \Delta w_i + \frac{1}{2} \sum_i h_{ii} \Delta w_i^2 + \frac{1}{2} \sum_{i \neq j} h_{ij} \Delta w_i \Delta w_j + O(\Delta w^3)

其中,( 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 )。

于是,上式可以简化为:

ΔE=12ihiiΔwi2\Delta E = \frac{1}{2} \sum_i h_{ii} \Delta w_i^2

因此,删除一个参数时,对目标函数的影响就是 ( \frac{1}{2} h_{ii} \Delta w_i^2 )。所以我们只要计算海森矩阵 ( h_{ii} ),就可以知道每个参数对目标的影响。然后就可以按照影响从小到大给参数排个序,这样就确定了参数剪枝的次序。

OBS: Optimal Brain Surgeon

OBS认为,参数之间的独立性不成立,还是要考虑交叉项,因此上式变成了

ΔE=12ihiiΔwi2+12ijhijΔwiΔwj\Delta E=\frac{1}{2} \sum_{i} h_{i i} \Delta w_{i}^{2}+\frac{1}{2} \sum_{i \neq j} h_{i j} \Delta w_{i} \Delta w_{j}

用向量/矩阵形式表达会更加简明:

ΔE=12ΔwTHΔw\Delta E=\frac{1}{2} \Delta \mathbf{w}^{\mathrm{T}} \mathbf{H} \Delta \mathbf{w}

删除一个权重 wqw_{q},那么 Δw\Delta \mathbf{w} 的第 qq 维固定为 wq-w_{q},但其他维度的值可变,可以用于减少删除该权重带来的目标偏离。

Δw\Delta \mathbf{w} 的第 qq 维固定为 wq-w_{q},这是一个约束条件,我们可以表示为一个等式:
[
\mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}=0
]
其中 eq\mathbf{e}_{\mathbf{q}} 是一个 one - hot 向量,第 qq 个位置是 1,其余位置是 0。

我们希望找到最合适的权重 wqw_{q},使得删除它对目标的影响最小。这可以表示为一个最优化问题:
[
\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}}
]

于是,我们也只需要求解海森矩阵的逆,就可以计算每个参数 wqw_{q} 对目标的影响 12wq2[H1]qq\frac{1}{2} \frac{w_{q}^{2}}{\left[\mathbf{H}^{-1}\right]_{q q}},然后就可以按照影响从小到大给参数排个序,这样就确定了参数剪枝的次序。同时,每次剪枝一个参数,其他的参数也按照 Δw\Delta \mathbf{w} 更新一次。

这里的思想一直沿用到了 GPTQ 算法:也就是对某个 block 内的所有参数逐个量化,每个参数量化后,需要适当调整这个 block 内其他未量化的参数,以弥补量化造成的精度损失。

OBC

OBD 和 OBS 都存在一个缺点,就是剪枝需要计算全参数的海森矩阵(或者它的逆矩阵)。但在动辄上亿参数的神经网络下求解海森矩阵显然不可能。于是,我们可以假设参数矩阵的同一行参数互相之间是相关的,而不同行之间的参数互不相关,这样,海森矩阵就只需要在每一行内单独计算就行啦。

为了求解海森矩阵,我们需要确定目标函数的具体形式。

令参数矩阵的维度为 (drow,dcol)(d_{row},d_{col}),OBC 论文的目标函数定义为:
[
E=\sum_{i = 1}^{d_{row}} |\mathbf{W}{\mathbf{i},:} \mathbf{X}-\hat{\mathbf{W}}{\mathbf{i},:} \mathbf{X}|_{2}^{2}
]
直观来看,就是参数量化前后,给同样的输入(具有代表性的校准数据),输出结果的差异要尽可能小。

由于每一行的参数独立,所以我们只需要对每一行的量化后参数 W^i,:\hat{\mathbf{W}}_{\mathbf{i},:} 求海森矩阵:

EW^i,:2=Hi=2XXT,i=1,2,,drow\frac{\partial E}{\partial \hat{\mathbf{W}}_{\mathbf{i},:}^{2}}=\mathbf{H}_{\mathbf{i}} = 2\mathbf{X}\mathbf{X}^{\mathbf{T}},\forall i = 1,2,\cdots,d_{row}

如果我们对一行共 dcold_{col} 个参数删除 k 个参数,那么 OBC 的算法流程如下:

从 for 循环开始逐行分析:

Line 1:找到对目标函数影响最小的参数 p
Line 2:对参数 p 剪枝,并更新其他参数
Line 3:删除海森矩阵的 p 行 p 列,再求逆(这里用了数学的等价表达,降低了计算复杂度)
该算法的时间复杂度为O(kdcol2)O(k\cdot{d_{col}^2})

OBQ

OBQ(和OBC是同一篇文章)指出,剪枝是一种特殊的量化(即剪枝的参数等价于量化到0点),我们只需要修改一下OBC的约束条件即可:
[
\mathbf{e}{\mathbf{q}}^{\mathrm{T}} \cdot \Delta \mathbf{w}+w{q}=\operatorname{quant}(w_{q})
]
也就是说,OBC推导结果中的 wqw_{q} 替换成 wqquant(wq)w_{q}-\operatorname{quant}(w_{q}),就能得到一般量化情况下的权重更新公式:

Δw=wqquant(wq)[H1]qqH1eq and L=12(wqquant(wq))2[H1]qq\Delta \mathbf{w}=-\frac{w_{q}-\operatorname{quant}(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}-\operatorname{quant}(w_{q}))^{2}}{\left[\mathbf{H}^{-1}\right]_{q q}}

OBQ对一行做量化的时间复杂度为 O(dcol3)O(d_{col}^{3}),因此对整个参数矩阵做量化的时间复杂度为 O(drowdcol3)O(d_{row} \cdot d_{col}^{3})

至此,GPTQ的前置介绍就基本结束了。

GPTQ

OBQ 的算法复杂度还是太高了,GPTQ 对 OBQ 做了一些算法和性能上的优化,因而可以实现大模型的高效量化。

GPTQ 的创新点有:

  • OBS 采用贪心策略,先量化对目标影响最小的参数;但 GPTQ 发现直接按顺序做参数量化,对精度影响也不大。
    • 这项改进使得参数矩阵每一行的量化可以做并行的矩阵计算,同时将时间复杂度从 O(drowdcol3)O(d_{row} \cdot d_{col}^{3}) 降低到 O(max{drowdcol2,dcol3})O(\max\{d_{row} \cdot d_{col}^{2}, d_{col}^{3}\}),在大模型环境下,这项改进使得量化速度快了一个数量级;
  • Lazy Batch - Updates,延迟一部分参数的更新,它能够缓解 bandwidth 的压力;
  • Cholesky Reformulation,用 Cholesky 分解求海森矩阵的逆(参数更新的一部分 ),在增强数值稳定性的同时,不再需要对海森矩阵做更新计算,进一步减少了计算量。

我们主要关注第二个创新点,也就是 Lazy Batch - Updates 是如何缓解 bandwidth 的压力的。

问题:虽然 GPTQ 降低了时间复杂度,但这个算法的计算/通信比太低,通信带宽成为了瓶颈。
例如在量化某个参数矩阵的情况下,每次量化一个参数,其他所有未量化的参数都要按公式全都要更新一遍:

wwH:,p11[H1]pp(wpq(wp))\mathbf{w} \leftarrow \mathbf{w}-\mathbf{H}^{-1}_{:, p} \frac{1}{\left[\mathbf{H}^{-1}\right]_{p p}} \cdot\left(w_{p}-q\left(w_{p}\right)\right)

如果每行的量化并行计算,那么每次更新过程就需要 read + write 一次参数矩阵。如果参数矩阵的维度为 k×kk \times k,那么量化这个参数矩阵就需要读写 k 次参数,总共的 IO 量为 k3k^3 个元素。当 k 比较大时 (>=4096>=4096),需要读写的元素就非常多了,运行时间大都被 IO 占据。

思路:由于参数量化是一列一列按次序进行的,第 i 列的参数的量化结果受到前 i-1 列量化的影响,但第 i 列的量化结果不影响前面列的量化。因此我们不需要每次量化前面的列,就更新一遍第 i 列的参数,而是可以先记录第 i 列的更新量,在量化到第 i 列时,再一次性更新参数,这样就可以减少 IO 的次数。

具体实现:将参数矩阵按每 128 列划分为一个个 group,量化某一列时,group 内的参数立即更新,而 group 后面的列只记录更新量,延迟更新。当一个 group 的参数全部量化完成,再统一对后面的所有参数做一次更新。这就是 Lazy Batch-Updates。

Lazy Batch-Updates 不减少实际的计算量,但它能有效解决吞吐的瓶颈问题。