Atom: Low-Bit Quantization for Efficient and Accurate LLM Serving

发表时间: 2023-10 · arXiv:2310.19102 (MLSys 2024)

原文: https://arxiv.org/abs/2310.19102

作者/机构：Yilong Zhao, Chien-Yu Lin, Kan Zhu, Zihao Ye, Lequn Chen, Size Zheng, Luis Ceze, Arvind Krishnamurthy, Tianqi Chen, Baris Kasikci

速读

一句话结论 Atom 提出了一种低比特权重-激活量化方法，通过混合精度、细粒度分组和动态量化技术，在保持极低精度损失的同时，将大语言模型的端到端服务吞吐量提升了最高 7.73 倍。

要解决什么问题 大语言模型推理服务通常通过批处理（batching）多个用户请求来提高吞吐量，这会摊薄加载权重的开销，使得稠密层（dense layer）的计算从访存密集型转变为计算密集型。现有的仅权重量化（weight-only quantization）方案虽然减少了显存占用，但在计算前仍需将权重反量化为 16 位浮点数（FP16），无法利用现代 GPU 强大的低比特整数计算单元（如 INT4 Tensor Cores），导致在批处理较大时计算能力成为瓶颈。另一方面，现有的权重-激活量化（weight-activation quantization）方案虽然能利用低比特算力，但直接将精度降到 4-bit（如 OmniQuant 和 QLLM）会导致模型困惑度显著飙升，精度损失难以接受；而停留在 8-bit（如 SmoothQuant）又无法榨干更低比特硬件的性能红利。此外，自注意力层（self-attention）在生成阶段无法跨请求共享数据，始终受限于 KV-cache 的显存带宽，仅靠加速计算无法解除这个卡点。因此，如何在维持模型精度的前提下，将权重和激活值同时压缩到极低比特，并真正转化为硬件执行效率，是当前量化部署的核心痛点。

怎么做的 Atom 的核心思路是结合多种高精度量化技术来弥补 4-bit 带来的信息损失，并通过深度的算子融合（kernel fusion）和通道重排（channel reordering）消除这些复杂技术带来的硬件开销。首先，针对激活值中存在极少数数值极大的异常值（outliers）这一导致量化误差的元凶，Atom 采用了混合精度设计。它通过离线校准找出这些异常通道，在运行时动态将它们重排到激活矩阵的末尾，并用 8-bit（INT8）进行量化，而其余正常值则用 4-bit 量化。权重矩阵也会离线进行对应的静态重排以保证计算等价。这种重排既隔离了异常值，又保证了内存访问的规则性。其次，为了进一步降低 4-bit 对正常值的量化误差，Atom 引入了细粒度分组量化（fine-grained group quantization），将每 128 个元素作为一个子组独立量化。为了解决分组带来的频繁反量化开销，Atom 设计了融合的 GEMM（通用矩阵乘法）算子：先用 Tensor Cores 计算低比特的分组乘法，然后在底层流水线内直接用 CUDA Cores 将临时结果反量化为 FP16 并累加，全程不产生额外的显存读写。第三，Atom 放弃了静态校准激活值参数，改为针对每次输入进行动态对称量化，以更好地捕捉局部数据分布。其量化公式定义为：

$$ s = \frac{2 \cdot \max(|X|)}{2^n - 1} \cdot c $$

$$ \bar{X} = \mathrm{clamp}(\lfloor \frac{X}{s} \rceil, -2^{n-1}, 2^{n-1} - 1) $$
其中 $X$ 是输入张量，$n$ 是目标位宽，$c$ 是用于抑制异常值影响的截断因子。动态量化的计算被融合到了前置算子中，开销微乎其微。最后，针对受限于显存带宽的自注意力层，Atom 对 KV-cache 采用了非对称低比特量化，在读取时直接反量化，大幅减少了显存搬运量。

