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+12∑ihiiΔwi2+12∑i≠jhijΔ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=12∑ihiiΔ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=12∑ihiiΔwi2+12∑i≠jhijΔ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[H−1]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},:} 求海森矩阵：

∂E∂W^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(k⋅dcol2)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} 替换成 wq−quant⁡(wq)w_{q}-\operatorname{quant}(w_{q})，就能得到一般量化情况下的权重更新公式：

Δw=−wq−quant⁡(wq)[H−1]qqH−1⋅eq and L=12(wq−quant⁡(wq))2[H−1]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(drow⋅dcol3)O(d_{row} \cdot d_{col}^{3})。

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

GPTQ

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

GPTQ 的创新点有：

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

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

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

w←w−H:,p−11[H−1]pp⋅(wp−q(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 不减少实际的计算量，但它能有效解决吞吐的瓶颈问题。