大模型量化技术原理:Atom、QuaRot

726 阅读24分钟

本文为稀土掘金技术社区首发签约文章,30天内禁止转载,30天后未获授权禁止转载,侵权必究!

近年来,随着Transformer、MOE架构的提出,使得深度学习模型轻松突破上万亿规模参数,从而导致模型变得越来越大,因此,我们需要一些大模型压缩技术来降低模型部署的成本,并提升模型的推理性能。 模型压缩主要分为如下几类:

  • 剪枝(Pruning)
  • 知识蒸馏(Knowledge Distillation)
  • 量化(Quantization)

本系列将针对一些常见大模型量化方案(GPTQ、LLM.int8()、SmoothQuant、AWQ等)进行讲述。

之前的文章主要集中在W4A16、W8A16、W8A8等量化方法上面。本文将介绍两篇针对W4A4KV4量化具有代表性的工作Atom、QuaRot。

Atom

背景

LLM 在内容生成、聊天机器人和情感分析等应用中的需求不断增长,给LLM服务提供商带来了显著的挑战。

为了有效利用GPU资源并提高吞吐量,批处理多个请求已成为一种流行的范式;为了进一步加速批处理,LLM量化技术通过减少内存消耗和增加计算能力来实现。现有的量化方案无法充分利用现代GPU的能力,如:4位整数运算,导致性能不是最优。

为了最大化LLM的服务吞吐量,作者提出了一种低比特量化方法Atom,通过使用低比特算子显著提高服务吞吐量,并通过低比特量化大幅减少内存消耗,同时通过应用混合精度和细粒度量化来减少模型的精度损失。对4位权重激活量化(W4A4)进行评估,Atom在保持相同延迟目标的同时,将端到端吞吐量(每秒Token生成数)提高了7.73倍(与FP16相比)和2.53倍(与INT8量化相比)。

LLM服务低比特量化性能分析

由于需求高,通常LLM服务以吞吐量为导向。然而,LLM推理的自回归解码阶段仅以一个Token作为输入并生成下一个Token,因此依赖于矩阵-向量乘法(GEMV)。由于GEMV需要加载一个大型的权重矩阵,同时只执行少量的乘法,因此它严重受到内存限制。这导致GPU利用率低下,从而导致计算强度(computation-to-IO ratio)和吞吐量底。为了缓解这个问题,批处理被广泛使用,通过结合多个请求的输入来执行稠密层(K,Q,V生成、O投影和MLP)的矩阵乘法,以增加计算强度,从而提高GPU利用率。

为了进一步利用批处理并提高吞吐量,解码和预填充阶段的稠密层的输入矩阵被批处理在一起形成更大的矩阵。由于大批量大小,稠密层为计算受限的矩阵-矩阵乘法(GEMM)。尽管解码阶段的自注意力层也是GEMV操作,但它们无法从批处理中受益。由于不同的推理请求无法共享上下文历史的KV缓存,跨请求的数据无法批处理重用,因此没有效率优势。即使有FlashAttention或GQA等优化,自注意力层仍然受到 KV Cache 内存移动较大的限制。

应用批处理技术后,通过对不同批处理大小下不同算子的运行时间进行分解。稠密层和自注意力层都是吞吐量的瓶颈,消耗了超过90%的处理时间。因此,采用量化机制来加速稠密层和自注意力层。

使用Roofline模型来评估不同量化方法在服务场景的效果。权重-激活量化由于高效的低比特硬件算术,稠密层具有更高的计算吞吐量。它还通过减小KV缓存的大小,从而减少内存移动,来增加自注意力层的吞吐量。然而仅对权重量化无法提高稠密层吞吐量,因为在矩阵乘法之前必须执行反量化,导致计算仍然以浮点格式进行。

下图通过Kernel进一步分析了不同量化技术的效果。

综上,低比特权重-激活量化在提高服务场景的吞吐量方面优于仅对权重量化,因为它加速了稠密层和自注意力层。

技术原理

要在低比特表示中保持高准确性是具有挑战的,为了在保持准确性的同时将LLMs量化到极低比特精度,需结合一系列针对LLM特性量身定制的量化机制。这些机制包括混合精度量化与通道重排序、细粒度分组量化和动态量化。

一、混合精度量化与通道重排序