效果如何 实验在 24GB RTX 4090 GPU 上进行，基于 Punica 服务框架搭建，请求数据分布来自 ShareGPT，评估了 7B 到 65B 规模的 Llama 模型。对比基线包括代表 8-bit 权重-激活量化路线的 SmoothQuant，代表 4-bit 权重-激活量化路线的 OmniQuant 和 QLLM，以及代表仅权重量化路线的 W4A16 和未量化的 FP16 基线。在精度方面，Atom 在 4-bit 设置下表现优异，在 Llama-65B 模型上的零样本（zero-shot）任务平均准确率仅比 FP16 下降 1.4%，WikiText2 困惑度仅增加 0.3，而同等条件下的 OmniQuant 和 SmoothQuant 均出现了严重的精度崩塌。在端到端吞吐量（token/s）方面，在相同的延迟目标和显存限制下，Atom 能够支持更大的批处理规模，其吞吐量相比 FP16 提升了最高 7.73 倍，相比 W4A16 提升了 5.5 倍，相比 INT8 提升了 2.53 倍。代价方面，混合精度和分组量化虽然通过算子融合缓解了开销，但相比纯粹的 INT4 矩阵乘法理论上限仍有约 20% 的吞吐量折损；此外，量化预处理阶段需要离线校准，对于 65B 模型需要耗费约 4 小时。

主要贡献

随着大型语言模型（LLMs）在内容生成、智能聊天机器人和情感分析等应用中的需求不断增长，LLM服务提供商面临着巨大的计算和成本挑战。为了高效利用GPU资源并提高吞吐量，将多个请求进行批处理（Batching）已成为主流范式。为了进一步加速批处理，LLM量化技术被广泛用于减少内存消耗并提高计算能力。然而，现有的量化方案（例如8位权重-激活量化）无法充分利用现代GPU的底层硬件能力（如4位整数算子），导致性能次优；而诸如OmniQuant或QLLM等极低位（如4位）量化方法则会带来显著的精度损失。

为了在最大化LLM服务吞吐量的同时保持极高的准确性，本文提出了Atom，这是一种低位权重-激活量化方法。Atom通过使用低位算子显著提升了服务吞吐量，并通过低位量化大幅降低了内存消耗。其核心创新点包括：
1. 采用混合精度和细粒度量化过程，在保持硬件效率的同时最大化准确性。
2. 结合了通道重排（Channel Reordering），以确保混合精度操作中内存访问的规律性。
3. 采用动态激活量化，以最佳方式捕捉每个输入的分布，从而最小化量化误差。
4. 引入了KV-cache量化，大幅减少了内存受限的自注意力层中的数据移动。

在服务场景下针对4位权重-激活量化进行评估，在保持相同延迟目标的情况下，与FP16相比，Atom将端到端吞吐量（token/s）提升了高达$7.73\times$；与INT8量化相比，吞吐量提升了$2.53\times$，且精度损失可以忽略不计。

图1. Atom设计概述。对于激活矩阵，我们动态重排通道以挑出异常值。然后，我们对正常值应用低位分组量化，同时对异常值使用高位精度。对于权重矩阵，量化过程可以静态完成。我们执行融合的GEMM和融合的FlashInfer来提升吞吐量。我们还采用了量化的KV-cache来减少内存移动。
图1. Atom设计概述。对于激活矩阵，我们动态重排通道以挑出异常值。然后，我们对正常值应用低位分组量化，同时对异常值使用高位精度。对于权重矩阵，量化过程可以静态完成。我们执行融合的GEMM和融合的FlashInfer来提升吞吐量。我们还采用了量化的KV-cache来减少内存移动。
图2. Llama模型在不同4位权重-激活量化机制下的WikiText2困惑度。Atom在所有模型大小上都保持了接近FP16基线的困惑度结果。
图2. Llama模型在不同4位权重-激活量化机制下的WikiText2困惑度。Atom在所有模型大小上都保持了接近FP16基线的困惑度结果。

背景知识与设计原则

量化基础与对称量化选择。量化技术使用离散的低位值来近似高精度的浮点数。典型的均匀非对称量化包含计算缩放因子$s$和零点$z$：$s = \frac{\max(X) - \min(X)}{2^n - 1} \cdot c$， $z = \lfloor \frac{-\min(X)}{s} \rceil$，其中$X$是输入张量，$n$是位宽，$c$是裁剪因子。量化后的张量计算为：$\bar{X} = \mathrm{clamp}(\lfloor \frac{X}{s} \rceil + z, 0, 2^n - 1)$。对于LLM，如果对权重和激活都应用非对称量化，矩阵乘法$W \cdot X = s_W(\bar{W} - z_W) \cdot s_x(\bar{X} - z_x)$会产生三个额外的交叉项，导致利用低位算术单元时产生额外计算。因此，为了硬件效率，Atom采用对称量化，其公式简化为：$s = \frac{2 \cdot \max(|X|)}{2^n - 1} \cdot c$，$\bar{X} = \mathrm{clamp}(\lfloor \frac{X}{s} \rceil, -2^{n-1}, 2^{n-1} - 1)$。

