文章目录
- 摘要
- 1 引言
- 2 相关工作
- 2.1 量化
- 2.2 大型语言模型的量化
- 3 Z-FOLD
- 3.1 新引入的参数 ζ
- 3.2 参数整合(ζ 折叠)
- 3.3 使用校准集的微调
- 4 实验
- 4.1 实验设置
- 4.2 与其他方法的比较
- 4.3 Z-FOLD 的泛化能力
- 4.4 Z-FOLD 的可移植性
- 4.5 消融研究
- 5 结论
- 6 限制
摘要
随着参数数量不断增加以提升性能,超大规模人工智能模型(包括大型语言模型)的高效推理变得尤为关键。无论是在移动设备还是云服务器等计算环境中,这一需求都同样存在。量化作为一种解决方案,可以缓解推理过程中的计算负担。通过使用更低位宽的表示方式,量化减少了对DRAM的访问频率,同时通过稠密矩阵格式充分利用计算的并行性。因此,量化模型在解决内存和计算瓶颈的同时,实现了较低的端到端延迟和资源利用的优化。本文提出了一种简单的后训练量化方案,称为Z-FOLD,该方法充分利用了广泛应用于大型语言模型中的Transformer结构的特性。相关代码将发布在:https://github.com/SamsungLabs/Z-Fold。
1 引言
Transformer(Vaswani 等,2017)模型通过利用注意力机制(Bahdanau 等,2014),在多个领域革新了机器学习技术。神经语言处理(NLP)任务,包括大型语言模型(LLMs)和神经机器翻译(NMT),在很大程度上依赖于Transformer架构,其影响力也扩展到了视觉任务(Khan 等,2022;Rombach 等,2022)。本质上,超大规模人工智能模型在很大程度上依赖于Transformer架构。
为了提升通用语言任务的性能,大型语言模型不断增加参数数量,并通过多种技术实现这一目标,如防止性能饱和或使超大规模模型能够被训练;采用层丢弃(Zhang and He,2020)、流水线并行(Huang 等,2019)、模型并行(Shoeybi 等,2019)、预归一化(Xiong 等,2020)以及AdamW优化器(Loshchilov 和 Hutter,2017)等方法,最终实现了参数规模达到上万亿的模型。
然而,这些超大规模模型面临着限制,妨碍了它们的实际应用。主要挑战包括处理庞大参数数量所带来的内存和计算成本,这导致无论计算资源如何,端到端延迟和功耗均显著增加。具体来说,大量参数需要在内存和处理器间频繁传输,并对长序列输入进行计算。因此,尽管参数扩展带来了性能提升,超大规模模型在推理时对硬件的负担依然沉重。由此,高效推理技术在资源受限(如边缘设备)和资源丰富(如云服务器)环境中,受到与模型扩展同样重要的关注。
量化作为众多提升模型效率技术(如剪枝(Frantar 和 Alistarh,2023)、低秩近似(Chen 等,2018)和知识蒸馏(Hsieh 等,2023))中的一个有前景且吸引人的方案,通过使用低位宽表示模型,减少DRAM访问频率,同时通过稠密矩阵格式充分利用计算并行性。因此,量化模型有效缓解了由内存和计算开销引起的瓶颈。
本文提出了一种针对大型语言模型的后训练量化方法——Z-FOLD,该方法利用Transformer架构(预归一化)的特点,在不增加额外量化参数或计算开销的前提下,提升量化模型性能。简而言之,我们使用更多的参数(ζ)将权重量化到更低的位宽(最低至2位),相较于现有方法进一步减少了量化引起的损失扰动,从而改善量化网络。然而,我们在推理前将这些额外参数(ζ)折叠或融合进其他已有参数(α或γ;见图1),避免了额外的硬件开销。因此,该方法在量化大型语言模型中实现了最先进的性能,同时没有引入任何额外的参数或硬件成本。
2 相关工作
2.1 量化
当整个数据集可用时,量化感知训练(Quantization Aware Training,QAT)(Jin 等,2021)可能更有用,且性能有望优于后训练量化(Post-Training Quantization,PTQ)。然而,考虑到大型语言模型(LLMs)中需要优化的参数数量庞大,使用整个数据集对其重新训练并非现实选择。使用少量样本的QAT也因搜索空间过大而易出现过拟合。因此,大多数量化LLMs的方法都依赖于通过校准集(few-shot)或无数据集(数据无关量化,data-free quantization)的PTQ。通过将搜索空间限制在预训练模型收敛点附近的区域,PTQ 能避免因样本数据少而导致的过拟合。
可以用泰勒展开(Taylor series)近似预训练模型的损失曲面。假设预训练模型已成功收敛并具备良好泛化能力,我们可以简化损失退化( Δ L \Delta L ΔL)与权重扰动 Δ w \Delta w Δw(即 Δ w ≜ v e c ( Δ W ) \Delta w ≜ \mathrm{vec}(\Delta W) Δw≜vec(ΔW),权重展开成向量)之间的关系,得到线性近似(LeCun 等,1989;Nagel 等,2020):
Δ L ≈ 1 2 Δ w ⊤ ⋅ H ( w ) ⋅ Δ w , (1) \Delta L \approx \frac{1}{2} \Delta w^\top \cdot H(w) \cdot \Delta w, \tag{1} ΔL≈21Δw⊤⋅H(w)⋅Δw,(1)
其中 H ( w ) H(w) H(w) 是权重的 Hessian 矩阵。
在没有数据集的情况下进行模型量化时,Hessian 可被近似为 c ⋅ I c \cdot I c⋅I,其中 c c c 是常数, I I I 为单位矩阵。因此,我们为权重 W ∈ R d i n × d o u t W \in \mathbb{R}^{d_{in} \times d_{out}} W∈Rdin×dout 设置通道维度的步长(量化步长) α ∈ R d o u t × 1 \alpha \in \mathbb{R}^{d_{out} \times 1} α∈Rdout×1,以最小化权重扰动,形式化为:
α ∗ = arg min α ∥ W − W i n t ⋅ A ∥ F , (2) \alpha^* = \arg\min_{\alpha} \| W - W_{int} \cdot A \|_F, \tag{2} α∗=argαmin∥W−Wint⋅A∥F,(2)
W i n t = c l i p ( ⌊ W ⋅ A − 1 ⌉ , n , p ) , (3) W_{int} = \mathrm{clip}\left(\left\lfloor W \cdot A^{-1} \right\rceil, n, p\right), \tag{3} Wint=clip(⌊W⋅A−1⌉,n,p),(3)
其中 ⌊ ⋅ ⌉ \lfloor \cdot \rceil ⌊⋅⌉ 表示四舍五入函数, A = d i a g ( α ) A = \mathrm{diag}(\alpha) A=diag(α), n n n 和 p p p 分别是整数权重 W i n t W_{int} Wint 的下界和上界。确定 α ∗ \alpha^* α∗ 后,我们采用最邻近舍入策略将权重量化到均匀分布的量化网格,这种方法称为最小均方误差量化(MMSE quantization)。需要注意的是,步长 α \alpha α 也可以简单地依据极值范围采用最小-最大量化(min-max quantization)计算:
α = max ( W ) − min ( W ) 2 b i t − 1 , (4) \alpha = \frac{\max(W) - \min(W)}{2^{bit} - 1}, \tag{4} α=2bit−1max(W)−min(W),(4)
其中 min ( ⋅ ) \min(\cdot) min(⋅) 和 max ( ⋅ ) \max(\cdot) max(⋅) 分别返回每列的最小值和最大值。
基于由 MMSE 或最小-最大方法设定的步长 α \alpha α,我们可以利用校准集(few-shot)近似 Hessian 来进一步优化量化网络,从而使量化网络比无数据量化更鲁棒(Nagel 等,2020;Li 等,2021;Hubara 等,2021;Frantar 和 Alistarh,2022;Jeon 等,2023)。其中,OPTQ(Frantar 等,2023)(也称 GPTQ)通过简化最优大脑量化(Optimal Brain Quantization,Frantar 和 Alistarh,2022),无需梯度优化且用较短时间即实现了量化LLMs的合理性能。
设转置权重矩阵为 W ⊤ ∈ R d o u t × d i n W^\top \in \mathbb{R}^{d_{out} \times d_{in}} W⊤∈Rdout×din,输入批次为 X ∈ R d i n × n X \in \mathbb{R}^{d_{in} \times n} X∈Rdin×n,其中 d i n d_{in} din 和 n n n 分别为输入维度和输入序列长度。假定 W ⊤ W^\top W⊤ 各行之间无交互,OPTQ 近似对应每行的 Hessian 矩阵为:
H ≈ 2 X X ⊤ + λ I ∈ R d i n × d i n . (5) H \approx 2 X X^\top + \lambda I \in \mathbb{R}^{d_{in} \times d_{in}}. \tag{5} H≈2XX⊤+λI∈Rdin×din.(5)
因此,整体 Hessian 为 H ( w ) = I ⊗ H H(w) = I \otimes H H(w)=I⊗H,其中 ⊗ \otimes ⊗ 表示 Kronecker 乘积。OPTQ 按列(从第0列到第 d i n − 1 d_{in}-1 din−1 列)依次量化 W ⊤ W^\top W⊤。在量化第 i i i 列权重 w i ∈ R d o u t × 1 w_i \in \mathbb{R}^{d_{out} \times 1} wi∈Rdout×1 时,计算误差 e e e:
e = ( w i − w i q ) H i i − 1 , (6) e = \frac{(w_i - w_i^q)}{H_{ii}^{-1}}, \tag{6} e=Hii−1(wi−wiq),(6)
其中 H − 1 = C h o l e s k y ( H − 1 ) ⊤ H^{-1} = \mathrm{Cholesky}(H^{-1})^\top H−1=Cholesky(H−1)⊤,随后对剩余未量化列权重 w j w_j wj ( j = i + 1 , … , d i n − 1 j = i+1, \ldots, d_{in}-1 j=i+1,…,din−1) 更新以反映量化扰动,
w j = w j − e ⋅ H i j − 1 , (7) w_j = w_j - e \cdot H_{ij}^{-1}, \tag{7} wj=wj−e⋅Hij−1,(7)
其中 H i j − 1 H_{ij}^{-1} Hij−1 是 Hessian 逆矩阵第 i i i 行 j j j 列的元素。
2.2 大型语言模型的量化
OPTQ 通过使用少量句子样本近似 Hessian 矩阵增强了量化 LLM 的鲁棒性,其中激活保持 FP16 精度,因为激活不是瓶颈。ZeroQuant(Yao 等,2022)采用权重的组级别量化和激活的令牌级别量化,提出了硬件友好的量化方案。尽管 ZeroQuant 相较其它方法拥有更多的缩放参数,但其会根据推理时使用的硬件架构调整缩放参数数量,从而将开销降至最低。LLM.int8()(Dettmers 等,2022)提出了矢量级量化方法并采用混合精度方案,其中异常向量保持 FP16。SmoothQuant(Xiao 等,2022)和 Quadapter(Park 等,2022)探讨了激活动态范围的挑战,并提出相应的缓解方法。QuIP(Chee 等,2023)引入额外参数(随机正交矩阵)尝试将LLMs量化至极低位宽(如2位),但 QuIP 需要在推理时存储或生成这些额外参数并承担计算开销。ZeroQuant-V2(Yao 等,2023)总结并分析了各方法及位宽配置(例如 E8A8、W4A16)的特性。尽管上述方法均采用后训练量化,LLM-QAT(Liu 等,2023)提出了针对 LLM 的 QAT 方法。
3 Z-FOLD
3.1 新引入的参数 ζ
为了进一步减少量化引起的损失扰动(公式 (1)),我们在现有的按输出通道的缩放因子 α ∈ R n × 1 \alpha \in \mathbb{R}^{n \times 1} α∈Rn×1 之外,新增与输入通道对应的缩放因子 ζ ∈ R m × 1 \zeta \in \mathbb{R}^{m \times 1} ζ∈Rm×1。该缩放因子 ζ \zeta ζ 会在推理前融合(fold)到上一层的 α \alpha α 中,以消除因增加参数 ζ \zeta ζ 带来的开销,我们称之为 Z-FOLD(见图 2)。本质上,我们的目标是获得双向(输入和输出)步长矩阵 S S S,可表述为:
S ∗ = arg min S Δ w ⊤ ⋅ H ( w ) ⋅ Δ w , (8) S^* = \arg\min_S \Delta w^\top \cdot H(w) \cdot \Delta w, \tag{8} S∗=argSminΔw⊤⋅H(w)⋅Δw,(8)
Δ w = v e c ( W − S ⊙ W i n t ) , (9) \Delta w = \mathrm{vec}(W - S \odot W_{int}), \tag{9} Δw=vec(W−S⊙Wint),(9)
W i n t = c l i p ( ⌊ W ⊘ S ⌉ , n , p ) , (10) W_{int} = \mathrm{clip}\left(\left\lfloor W \oslash S \right\rceil, n, p \right), \tag{10} Wint=clip(⌊W⊘S⌉,n,p),(10)
其中 ⊘ \oslash ⊘ 和 ⊙ \odot ⊙ 分别表示哈达玛除法(Hadamard division)和哈达玛乘积(Hadamard product)。
为求解公式 (8) 中的 S ∗ S^* S∗,我们将矩阵 S S S 分解为两个向量的乘积(秩-1 近似),采用交替最小二乘法(Alternating Least Squares, ALS)分解(Zachariah 等,2012)。具体地,我们确定 α ∈ R n × 1 \alpha \in \mathbb{R}^{n \times 1} α∈Rn×1 和 ζ ∈ R m × 1 \zeta \in \mathbb{R}^{m \times 1} ζ∈Rm×1,使得:
Δ W = ∥ W − Z ⋅ W i n t ⋅ A ∥ 2 2 , (11) \Delta W = \| W - Z \cdot W_{int} \cdot A \|_2^2, \tag{11} ΔW=∥W−Z⋅Wint⋅A∥22,(11)
其中
Z ≡ d i a g ( ζ ) , A ≡ d i a g ( α ) , (12) Z \equiv \mathrm{diag}(\zeta), \quad A \equiv \mathrm{diag}(\alpha), \tag{12} Z≡diag(ζ),A≡diag(α),(12)
该问题可通过向量级的最小二乘法求解。
考虑 W ∈ R m × n W \in \mathbb{R}^{m \times n} W∈Rm×n 的第 i i i 列向量 w i ∈ R m × 1 w_i \in \mathbb{R}^{m \times 1} wi∈Rm×1,可用加权最小二乘法利用 Hessian 矩阵求解 α i \alpha_i αi:
α i = ( w ~ i ⊤ H w ~ i ) − 1 w ~ i ⊤ H w i , (13) \alpha_i = (\tilde{w}_i^\top H \tilde{w}_i)^{-1} \tilde{w}_i^\top H w_i, \tag{13} αi=(w~i⊤Hw~i)−1w~i⊤Hwi,(13)
其中 w ~ i = d i a g ( ζ ) ⋅ w i n t , i \tilde{w}_i = \mathrm{diag}(\zeta) \cdot w_{int,i} w~i=diag(ζ)⋅wint,i。
加权最小二乘法仅用于求解 α \alpha α,并与 ζ \zeta ζ 交替迭代优化。需要注意的是,由于共享 ζ \zeta ζ 的权重也共享 Hessian 矩阵的条目,因此是否使用 Hessian 来确定 ζ \zeta ζ 不影响最终结果。
算法 1 详细描述了该分解算法。我们以全1向量初始化 ζ \zeta ζ,然后按公式 (4) 使用最小-最大量化初始化 α \alpha α(第2行)。随后, α \alpha α 和 ζ \zeta ζ 交替使用最小二乘解更新,直到收敛(第3–8行),以最小化公式 (8) 中的损失。值得注意的是,按列的最小二乘过程(第12–14行)支持循环级并行处理,可以通过批量矩阵乘法(bmm)并行加速。
从权重矩阵 W W W 分解得步长矩阵后,量化权重矩阵 W q W^q Wq 计算如下:
W q = S ⊙ W i n t , (14) W^q = S \odot W_{int}, \tag{14} Wq=S⊙Wint,(14)
= d i a g ( ζ ) ⋅ W i n t ⋅ d i a g ( α ) , (15) = \mathrm{diag}(\zeta) \cdot W_{int} \cdot \mathrm{diag}(\alpha), \tag{15} =diag(ζ)⋅Wint⋅diag(α),(15)
= Z ⋅ W i n t ⋅ A . (16) = Z \cdot W_{int} \cdot A. \tag{16} =Z⋅Wint⋅A.(16)
该方案还支持基于输出通道的非对称量化,其中每列权重 w i w_i wi 量化为:
w i q = α ( w i , i n t − o ) ⊙ ζ , (17) w_i^q = \alpha (w_{i,int} - o) \odot \zeta, \tag{17} wiq=α(wi,int−o)⊙ζ,(17)
w i , i n t = c l i p ( ⌊ w i ζ ⊘ α + o ⌉ , n , p ) , (18) w_{i,int} = \mathrm{clip} \left( \left\lfloor \frac{w_i}{\zeta} \oslash \alpha + o \right\rceil, n, p \right), \tag{18} wi,int=clip(⌊ζwi⊘α+o⌉,n,p),(18)
其中 o o o 是零点向量(全零向量),通过
o = − min ( w i ζ ⊘ α ) o = - \min \left( \frac{w_i}{\zeta} \oslash \alpha \right) o=−min(ζwi⊘α)
计算得到。对于非对称量化,我们可以用网格搜索寻找最佳 α \alpha α 和 o o o,以最小化量化带来的损失扰动。该搜索替代了算法 1 中第7行的最小二乘解。
3.2 参数整合(ζ 折叠)
图 2、3 和 5 展示了 Z-FOLD 的机制。图中说明了输入通道维度的参数 ζ \zeta ζ 如何被整合到上一层的仿射变换参数 γ \gamma γ 和 β \beta β(归一化层到线性层,Norm-to-Linear)或上一层的输出通道维度缩放参数 α \alpha α(线性层到线性层,Linear-to-Linear)。
Norm-to-Linear
考虑一个作为归一化层输入的潜在令牌向量 x ∈ R m × 1 x \in \mathbb{R}^{m \times 1} x∈Rm×1,归一化后向量 x ^ \hat{x} x^ 计算为:
y = γ ⊙ x ^ + β , (19) y = \gamma \odot \hat{x} + \beta, \tag{19} y=γ⊙x^+β,(19)
其中 γ ∈ R m × 1 \gamma \in \mathbb{R}^{m \times 1} γ∈Rm×1 和 β ∈ R m × 1 \beta \in \mathbb{R}^{m \times 1} β∈Rm×1 是仿射变换参数(缩放和偏移)。
公式 (19) 可写成矩阵乘法形式:
y ⊤ = x ^ ⊤ ⋅ 1 ⋅ Γ , (20) y^\top = \hat{x}^\top \cdot \mathbf{1} \cdot \Gamma, \tag{20} y⊤=x^⊤⋅1⋅Γ,(20)
其中
Γ = [ d i a g ( γ ) β ⊤ ] ≡ [ γ 0 0 ⋯ 0 0 γ 1 ⋯ 0 ⋮ ⋮ ⋱ ⋮ 0 0 ⋯ γ m − 1 β 0 β 1 ⋯ β m − 1 ] . (21) \Gamma = \begin{bmatrix} \mathrm{diag}(\gamma) \\ \beta^\top \end{bmatrix} \equiv \begin{bmatrix} \gamma_0 & 0 & \cdots & 0 \\ 0 & \gamma_1 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \gamma_{m-1} \\ \beta_0 & \beta_1 & \cdots & \beta_{m-1} \end{bmatrix}. \tag{21} Γ=[diag(γ)β⊤]≡ γ00⋮0β00γ1⋮0β1⋯⋯⋱⋯⋯00⋮γm−1βm−1 .(21)
归一化层输出 y y y 随后与后续线性层的量化权重矩阵 W q ∈ R m × n W^q \in \mathbb{R}^{m \times n} Wq∈Rm×n 相乘:
y ⊤ ⋅ W q = y ⊤ ⋅ ( Z ⋅ W i n t ⋅ A ) , (22) y^\top \cdot W^q = y^\top \cdot (Z \cdot W_{int} \cdot A), \tag{22} y⊤⋅Wq=y⊤⋅(Z⋅Wint⋅A),(22)
= x ^ ⊤ ⋅ 1 ⋅ Γ ⋅ ( Z ⋅ W i n t ⋅ A ) , (23) = \hat{x}^\top \cdot \mathbf{1} \cdot \Gamma \cdot (Z \cdot W_{int} \cdot A), \tag{23} =x^⊤⋅1⋅Γ⋅(Z⋅Wint⋅A),(23)
= x ^ ⊤ ⋅ 1 ⋅ ( Γ ⋅ Z ) ⋅ ( W i n t ⋅ A ) . (24) = \hat{x}^\top \cdot \mathbf{1} \cdot (\Gamma \cdot Z) \cdot (W_{int} \cdot A). \tag{24} =x^⊤⋅1⋅(Γ⋅Z)⋅(Wint⋅A).(24)
正如公式 (24) 所示,我们可以用下一层的 ζ \zeta ζ 更新仿射变换参数:
Γ ′ = Γ ⋅ Z = [ γ 0 ζ 0 0 ⋯ 0 0 γ 1 ζ 1 ⋯ 0 ⋮ ⋮ ⋱ ⋮ 0 0 ⋯ γ m − 1 ζ m − 1 β 0 ζ 0 β 1 ζ 1 ⋯ β m − 1 ζ m − 1 ] . (25) \Gamma' = \Gamma \cdot Z = \begin{bmatrix} \gamma_0 \zeta_0 & 0 & \cdots & 0 \\ 0 & \gamma_1 \zeta_1 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \gamma_{m-1} \zeta_{m-1} \\ \beta_0 \zeta_0 & \beta_1 \zeta_1 & \cdots & \beta_{m-1} \zeta_{m-1} \end{bmatrix}. \tag{25} Γ′=Γ⋅Z= γ0ζ00⋮0β0ζ00γ1ζ1⋮0β1ζ1⋯⋯⋱⋯⋯00⋮γm−1ζm−1βm−1ζm−1 .(25)
本质上,我们无需依赖数据集,即可更新仿射变换参数以补偿量化带来的误差。通过将 ζ \zeta ζ 折叠进前一归一化层的 γ \gamma γ 和 β \beta β,我们在不增加参数和计算成本的情况下,减少了量化损失。
值得注意的是,通过拼接查询(Q)、键(K)和值(V)的权重矩阵,可以提取共享的 ζ \zeta ζ,从而无缝整合进前一归一化层的 γ \gamma γ 和 β \beta β(参见图 4)。
Linear-to-Linear
在不含除 ReLU 之外的非线性函数的相邻两线性层间,输入通道级别的步长 ζ \zeta ζ 可被融合进上一层的输出通道级别步长 α \alpha α 中。设两连续线性层的权重分别为 W 1 q ∈ R n × m W_1^q \in \mathbb{R}^{n \times m} W1q∈Rn×m 和 W 2 q ∈ R m × n W_2^q \in \mathbb{R}^{m \times n} W2q∈Rm×n,第一层输入为 x ∈ R n × 1 x \in \mathbb{R}^{n \times 1} x∈Rn×1,则输出计算为:
( x ⊤ ) ⋅ ( W 1 q A 1 ) ⋅ ( Z 2 W 2 q A 2 ) , (26) (x^\top) \cdot (W_1^q A_1) \cdot (Z_2 W_2^q A_2), \tag{26} (x⊤)⋅(W1qA1)⋅(Z2W2qA2),(26)
= ( x ⊤ ) ⋅ ( W 1 q A 1 Z 2 ) ⋅ ( W 2 q A 2 ) . (27) = (x^\top) \cdot (W_1^q A_1 Z_2) \cdot (W_2^q A_2). \tag{27} =(x⊤)⋅(W1qA1Z2)⋅(W2qA2).(27)
如图 1 所示,除了 W 1 W_1 W1 和 W 2 W_2 W2,这一方式也适用于融合注意力层中的权重矩阵 W O W_O WO 和 W V W_V WV。
设 Y Y Y 和 X X X 分别表示 Y = s o f t m a x ( Q K ⊤ ) Y = \mathrm{softmax}(Q K^\top) Y=softmax(QK⊤) 和注意力层输入,则线性层 W O W_O WO 的输出在融合参数后保持一致:
( Y ⊤ ) ⋅ ( X ⊤ ⋅ W V q A V ) ⋅ ( Z O W O q A O ) , (28) (Y^\top) \cdot (X^\top \cdot W_V^q A_V) \cdot (Z_O W_O^q A_O), \tag{28} (Y⊤)⋅(X⊤⋅WVqAV)⋅(ZOWOqAO),(28)
= ( Y ⊤ ) ⋅ ( X ⊤ ⋅ W V q A V Z O ) ⋅ ( W O q A O ) . (29) = (Y^\top) \cdot (X^\top \cdot W_V^q A_V Z_O) \cdot (W_O^q A_O). \tag{29} =(Y⊤)⋅(X⊤⋅WVqAVZO)⋅(WOqAO).(29)
3.3 使用校准集的微调
在算法 1 得到双向步长矩阵 S S S 后,我们利用小规模校准集通过近似 Hessian 的优化(Nagel 等,2020;Li 等,2021;Hubara 等,2021;Wei 等,2022;Frantar 等,2023;Jeon 等,2023)进一步优化量化后的 LLM。我们采用 OPTQ,它是适用于超大规模模型(如 LLMs)的高效方案。具体地,我们以 Z-FOLD 作为初始量化状态,再用少量样本进行细调。算法 2 详细说明了该流程。
我们先用 Hessian 优化 α \alpha α,最小化损失退化(公式 (1))(第2–3行),然后再用 OPTQ 进一步细化量化模型(第4–9行)。
本方案不仅适用于均匀量化,也支持多级二值量化(Xu 等,2018;Jeon 等,2022)。在推理时不量化激活的前提下,两种格式均可由专用核(Jeon 等,2020;Frantar 等,2023)加速执行。Z-FOLD 同样可应用于 QLoRA(Dettmers 等,2023)量化场景。
4 实验
4.1 实验设置
为了评估所提方法的性能,我们使用 Z-FOLD 对公开可用的大型语言模型(如 OPT(Zhang 等,2022)、LLaMA(Touvron 等,2023)和 BLOOM(Scao 等,2022))进行了量化。作为估计 Hessian 矩阵的校准数据集,我们采用了来自 C4 数据集(Raffel 等,2019)中随机抽取的128段2048个token的片段,和 OPTQ(Frantar 等,2023)中使用的校准集一致。所有实验均在单块 NVIDIA A100 GPU(80GB)上进行,并使用 PyTorch 实现 Z-FOLD。
在实验中,我们只量化权重,保持激活为全精度(FP16),因为激活在 LLM 推理中并非显著瓶颈(Frantar 等,2023)。当 Z-FOLD 中 α \alpha α 和 ζ \zeta ζ 的交替更新导致损失扰动 Δ L \Delta L ΔL(见公式 (1))增加,或者总更新次数超过30次时,我们会中断交替更新(见算法1第4-7行)。
量化完成后,我们使用三个基准数据集(WikiText-2(Merity 等,2016)、C4(Raffel 等,2019)和 PTB(Marcus 等,1993))以及若干零样本任务,评估量化后的 LLM 性能。
4.2 与其他方法的比较
表1总结了不同规模(125M至30B)OPT模型的量化结果。为对比,我们也列出了传统 RTN 和 OPTQ 方案的结果。总体来看,Z-FOLD 和 OPTQ 由于利用估计的 Hessian 矩阵以最小化损失退化 Δ L \Delta L ΔL,在性能上远超仅考虑权重量化误差 Δ w \Delta w Δw 的 RTN。此外,Z-FOLD 凭借引入额外参数 ζ \zeta ζ 对量化水平进行更精细划分,在低位宽量化时,相较于 OPTQ 性能优势更加明显。
例如,在2位量化情况下,OPTQ 完全失效(无论模型大小及测试集,困惑度均超过2000),而 Z-FOLD 仍表现出合理的困惑度。Z-FOLD 优异的性能源于引入额外参数 ζ \zeta ζ 精细划分量化水平,但通过折叠技术(见第3.2节)避免了额外开销。类似优越表现也可见于 LLaMA(见表2)和 BLOOM 模型(详见附录 A.1 的表5)。
Z-FOLD 在几乎所有实验中均优于 OPTQ,低位宽量化时优势尤为显著。值得注意的是,OPTQ 量化后模型性能显著下降,例如 2 位 LLaMA-30B 模型困惑度高达 2065,而对应的全精度模型仅为 4.10。相比之下,Z-FOLD 量化模型困惑度为 9.65,表现非常接近全精度模型。
4.3 Z-FOLD 的泛化能力
我们通过零样本任务测试 Z-FOLD 量化的 LLM 泛化能力,任务包括 LAMBADA(Paperno 等,2016)、PIQA(Bisk 等,2020)、StoryCloze(Mostafazadeh 等,2017)和 ARC(Easy 和 Challenge)(Clark 等,2018)。需要说明的是,我们用于量化的校准数据(C4数据集随机token片段)来自随机爬取的网页,非特定任务数据,实验环境保持零样本设置。
附录 A.2(表6至表8)汇总了量化 LLM 的零样本性能。结果显示,Z-FOLD 量化模型比传统 RTN 和 OPTQ 更加鲁棒。特别是在 LAMBADA 任务中,RTN 和 OPTQ 在2位量化时即使是30B模型也完全失败,而 Z-FOLD 由于引入了参数 ζ \zeta ζ,表现出合理的性能。
4.4 Z-FOLD 的可移植性
至今我们主要将 Z-FOLD 与 OPTQ 联合使用量化 LLM。为了展示 Z-FOLD 的通用性,我们将其与其他知名量化方法如 AdaRound(Nagel 等,2020)和 BRECQ(Li 等,2021)结合(见表3)。由于 AdaRound 和 BRECQ 对大规模模型的量化计算资源和时间需求高,我们只考虑了 OPT-125M 模型。
表3显示,Z-FOLD 大幅提升了 AdaRound 和 BRECQ 的量化性能。尤其是与 BRECQ 结合时,Z-FOLD 量化模型性能接近全精度模型。
4.5 消融研究
我们回顾 Z-FOLD 的两大核心特性:1)额外的(可折叠)参数 ζ \zeta ζ,2)基于 Hessian 的量化网格计算(即 ζ \zeta ζ 和 α \alpha α)。需要注意的是,Z-FOLD 在确定步长时利用了 Hessian 信息,而现有方法如 AdaRound 和 OPTQ 仅使用 MMSE 或最小-最大方法。
通过消融实验(见表4),我们发现引入参数 ζ \zeta ζ 能有效提升量化性能,说明 ζ \zeta ζ 对细化量化网格的作用显著。同时,利用 Hessian 计算 ζ \zeta ζ 和 α \alpha α 可大幅提升性能,即借助 Hessian 优化量化网格,有助于减少损失扰动 Δ L \Delta L ΔL,而非仅减少量化误差 Δ W \Delta W ΔW。
5 结论
本文提出了一种基于后训练的量化方案 Z-FOLD。该方法充分利用 Transformer(预归一化结构和线性层的连续性)的独特架构特性,发展出一种简单且高效的量化方法,显著提升了量化模型的性能。相比现有方法,Z-FOLD 使用了更多参数,实现了更精细的量化,同时通过折叠技术将额外参数智能融合进已有参数,避免了额外的硬件成本。实验结果表明,Z-FOLD 在量化大型语言模型上达到了最先进的性能水平。
6 限制
尽管我们针对的预归一化 Transformer 架构(如 OPT、BLOOM 和 LLaMA)在生成型大型语言模型中应用广泛,但所提方案对目标架构存在依赖性。