由于激活中存在异常值现象,因此,在低比特量化中有效处理异常值至关重要。

一个直观的方法是将异常值和正常值分开量化,分别使用低比特(移除异常值后,剩余的通道更加均匀,可以用低比特值有效表示)和高比特,这被称为混合精度方法。

本文结果表明,8位(如:FP8和INT8)足以表示异常值(如表3所示)。由于INT8得到硬件实现的广泛支持(例如,NVIDIA Tensor Core),Atom对异常值应用INT8量化

混合精度量化的主要问题是其不规则的内存访问,这导致硬件效率低下。为了保持规则的内存访问的同时应用混合精度量化,Atom 重新利用了RPTQ中引入的重排序技术,其目标是提高量化准确性

如图7所示,Atom将激活中分散的异常值通道重排序到矩阵的末端,从而实现了混合精度的有效实现。为确保计算结果的等效性,权重矩阵需要用激活的相应重排序索引进行重排序。由于异常值通道可以通过使用校准数据进行离线识别,权重矩阵的重排序仅产生一次性成本。然而,激活矩阵的重排序仍然需要在线执行,这可能代价高昂。为了缓解这一点,Atom将激活矩阵重排序操作融合到先前的操作中,显著减少了重排序开销,使重排序运行时间小于0.5%。

Atom动态地重新排序激活(A),将离群通道移动到矩阵的末端,重新排序索引在离线校准中确定。权重矩阵(W)被静态地重新排序,以保持与相应的激活通道对齐,从而保证输出结果的正确性。

二、细粒度分组量化

即使Atom将异常值和正常值分开量化,仍然难以准确执行,因为4位精度的表示能力有限。

为了进一步提高准确性,广泛采用分组量化,即将矩阵分成子组并在每个子组内独立进行量化。例如,组大小为128意味着每连续128个元素被视为一个单独的组进行独立量化。分组量化在准确性提升和反量化开销之间提供了权衡,特别是在权重-激活量化中。

Atom 提出了一种融合技术,如图8所示,这有助于实现具有实际加速的 GEMM Kernel。

  • 首先使用高效的低比特硬件(即Tensor Cores)来计算激活组与相应权重组的矩阵乘法,并获取临时结果(步骤1)。
  • 然后,Atom将多个临时结果相加得到GEMM结果。然而,由于Atom对每个激活和权重组进行了细粒度量化,每个临时结果具有不同的量化参数。因此,Atom首先使用 CUDA Cores 将所有临时结果反量化为FP16表示(步骤2)。
  • 然后执行加法(步骤3)。

为了管理开销,将反量化和求和操作融合到GEMM Kernel中,具体来说,融合到 MMA 流水线中。因此,额外的操作可以在没有额外内存移动的情况下就地执行,并与原始 MMA 指令重叠。

使用128的组大小和128的通道大小,Atom 在Llama-7b上实现了4.251的有效比特。有效比特定义为每个元素使用的平均比特数,包括量化参数。这个指标在以前的权重量化工作中广泛使用,主要是因为它代表了实际的压缩比,因此,在内存受限场景中可以提高速度。然而,权重-激活量化在服务场景中的主要好处是利用低比特算术单元的计算效率,而不是减少内存

三、动态量化

尽管细粒度量化可以更好地保留每个通道内激活的局部变化,但如果基于校准数据静态计算量化参数,这一优势将消失,因为实际输入可能具有不同的局部分布。因此,Atom采用动态量化,为每次推理中的每个激活矩阵定制量化参数。

为了控制动态量化的开销,作者像ZeroQuant的实现一样,将量化操作融合到先前的操作中。由于额外的操作是逐元素(element-wise)的(包括一个reduction和一个逐元素 division),融合操作的运行时间与耗时的稠密层和自注意力层相比,仍然可以忽略不计。

然而,非对称量化可能会导致显著的运行时开销,因为需要大量的额外计算。为了在吞吐量和准确性之间取得平衡,作者采用了具有精心选择的裁剪阈值进行对称量化。同时,还在量化权重矩阵时引入了GPTQ,因为这纯粹是一个离线过程,可以在不牺牲运行时效率的情况下提升准确性。

四、KV Cache量化