低位LLM服务的性能瓶颈分析。LLM推理的自回归解码阶段依赖于矩阵-向量乘法（GEMV），由于需要加载庞大的权重矩阵却只执行少量乘法，该阶段严重受限于内存带宽，导致计算强度低。批处理（Batching）通过组合多个请求，将密集层（K,Q,V生成、O投影和MLP）的计算转化为矩阵-矩阵乘法（GEMM），从而使其变为计算受限（Compute-bound）。然而，自注意力层由于无法跨请求共享具有不同上下文历史的KV-cache，依然是受限于KV-cache内存移动的GEMV操作。如图3所示，密集层和自注意力层占据了超过90%的处理时间。

权重-激活量化优于仅权重（Weight-only）量化。通过Roofline模型分析（图4），权重-激活量化由于使用了高效的低位硬件算术单元，提高了密集层的计算吞吐量；同时通过减小KV-cache的尺寸减少了内存移动，提升了自注意力层的吞吐量。相反，仅权重量化在矩阵乘法前必须进行反量化，计算仍采用浮点格式，未能提升密集层吞吐量，且未量化KV-cache，对自注意力层无益。

图3. 具有不同批量大小的Llama-7b推理的运行时间细分。密集层代表批处理的K、Q、V生成、O投影和MLP。自注意力层由集成了PageAttention的FlashInfer实现。结果表明，密集层和自注意力层合计占执行时间的90%以上，从而限制了吞吐量。
图3. 具有不同批量大小的Llama-7b推理的运行时间细分。密集层代表批处理的K、Q、V生成、O投影和MLP。自注意力层由集成了PageAttention的FlashInfer实现。结果表明，密集层和自注意力层合计占执行时间的90%以上，从而限制了吞吐量。
图4. 不同量化方法的Roofline模型，通过算术强度（定义为Ops/Elements）来表征算子。在大批量大小时，密集层是计算受限的，具有较大的算术强度，而自注意力层始终表现出较低的算术强度。
图4. 不同量化方法的Roofline模型，通过算术强度（定义为Ops/Elements）来表征算子。在大批量大小时，密集层是计算受限的，具有较大的算术强度，而自注意力层始终表现出较低的算术强度。

方法细节

激活异常值挑战与混合精度分离量化。LLM量化的一个关键挑战是激活中的异常值现象【1，LLM.int8(): 8-bit matrix multiplication for transformers at scale+2022】。如图5(a)所示，少数通道的幅度比其他通道大几个数量级，这种巨大的动态范围会显著增加量化误差。为了有效缓解这一问题，Atom采用混合精度方法，将异常值和正常值分别量化为低位和高位。如图5(b)所示，在去除异常值后，剩余通道变得更加均匀，可以有效地用低位值表示。由于8位表示（如FP8和INT8）足以表达异常值，且INT8得到了底层硬件（如NVIDIA Tensor Core）的广泛支持，Atom对异常值应用INT8量化。

图5. Llama-7b中激活矩阵的采样值。(a) 激活矩阵包含异常值通道，导致较大的量化误差。(b) Atom将这些异常值通道重排到矩阵末尾，并使用更高的精度对它们进行量化，同时保持规律的内存访问。
图5. Llama-7b中激活矩阵的采样值。(a) 激活矩阵包含异常值通道，导致较大的量化误差。(b) Atom将这些异常值通道重排到矩阵末尾，并使用更高的精度对它们进行量化，同时保持规律的内存访问。

