专栏名称: 吃果冻不吐果冻皮
专注于AI工程化(LLM、MLOps、LLMOps、RAG、Agent)落地。
目录
相关文章推荐
首席商业评论  ·  头最铁的年轻人,买房只买顶楼 ·  昨天  
神嘛事儿  ·  我回答了 @用户7699215969 ... ·  2 天前  
第一财经  ·  阿里,跳空上涨! ·  3 天前  
魔法K线一王焕昌  ·  【休息一下也无妨】 ·  3 天前  
李楠或kkk  ·  我就不明白,同样价格,想要 c ... ·  3 天前  
51好读  ›  专栏  ›  吃果冻不吐果冻皮

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

吃果冻不吐果冻皮  · 公众号  ·  · 2024-11-12 09:13

正文

近年来,随着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( )移除缩放操作并将其融入到后续的权重矩阵中。

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

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

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

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

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

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

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

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

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

其中, 是通过对键和值进行softmax运算得到的序列长度大小的方阵, 是关于一个头的Value矩阵。

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

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

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

由于已经将 应用于 ,因此仍然需要应用 ,这导致 的完整变换变成了 ,并且插入一个块到前向传播中计算 ,其中,Z是注意力激活。

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







请到「今天看啥」查看全文