解码阶段的自注意力层在内存受限下运行。为了缓解这个问题,Atom还对KV缓存应用了低比特量化。Atom以低比特精度加载KV缓存,并在执行FP16计算之前直接将其反量化,通过减少内存来显著提高吞吐量。另一方面,由于非对称和对称量化的KV缓存的内存移动相似,因此它们在内存受限的自注意力层上表现相似。因此,Atom对KV缓存使用非对称量化,因为它可以提供准确性优势。

与激活矩阵相比,作者认为KV缓存更适合量化。要执行自注意力,传入Token的Query向量与Key缓存相乘,结果使用Softmax进行归一化,然后与Value缓存相乘以获得输出。由于Softmax的归一化,Key缓存的量化误差对输出的影响较小。此外,在图9中的分析表明,Value缓存出现的异常值现象较少,使其更适合量化。因此,Atom直接对KV缓存以注意力头的粒度应用非对称低比特量化,以保持高精度。

为了验证该设计方案的可行性,在Llama模型上实现了Atom,并进行了多种优化,如图6所示。

  • 为了充分利用量化的好处,Atom通过Kernel融合管理额外操作的开销:Atom融合了包括重排序、量化和反量化在内的量化操作到现有操作中。
  • 对于计算受限的稠密层,Atom利用低比特来提高吞吐量。
  • 对于内存限制的自注意力层,Atom将反量化与LLM服务的Kernel库FlashInfer融合,以便只加载来自KV缓存的低比特值。
  • Atom还采用了PageAttention进行高效的内存使用,以支持大批量大小。

实验细节

作者通过在流行数据集上的实验表明,Atom在将模型量化到4位(对于权重和激活)时,精度损失可以忽略不计(在Llama-65B上,平均零样本(zero-shot )准确率下降1.4%,针对WikiText2困惑度增加0.3),而先前的工作在相同精度下导致更大的精度损失。

将端到端服务吞吐量与不同的精度和量化方案进行比较时,Atom在保持类似延迟的同时,将吞吐量提高了7.7倍、5.5倍和2.5倍(分别与FP16、W4A16和W8A8相比)。这些结果表明,Atom可以准确地将LLMs量化到低比特,同时实现高服务吞吐量。

通过消融实验,作者对比了在Atom中使用的不同的量化技术在准确性方面的情况。

首先使用RTN并采用标准的量化方法,即对权重进行逐通道量化,对激活进行逐Token量化,将模型量化到W4A4。然后,应用Atom中的其他量化技术,即混合精度、量化异常值、分组量化、裁剪、GPTQ和KV缓存量化,并逐个检查困惑度。

如下表所示,将异常值通道保留在FP16中显著降低了困惑度。进一步将异常值量化到INT8仅导致非常小的0.05困惑度增加,这表明混合精度有效地解决了异常值问题。

此外,细粒度分组量化是另一个使困惑度降低的主要方法。使用裁剪和GPTQ各自降低了0.09的困惑度。最终,量化KV缓存导致轻微的0.12困惑度增加。

同时,作者对比了使用不同融合量化技术后,GEMM Kernel的吞吐量。一个没有任何量化操作的纯INT4 GEMM 将近 980 TOPS。混合精度的融合,保留128通道的计算在INT8 Tensor Cores 中,导致了8%的额外开销,达到了900 TOPS的吞吐量。细粒度分组量化是主要的开销来源,因为它深刻影响了计算流水线。分组反量化的融合将性能降低到770 TOPS。然而,融合的 GEMM Kernel 仍然比INT8的理论峰值限制高出近18%。

此外,为了证明通道重排序的效率,作者还对Atom和基线进行了消融研究。基线是按照先前的工作(LLM.int8())实现的,使用矩阵分解进行混合精度量化。同时,Atom融合了包括重排序和量化在内的量化操作到现有操作中。通过评估从16到256的批处理大小,并测量了一个层归一化和一个GEMM操作的推理延迟。结果表明Atom在每个批处理大小上的性能都比基线高出25%到35%。

QuaRot

背景

LLM 推理需要大量的计算、内存。LLMs的推理有两个主要阶段:在预填充阶段,提示被处理并缓存每一层的KV,然后在解码阶段逐个生成输出Token。预填充阶段通常是计算密集型的,而解码阶段则以内存密集型而闻名。