通道重排以保持内存访问规律。混合精度量化的主要问题是其会导致不规则的内存访问，从而降低硬件效率。为了在保持规律内存访问的同时应用混合精度量化，Atom重新利用了RPTQ【2，Rptq: Reorder-based posttraining quantization for large language models+2023】中引入的重排技术。如图7所示，Atom将分散的激活异常值通道重排到矩阵的末尾，从而实现了混合精度的高效实现。为了保证计算结果的等效性，权重矩阵必须使用激活矩阵对应的重排索引进行重排。由于异常值通道可以使用校准数据离线识别，权重矩阵的重排是一次性成本。然而，激活矩阵的重排仍需在线执行。为了减轻这一开销，Atom将激活矩阵重排算子融合到前置算子中，这显著地将重排开销降低到不到运行时间的0.5%。

图7. Atom动态重排激活（A），将异常值通道移动到矩阵末尾，重排索引在离线校准中确定。权重矩阵（W）被静态重排以保持与相应的激活通道对齐，从而保证输出结果的正确性。
图7. Atom动态重排激活（A），将异常值通道移动到矩阵末尾，重排索引在离线校准中确定。权重矩阵（W）被静态重排以保持与相应的激活通道对齐，从而保证输出结果的正确性。

细粒度分组量化提升精度。即使Atom将异常值和正常值分开量化，由于4位精度的表示能力有限，对正常值的准确量化仍然具有挑战性。为了进一步提高准确性，Atom采用了分组量化【3，Awq: Activation-aware weight quantization for llm compression and acceleration+2023】【4，A white paper on neural network quantization+2021】，即将矩阵划分为子组并在每个子组内独立执行量化。例如，组大小为128意味着每连续的128个元素被视为一个组并独立量化。

融合反量化到GEMM流水线。分组量化在权重-激活量化中带来了反量化开销。Atom提出了一种融合技术（如图8所示），实现了一个高效的GEMM内核。Atom首先利用高效的低位硬件（即Tensor Cores）计算激活组与对应权重组的矩阵乘法，获得临时结果（步骤1）。接着，Atom将多个临时结果相加得到GEMM结果。然而，由于Atom对每个激活和权重组执行细粒度量化，每个临时结果具有不同的量化参数。因此，Atom首先使用CUDA Cores将所有临时结果反量化为FP16表示（步骤2），然后再执行加法（步骤3）。为了管理开销，Atom将反量化和求和融合到GEMM内核的MMA流水线【5，CUTLASS+2023+URL: https://github.com/NVIDIA/cutlass】中。因此，额外的操作可以在原地执行，无需额外的内存移动，并与原始的MMA指令重叠 。

图8. 融合GEMM算子概述。每组的乘法首先由具有高效低位支持的单元（即Tensor Cores）计算（步骤1）。然后使用典型的FP16单元对结果进行反量化并随后进行累加（步骤2、3）。请注意，所有操作都融合在单个流水线中。
图8. 融合GEMM算子概述。每组的乘法首先由具有高效低位支持的单元（即Tensor Cores）计算（步骤1）。然后使用典型的FP16单元对结果进行反量化并随后进行累加（步骤2、3）。请注意，所有操作都融合在单个流水线中。

动态量化捕捉局部输入分布。尽管细粒度量化可以更好地保留激活通道内的局部变化，但如果基于校准数据静态计算量化参数，这种优势就会减弱，因为实际输入可能具有不同的局部特征。因此，Atom采用动态量化，在推理期间为每个激活矩阵定制量化参数。为了控制动态量化的开销，Atom将量化操作融合到前置算子中（类似于ZeroQuant【6，Zeroquant: Efficient and affordable post-training quantization for large-scale transformers+2022】）。由于附加算子是逐元素的（包含一次归约和一次逐元素除法），其运行时间相比密集的dense和self-attention层可以忽略不计。为了在吞吐量和精度之间取得平衡，Atom选择具有精心挑选的裁剪阈值的对称量化。在量化权重矩阵时，Atom还结合了GPTQ【7，Gptq: Accurate post-training quantization for generative pretrained transformers+2023】，因为这是一个纯离线过程，可以在不牺牲运行时效率的情况下提升准确性。

KV-cache量化缓解内存瓶颈。为了缓解解码阶段自注意力层高度受限于内存的问题，Atom对KV-cache应用了低位量化。Atom以低位精度加载KV-cache，并在执行FP16计算之前直接对其进行反量化，这通过大幅减少内存使用量显著提升了吞吐量。由于非对称和对称量化的KV-cache的内存移动相似，Atom对KV-cache使用非对称量化以提供精度优势。如图9分析所示，V cache极少展现出异常值现象，使其更适合量化。由于Softmax的归一化作用，K cache的量化误差对输出影响较小。因此，Atom直接以注意力头为粒度应用非对称低位量化。

图9. Llama-7b中单个注意力头内V cache的采样值。与图5(a)中显示的采样激活相比，V cache显示出小得多的动态范围和更少的异常值通道，这更容易量化。
图9. Llama-7b中单个注意力头内V cache的采样值。与图5(a)中显示的采样激活相比，V cache显示出小得多的动态范围和更少的异常值通道，这更容易量化。

端到端工作流实现。为了证明设计选择的可行性，Atom在Llama模型上进行了实现（如图6所示）。Atom通过内核融合管理额外算子的开销：将重排、量化和反量化等量化算子融合到现有算子中。对于计算受限的密集层，Atom利用低位单元提升吞吐量。对于内存受限的自注意力层，Atom将反量化与LLM服务内核库FlashInfer【8，Accelerating self-attentions for llm serving with flashinfer+2024+URL: https://flashinfer.ai/2024/02/02/introduce-flashinfer.html】融合，使得仅加载来自KV-cache的低位值。Atom还结合了PageAttention【9 ，Efficient memory management for large language model serving with pagedattention+2023】进行高效的内存使用，以支持大批量大小。

图6. Llama模型系列上Atom工作流的概述。Atom通过将量化算子融合到现有算子中来仔细管理其开销。对于计算受限的算子，Atom利用高效的低位硬件支持。对于内存受限的自注意力层，Atom量化KV-cache以进一步提升吞吐量。我们为每个融合算子实现了专用内核。
图6. Llama模型系列上Atom工作流的概述。Atom通过将量化算子融合到现有算子中来仔细管理其开销。对于计算受限的算子，Atom利用高效的低位硬件支持。对于内存受限的自注意力层，Atom量化KV-cache以进一步提升吞吐量。我们为每个融合算子实现了专用内核。

实验环境

实验结果

1. Zero-shot 准确率评估
* 实验内容：在Llama模型系列上，对比Atom与SmoothQuant、OmniQuant、QLLM在W4A4和W3A3配置下的Zero-shot准确率。
* 实验结果：在W4A4配置下，Atom在Llama 7B、13B、30B和65B模型上相比FP16的平均准确率下降仅为$2.3\%$、$1.7\%$、$0.4\%$和$1.4\%$，而基线方法在相同设置下准确率下降幅度达到$9.6\%$至$23.8\%$。
* 分析结论：Atom在极低位量化下显著优于其他权重-激活量化方法，几乎无损地保留了模型的常识推理能力。

2. 困惑度（Perplexity）评估
* 实验内容：在WikiText2、PTB和C4数据集上评估各量化方法的生成困惑度。
* 实验结果：对于W4A4，Atom在Llama-65B模型上的三个数据集困惑度增加均小于0.4。对于W3A3，Atom在Llama-65B上的平均困惑度增加为2.3，而现有的基线方法（如SmoothQuant和OmniQuant）在W3A3下困惑度爆炸，无法达到可接受的水平。
* 分析结论：Atom有效降低了量化误差，且在量化更大参数量模型时精度损失更小。

3. 内核级效率评估（Kernel Evaluation）
* 实验内容：在Llama-7b配置下，对比Atom (W4A4)与FP16、W8A8、W4A16的融合GEMM算子和自注意力算子的吞吐量。
* 实验结果：
* 密集层（图11a）：在批量大小较小时，W4A16由于内存减少而有效；但随着批量大小增加（计算受限），W4A16效率下降。在批量大小为512时，Atom的矩阵乘法实现了比FP16高$3.4\times$、比INT8内核高$1.9\times$的加速。
* 自注意力层（图11b）：在批量大小为128时，Atom实现了比INT8量化高$1.8\times$、比FP16基线高$3.5\times$的加速。

图11. Atom和基线内核上不同量化方法的性能评估。我们设置了与Llama-7b配置和1024序列长度对齐的评估配置。内核由NVBench评估。
图11. Atom和基线内核上不同量化方法的性能评估。我们设置了与Llama-7b配置和1024序列长度对齐的评估配置。内核由NVBench评估。