量化是解决LLMs推理期间内存和计算问题的最重要技术之一。量化的目标是降低权重参数和KV缓存(使内存使用量减少)以及输入(即激活)的精度,并在低精度下计算前向传播。然而,量化激活很困难,因为它们有较大的异常值(如图1所示),这些异常值的数值要大得多,使得激活量化比权重量化更加困难,尤其是针对4位量化的情况。

以前的工作依赖于使用校准集来表征异常特征,并在推理期间将它们保留在更高的精度。而本文通过使用随机Hadamard变换旋转模型输入来解决异常特征问题。作者利用计算不变性的思想,将Hadamard变换融合到权重矩阵中,从而得到一个没有异常特征的等效网络。这使得权重、激活和KV缓存可以量化到4位精度,而精度损失很小。

技术原理

QuaRot由两个阶段组成。

  • 在第一阶段,对模型权重进行运算(以全精度),并在模型的前向传播中插入两个额外的哈达玛(Hadamard)运算(哈达玛矩阵是一种方块矩阵。它的矩阵元素仅包含1或-1。其矩阵中不同的行具备正交性质),有效地抑制了异常值通道。
  • 在第二阶段,使用某种现有方法量化权重,并在前向传播中添加量化操作以实现激活的在线量化。

默认情况下,使用GPTQ量化权重,而激活则使用简单的RTN方案进行在线量化。图4和图5显示了带有QuaRot修改的前向传播更新后的FFN块和Attention块,包括更新后的权重矩阵、插入的块以及权重和激活的位宽。

第一阶段第1a步:权重修改

首先利用计算不变性将每个权重矩阵乘以一个正交矩阵。为了实现这一点,将 LayerNorm 或 RMSNorm 的线性部分融合到相邻的权重矩阵中。图4显示了通过从 RMSNorm(diag(α)\textrm{diag}(\boldsymbol\alpha))移除缩放操作并将其融入到后续的权重矩阵中。

选择一个与模型的隐藏维度大小匹配的随机Hadamard矩阵,并在之前或之后乘以每个权重矩阵。在图4和图5中,这个矩阵被表示为Q。例如,key投影权重矩阵 WkW_k 被修改为

WkQdiag(α)WkW_k \leftarrow Q^\top \textrm{diag}(\boldsymbol\alpha) W_k

同时,其他权重矩阵也进行类似的修改。

出现在块输出侧的矩阵在之后乘以Q。这种权重修改不会影响模型的输出(假设精度足够),根据计算不变性定理。修改后的权重类似于 QUiP# 中使用的修改,减少了权重的互不相干性,虽然QuaRot的修改不需要任何额外的运行时处理。

此外,传递给Transformer块之间的激活矩阵也经过了互不相干性处理,变成了 XXQX \leftarrow X Q。上图1显示了这种处理的结果,可以看到处理后的激活不再包含任何异常值。

第一阶段第1b步:旋转FFN激活。在上述权重修改后,已经通过将许多权重矩阵的一侧乘以Hadamard矩阵来改变,并且改变了Transformer块之间的计算激活。

若需要改善每个块内激活的量化能力,通过插入在线Hadamard操作来实现

首先在feed-forward网络中插入一个Hadamard操作,在下投影(down-projection)矩阵之前。这个操作以全精度执行,并使用 Quip# 中的快速 Kernel 实现。这个操作被隐式地通过将Hadamard矩阵融合到网络的下投影(down-projection)矩阵中来反转:WdownHWdownW_\textrm{down} \leftarrow H W_\textrm{down}。结合全局矩阵Q,这意味着下投影矩阵现在变成了HWdownQH W_\textrm{down}Q(如图4)。

第一阶段第1c步:注意力中 Value 投影。接下来,对每个注意力块应用额外的Hadamard操作。这种修改是部分在线,部分融合到权重矩阵中。

首先,在计算注意力时,WvW_vWoutW_{out}矩阵在每个头内隐式地相乘。注意力计算包括

Y=concat[(P1V1)(PnhVnh)]Wout=h=1HPhXWv(h)Wout(h)Y = \textrm{concat}[(P_1 V_1) \ldots (P_{n_h} V_{n_h})]W_\textrm{out} = \sum_{h=1}^H P_h XW_v^{(h)} W_\textrm{out}^{(h)}

其中, PhP_h 是通过对键和值进行softmax运算得到的序列长度大小的方阵,Vh=XWv(h)V_h=XW_v^{(h)}是关于一个头的Value矩阵。