4. 端到端服务评估
* 实验内容：将Atom集成到Punica框架中，使用ShareGPT工作负载，在连续批处理（Continuous batching）模式下测试端到端吞吐量和延迟。
* 实验结果：如图10(a)和10(b)所示，Atom在所有批量大小上均优于其他量化方法。在固定可用GPU内存的条件下（图10(c)），Atom能够支持更大的批量大小，使其吞吐量达到FP16基线的$7.73\times$，以及INT8量化的$2.53\times$。在延迟方面，即使在批量大小为256时，Atom的单Token解码延迟仍低于100ms。
* 分析结论：仅权重量化受限于密集层中的FP16计算能力和自注意力层中KV-cache的巨大内存移动，而Atom通过全面低位化打破了这些瓶颈，实现了极高的服务吞吐量和极低的延迟。

图10. Atom的端到端评估。实线是精确测量值，而虚线是由于内存容量有限而得出的估计值。(a) 每秒生成的token数。(b) 每个token的平均解码延迟。Atom在吞吐量和延迟方面都超过了所有其他量化方法。(c) 在固定数量的GPU内存下评估的性能。请注意，Atom将吞吐量提升了比W8A8多2.5倍，因为它支持更大的批量大小，从而利用了批处理效应。
图10. Atom的端到端评估。实线是精确测量值，而虚线是由于内存容量有限而得出的估计值。(a) 每秒生成的token数。(b) 每个token的平均解码延迟。Atom在吞吐量和延迟方面都超过了所有其他量化方法。(c) 在固定数量的GPU内存下评估的性能。请注意，Atom将吞吐量提升了比W8A8多2.5倍，因为它支持更大的批量大小，从而利用了批处理效应。

5. 量化技术的消融实验
* 实验内容：评估混合精度、分组量化、裁剪、GPTQ和KV-cache量化对精度（WikiText2 PPL）和效率（TOPS）的独立影响。
* 实验结果：
* 精度：基础W4A4 RTN的PPL高达2315.52。将128个异常值保留在FP16使PPL骤降至11.34；进一步将异常值量化为INT8仅带来0.05的PPL轻微增加。引入Group size 128使PPL降低5.17。裁剪和GPTQ各自降低0.09。最后量化KV-cache至INT4仅导致0.12的PPL增加。
* 效率：纯INT4 GEMM达到近980 TOPS。融合INT8混合精度计算带来$8\%$开销（900 TOPS）。融合分组反量化是主要开销，性能降至770 TOPS，但仍比INT8吞吐量的理论极限高出近$18\%$。将重排与量化算子融合的方法比基线（矩阵分解）在推理延迟上始终保持$25\%$到$35\%$的性能优势。

补充细节

模型通用性。Atom实现高精度的主要技术是针对异常值的混合精度量化和针对正常值的细粒度量化。经验证明，这些技术可推广到较新的基于Transformer的LLMs。在Llama-2和Mixtral 8x7B模型上的评估表明，Atom依然保持了高精度。为了适应混合专家（MoE）模型，Atom只需为不同专家的前馈神经网络（FFN）使用不同的重排索引即可。

数据格式通用性。随着新硬件（如NVIDIA Blackwell GPU）对新兴数据格式（如FP4和MX）的支持，Atom在FP4数据格式下也进行了评估。结果表明，将权重和激活量化为FP4时，Atom保持了与INT4相似的精度，证明了INT4和FP4之间的表示能力相似。此外，新硬件支持MX格式的分组量化，预计这一硬件特性将进一步缓解Atom的分组量化开销。

结论

本文提出了Atom，这是一种低位量化方法，能够高效利用底层硬件在LLM服务中实现高精度和高吞吐量。Atom结合了带有重排的混合精度量化、细粒度分组量化、动态量化以及KV-cache量化，在充分利用新兴低位硬件支持的同时保留了模型精度。将Atom集成到端到端服务框架中，与FP16基线相比，实现了高达$7.73\times$的吞吐量提升，同时保持了低于$1.4\%$的Zero-shot准确率损失。未来工作可进一步结合新型硬件特性（如FP4/MX格式的硬件级分组量化支持）来彻底消除反量化开销。