这提供了一个机会,可以使用与每个头的维度匹配的Hadamard矩阵HdhH_{d_h}WvW_vWoutW_{out}进行额外处理:

Wv(h)Wv(h)Hdh,Wout(h)HdhWout(h)W_v^{(h)} \leftarrow W_v^{(h)}H_{d_h},\qquad\qquad W_\textrm{out}^{(h)} \leftarrow H_{d_h}W_\textrm{out}^{(h)}

将这些修改代入方程 h=1HPhXWv(h)Wout(h)\sum_{h=1}^H P_h XW_v^{(h)} W_\textrm{out}^{(h)},看到注意力的计算结果保持不变。由于每个头的权重在权重表示中是串联的,可以等效地执行单个Kronecker(Kronecker积在张量计算中非常常见,是衔接矩阵计算和张量计算的重要桥梁)结构乘法:

WvWv(IHdh),Wout(IHdh)WoutW_v \leftarrow W_v ( I \otimes H_{d_h}),\qquad \qquad W_\textrm{out} \leftarrow (I \otimes H_{d_h}) W_\textrm{out}

这种变换现在已经应用于权重矩阵进行了逐头(head-wise)处理,并且导致计算激活也被逐头(head-wise)旋转。为了完成对注意力激活的“完整”Hadamard操作,利用恒等式Hnh×dh=(IHdh)(HnhI)H_{n_h \times d_h} = (I \otimes H_{d_h}) (H_{n_h}\otimes I),跨注意力头共享转换。当头数nhn_h和每个头的维度dhd_h都是2的幂时成立。

由于已经将(IHdh)(I \otimes H_{dh})应用于WvW_vWoutW_{out},因此仍然需要应用(HdhI)(H_{dh} \otimes I)WoutW_{out},这导致WoutW_{out}的完整变换变成了WoutHWoutW_{\text{out}} \leftarrow HW_{\text{out}},并且插入一个块到前向传播中计算ZZ(HnhI)Z \leftarrow Z(H_{nh} \otimes I),其中,Z是注意力激活。

这个块在图5中被称为Hadamard heads,可以使用reshape处理Kronecker结构,并且在reshaped数据上进行Walsh-Hadamard变换(Walsh-Hadamard变换是实时的、对称的正交变换。因为该变换只包括加、减,没有任何乘、除运算,所以有快速算法。Walsh-Hadamard变换在图像处理中的主要用途是压缩编码)来有效计算。

第一阶段第1d步:注意力中Key旋转。使用上述方法,可以成功量化Value向量。然而,注意力模块中的Key向量也已知会受到异常值的影响。类似于上述,可以使用Hadamard旋转来缓解这个问题,允许完全量化KV缓存。

首先,注意力分数P1,...,PhP_1, ..., P_h被计算为:

Q:=Pos(XWq)=:concat[Pos(Q1),,Pos(Qnh)]Q := \operatorname{Pos}(XW_q) =: \textrm{concat}[\operatorname{Pos}(Q_1), \ldots, \operatorname{Pos}(Q_{n_h})]
K:=Pos(XWk)=:concat[Pos(K1),,Pos(Knh)]K := \operatorname{Pos}(XW_k) =: \textrm{concat}[\operatorname{Pos}(K_1), \ldots, \operatorname{Pos}(K_{n_h})]
Ph:=Softmax(αPos(Qh)Pos(KhT)M)P_h := \operatorname{Softmax}(\alpha \operatorname{Pos}(Q_h) \operatorname{Pos}(K_h^T) \odot M)

其中,α\alpha是 Softmax 缩放因子,通常设置为[1/dh][1/\sqrt{d_h}],M是注意力掩码(例如因果关系),Pos表示位置嵌入。

以前,位置嵌入只添加到输入的第一层,这种情况下Pos是一个恒等函数。然而,像 RoPE 这样的最新方法直接将位置信息添加到Key和Query向量。

现在可以观察到Q和K之间存在相同的相互作用,就像WvW_vWoutW_{out}之间的作用一样。然而,Pos的存在阻止直接将Hadamard矩阵融合到W_qW\_qW_kW\_k中。

因此,使用在线逐头(head-wise)Hadamard旋转来旋转Query和Key。因此,Query和Key矩阵的计算被改变如下:

QPos(XWq)(IHdh)=concat[Pos(Q1)Hdh,,Pos(Qnh)Hdh]Q \leftarrow \operatorname{Pos}(XW_q)(I \otimes H_{d_h}) = \textrm{concat}[\operatorname{Pos}(Q_1)H_{d_h}, \ldots, \operatorname{Pos}(Q_{n_h})H_{d_h}]
KPos(XWk)(IHdh)=concat[Pos(K1)Hdh,,Pos(Knh)Hdh]K \leftarrow \operatorname{Pos}(XW_k)(I \otimes H_{d_h}) = \textrm{concat}[\operatorname{Pos}(K_1)H_{d_h}, \ldots, \operatorname{Pos}(K_{n_h}) H_{d_h}]

由于Query和Key都经过旋转,最终的注意力分数P1,...,PhP_1, ..., P_h保持不变。上述过程的替代方法是在应用位置编码之前缓存Key。这种方法(称为Pre-RoPE Caching)需要在应用位置编码之前在线应用逆旋转,但消除了旋转Query向量的需要。同时,它也增加了旋转每个query的Key和Value的开销。

鉴于在解码时有一个query向量和许多缓存的key向量,因此,作者决定使用 Post-RoPE Caching。

总的来说,对前向传播的修改,包括插入特殊的Hadamard块以及调整权重(不会改变模型的前向传播)。其效果是块之间的激活乘以Hadamard矩阵,并且使用Hadamard变换在线处理块内的激活,这种方式通过对应的权重矩阵修改来消除影响。现在已准备量化权重和激活。

第二阶段第2a步:权重量化。对网络的权重应用 GPTQ 进行量化。在上述前向传播修改之后,可以应用任何量化方法。比如:使用简单的 RTN 方案代替 GPTQ,但会牺牲一些准确性。

第二阶段第2b步:在线量化操作

权重量化后,准备应用前向传播中的操作来量化激活。

先将RMSNorm(没有缩放)的计算保留在FP32中。使用对称逐token(输入矩阵的行)量化输入。在对称量化期间,通过将每个token的最大绝对值除以7来计算行缩放因子。然后,将每行除以其对应的缩放因子,并将结果四舍五入到最近的整数。

反量化也是通过将 GEMM 的 INT32 输出转换为 FP16 来完成的,乘以行(来自输入 scale)和列(来自权重 scale)的相应缩放因子。

第二阶段第2c步:量化注意力

对于更长的序列和更大的批量大小,注意力显著受到内存限制。

旋转了key和value之后,可以成功地将缓存量化到低位宽。这减少了所需的IO操作数量。

通过保持Query在FP16中,并使用与Flash Attention 类似的在线softmax计算。在从内存中加载一段key和value向量之后,对其进行反量化,并在FP16中计算点积。

模型效果

在语言生成任务上评估QuaRot的准确性,针对LLAMA-2模型而言,通过将使用GPTQ进行权重量化的QuaRot与SmoothQuant(4位)、OmniQuant(4位)、QUIK(其保留了所有层在4位)在WikiText-2上进行困惑度比较,QuaRot的表现优于所有先前的工作,困惑度损失最多为0.63(而LLAMA2-70B模型仅为0.47)。

QuaRot不需要像OmniQuant中那样重新训练的步骤,也不需要像QUIK中那样使用更高精确度的异常特征和非对称量化。

同时,应用组量化与 Atom 进行了比较,在权重和激活上使用了相同数量的组。在这种情况下,QuaRot不需要保留任何更高精确度的特征和相关操作(如:重排序)。QuaRot在7B模型上比Atom提高了0.1困惑度。

在13B模型上,得到了与Atom相同的困惑度。

结果证明了在生成任务中应用Hadamard变换抑制异常值的有效性

通过在六个重要的Zero-Shot任务(PIQA、WinoGrande 、HellaSwag 、LAMBADA(OpenAI)、Arc(Easy和Challenge))上评估QuaRot的准确性。 使用 lm-evaluation-harness 的默认参数进行实验。QuaRot在保持准确性的同时,平均分数下降最多为4.18%(而70B模型下降仅为1.09%)。

性能分析

作者评估QuaRot在预填充阶段的速度提升(首个Token时间),该阶段通常是计算受限的。QuaRot在LLAMA2-7B上实现了高达2.16倍的速度提升。这比线性层单独实现的3.2倍要少,因为大部分注意力计算仍然是在FP16或FP32中完成的,即使较小的缓存大小减少了IO需求。此外,QuaRot的预填充还包括量化KV缓存,这是解码阶段所需的。

使用QuaRot时,除了在预填充期间获得的速度提升外,在解码期间内存使用量也显著减少,高达3.71倍。这种改进来自于量化KV缓存以及量化模型的权重。

为了展示减少的内存使用量,下表展示了在解码一个Token的每个步骤中的峰值内存使用量(在预填充完成后)。在所有情况下,内存使用量显著的减少。这种减少在消费类硬件上尤为重要,例如:在 Nvidia RTX 3090 GPU 上,使用QuaRot可以在序列长度为2048上进行推理。

消融实验

作者通过大量的消融实验评估了QuaRot的不同方面。

QuaRot通过在Hadamard变换期间消除异常特征来提高量化模型的质量。在4位仅权重量化中,使用GPTQ量化时,QuaRot的困惑度提高了最多2.65(越低越好)。此外,将QuaRot应用于所有模型中,在更低精度(2-3位)进行量化,同样提高了质量。

QuaRot将Hadamard变换融合到权重矩阵中以消除异常值。然而,由于LLMs中的计算不变性属性,任何正交矩阵都可以融合到模型中。通过实验发现随机正交矩阵不如随机Hadamard变换好,在LLAMA2-7B上存在高达1.35困惑度差距。然而,随着模型大小的增加,差距减小,在LLAMA2-70B模型上的困惑度差距仅为0.28。

在整个实验中使用FP32在线Hadamard变换。作者通过消融实验对比了在推理期间使用FP16 Hadamard变换的结果(用于down-projection 和 out-projection 层)。在LLAMA2-7B模型上,结果显示WikiText-2的困惑度变化小于0.1,零样本任务的平均准确度变化小于0.6%,作者认为是噪声。

在LLAMA2-13B模型上,不同的Hadamard精度具有相同的困惑度,零样本结果的差异为0.07%。

因此,使用不同的Hadamard精度(FP16/FP32)不会改变模型。

GPTQ是QuaRot中权重量化的默认选择。作者通过消融实验研究了在权重量化中使用RTN的作用。实验表明,使用GPTQ和RTN在权重量化期间保持了FP16模型准确性的99%和98%。在4位模型量化中,QuaRot-RTN和QuaRot-GPTQ之间的差距随着模型大小的增加而减小(LLAMA2-7B上为2.27,LLAMA2-70B上为0.34),这表明GPTQ在较小模型中是更好的选择

通过在激活和权重上应用不同组大小来评估QuaRot的准确性。结果表明准确性和组大小之间存在明确的权衡,较小的组大小提供了更好的准确性(但需要更多的位数来存储每个组的scales,并且需要更复杂的矩阵乘法kernels)。

通过保持模型的其余部分(包括权重和激活)在高精度中,然后应用QuaRot中组(group-wise)异步量化(组大小为128)以及不同精度的Key和Value。结果表明,3位 KV Cache中的模型困惑度退化可以忽略不计(在LLAMA2-70B模型上仅0.07的困惑度差距)。

此外,通过对比3位和4位量化,与Value相比,Key对量化更敏感,对于LLaMA2-7b而言,将key保留在4位和value量化为3位有0.03困惑度损失(而3位Key量化和4位Value量化困惑度损失为0.18)。结果表明,与将key保留在4位并将value量化为2位相比(在LLAMA2-7B模型上困惑度为5.75),使用3位KV Cache的结果在准确性上更好(LLAMA2-7B模型上困惑度为5.68)。

结语

本文介绍了两篇W4A4KV4量化方法Atom和QuaRot。其中,Atom通过一系列方法(如:混合精度、量化异常值、分组量化、裁剪、GPTQ和KV缓存量化等)尽可能减少模型精度损失,尽可能大的提升模型吞吐量。而 QuaRot 则主要通过引入Hadamard变换来消除异常特征,从而提高量化模型的质量,可以使用RTN、GPTQ量化方法。默认采用GPTQ量化方案,其在小模型上是更好的选择。

码字不易,如果觉得我的文章能够能够给您带来帮助,期待您的点赞收藏加关注~~

参考文档: