像 LLaMA-2 和 Falcon 这样的大语言模型 (LLM) 彻底改变了人工智能领域,但它们伴随着高昂的代价: 计算成本。运行这些庞大的模型需要大量的 GPU 显存和处理能力。为了让这些模型能够在本地硬件或更便宜的云实例上使用,研究人员一直在竞相压缩它们。

最常见的技术是量化 (Quantization) ——即降低用于表示模型数字的精度。我们尝试将其压缩为 8-bit 甚至 4-bit,而不是使用 16-bit 浮点数 (FP16) 。

直到最近,4-bit 量化主要还是一种存储技巧。你可以用 4-bit 存储模型,但在进行实际数学运算 (推理) 时,必须将其解包回 16-bit。这有助于节省显存,但并不能让计算变得更快。事实上,解包过程有时反而会使其变慢。

在这篇文章中,我们将深入探讨 QUIK (Quantization to INT4 with Kernels,带内核的 INT4 量化) ,这是一种开创性的方法,允许 LLM 以 4-bit 精度执行实际的计算。这不仅仅是为了节省 RAM,更是为了追求纯粹的速度。QUIK 在像 RTX 3090 这样的消费级 GPU 上实现了高达 3.4 倍的端到端加速 , 且精度损失可以忽略不计。

我们将剖析 QUIK 如何解决“离群值问题”,它是如何优化 GPU 核函数 (kernels) 的,以及它在 LLaMA、OPT 和 Falcon 模型家族中取得的令人印象深刻的成果。

瓶颈: 内存 vs. 计算

要理解为什么 QUIK 是必要的,我们首先需要了解 LLM 推理的行为模式。推理通常分为两个阶段:

  1. 提示词处理 (预填充/Prefill) : 模型一次性读取你输入的提示词。这涉及同时处理许多 Token (批量处理) 。
  2. Token 生成: 模型逐个单词地生成答案。

大多数现有的量化技术 (如 GPTQ) 专注于“仅权重”量化。这减小了模型参数 (权重) 的大小。这对于内存受限 (memory-bound) 的任务非常有用。当你一次生成一个 Token 时,瓶颈通常是将巨大的权重矩阵从 GPU 显存移动到计算核心。如果权重更小,它们的移动速度就更快。

然而,当你处理长提示词或为多个用户运行批量请求时,瓶颈就会发生转移。GPU 有太多的数学运算要做,以至于显存速度不再像计算速度那么重要。这就是计算受限 (compute-bound) 场景。

Figure 2. Roofline analysis of a standard LLM MatMul operation. 图 1: Roofline 分析展示了从内存受限到计算受限性能的转变。

图 1 (源自论文的图 2) 所示,较少的 Token 数量 (1 到 16) 位于图表的斜坡部分——它们是内存受限的。但随着 Token 数量的增加 (128+) ,性能在顶部趋于平稳。这就是计算上限。仅权重的量化对于提升这个上限毫无作用。要在这里跑得更快,我们需要更快地进行数学运算。

4-Bit 数学的潜力

现代 NVIDIA GPU (Ampere 和 Lovelace 架构) 对低精度算术有硬件加速支持。与 FP16 相比,它们使用 INT8 或 INT4 数据类型执行矩阵乘法 (MatMul) 的速度要快得多。

Figure 3. Ideal matrix multiplication performance for different layer sizes and data precision on RTX3090. 图 2: 不同精度下的理论加速比。

图 2 展示了“理想”情况。与绿色的 FP16 线相比,蓝线 (W4A4——意味着 4-bit 权重和 4-bit 激活) 直冲云霄。如果我们能用 INT4 进行矩阵乘法,理论吞吐量是巨大的。

然而,这里有个陷阱。要利用这些快速的 INT4 核函数, 权重 (静态模型参数) 和激活 (流经网络的动态输入数据) 必须被量化为 4-bit。

挑战: 离群值问题

量化权重相对“容易”,因为它们是静态的。我们可以离线花时间找到压缩它们的最佳方法。激活则较难,因为它们随每个输入而变化。

4-bit 激活量化的真正杀手是离群值 (Outliers) 。 在 LLM 中,某些特征 (激活矩阵中的列) 往往具有比其他特征大几个数量级的值——有时大 100 倍。

想象一下,试图将 -100 到 100 范围内的数字放入一个 4-bit 整数中 (它只能表示 16 个不同的值) 。如果你的大部分数据在 -1 到 1 之间,但有一个“离群”值为 100,你就必须拉伸你的刻度以包含 100。这会彻底破坏小数值的精度,而大部分信息都存在于这些小数值中。

试图进行 W4A4 (4-bit 权重和激活) 量化的现有方法通常会导致精度大幅下降,正是因为它们破坏了这些离群特征。

QUIK 解决方案: 混合量化

研究人员提出了 QUIK , 一种混合策略。其核心洞察简单而强大: 不要量化离群值。

QUIK 不是强行将整个矩阵压缩为 4-bit,而是将计算分开进行:

  1. 基础部分 (Base) : 绝大多数权重和激活被压缩为 4-bit。
  2. 离群值部分 (Outliers) : 一小部分敏感特征保持较高精度 (FP16 或 INT8) 。

第一步: 识别与重排 (离线)

由于 LLM 中的离群特征倾向于出现在特定的通道 (列) 中,无论输入如何,QUIK 使用校准数据集离线识别这些“捣乱”的列。

一旦识别出来,算法就会对权重矩阵进行重新排序 (排列) 。它将对应于这些离群值的列移动到矩阵的末尾。

Figure 4. Outlier-aware quantization with QUIK. 图 3: QUIK 预处理流程。

图 3 所示,橙色 (输入) 和蓝色 (权重) 高亮的列是离群值。

  1. 提取: 识别离群列的索引。
  2. 排列: 将这些列移至矩阵的右侧。
  3. 量化: 对左侧的“基础”部分应用 GPTQ (一种复杂的量化算法) 。右侧的离群值部分保留为 FP16。

这种分离允许量化算法优化“正常”数据的 4-bit 比例,而不受巨大离群值的干扰。

第二步: 前向传播 (在线)

在推理过程中 (当模型实际运行时) ,系统执行一种分裂计算。

Figure 5. Schematic for the forward pass of a linear layer with QUIK-4B. 图 4: QUIK 中线性层的计算流程。

图 4 详细描述了运行时过程:

  1. 输入拆分: 输入矩阵 \(X\) 进来。系统立即将其一分为二: “基础”列和“离群”列。
  2. 双路径:
  • 路径 A (FP16): 离群列使用标准的 FP16 数学运算与未量化的离群权重相乘。这保留了最敏感数据的高精度。
  • 路径 B (INT4): 基础列被在线量化为 INT4。然后使用快速 GPU 核函数将它们与预量化的 INT4 权重相乘。
  1. 合并: INT4 结果被反量化回 FP16,并加到 FP16 离群值结果上,得到最终输出 \(Y\)。

反量化背后的数学

为了让 INT4 数学运算奏效,系统对激活使用非对称量化 (缩放因子和零点) ,对权重使用对称量化 (仅缩放因子) 。

反量化步骤——将整数结果转换回浮点数——至关重要。单个元素的计算如下所示:

Equation for dequantization

这里,\(w\) 是权重,\(x\) 是输入。项 \(z\) 代表“零点”偏移。作者的关键观察是,项 \(z \cdot \sum w_i\) 仅取决于权重和零点。权重的总和 (\(\sum w_i\)) 可以离线预计算。这显著减少了前向传播期间的计算开销。

高效实现: 速度至上

如果管理它的开销高于它提供的加速,那么理论算法就是无用的。拆分矩阵和在线量化数据需要时间。如果不小心处理,这些操作 (元数据处理、内存读/写) 可能会成为新的瓶颈。

为了解决这个问题,作者使用 NVIDIA CUTLASS 库开发了高度优化的自定义 GPU 核函数。

核融合 (Kernel Fusion)

标准方法是为每个步骤启动单独的 GPU “kernel” (函数) : 一个用于查找输入的最小/最大值,一个用于量化它,一个用于分离离群值等。每个 kernel 的启动都会产生开销,并需要对慢速的全局内存 (VRAM) 进行读/写。

QUIK 融合了这些操作。

  1. 量化融合: 单个 kernel 读取输入,计算最小/最大统计数据,量化基础部分,并提取离群部分,所有这些一气呵成。
  2. 反量化 Epilogue (收尾操作) : 不是将 INT4 MatMul 结果写入内存,然后再读回来转换为 FP16,而是在矩阵乘法结束时,当数据仍在快速 GPU 寄存器中时,立即进行转换。

Figure 6. Operation timings in different QUIK-4B versions. 图 5: 核融合对处理时间的影响。

图 5 展示了这些优化的巨大影响。“V1”柱状图显示了一个朴素的实现,其中元数据和量化是分开的。“V2”融合了它们 (绿色阴影柱) ,显著减少了时间。“V3”融合了反量化 (棕色阴影柱) 。对于较小的矩阵尺寸 (例如 8192x1024) ,与朴素方法相比,这些优化将总运行时间减少了近一半。

实验结果

那么,它有效吗?作者在三个主要模型家族上测试了 QUIK: OPT、LLaMA-2 和 Falcon。

精度恢复

4-bit 激活量化的主要担忧是“困惑度 (perplexity) 恶化” (即模型开始输出乱码) 。

Table 2. Perplexity results of QUIK for 4-bit LLaMA-2 and Falcon models. 图 6: 困惑度得分 (越低越好) 。QUIK 保持了接近基线的精度。

图 6 , 我们将基线 (FP16) 与 QUIK-4B 进行比较。

  • LLaMA-2 70B: 基线困惑度为 3.20。QUIK-4B 达到 3.74。
  • Falcon 180B: 基线为 3.30。QUIK-4B 为 3.61。

恶化程度极小——大约 0.3 到 0.5 个困惑度点。相比 SmoothQuant 或 OmniQuant 等以前的方法在 4-bit 设置下的表现 (通常会完全破坏模型或导致高得多的困惑度) ,这是一个巨大的进步。

加速比

主要目标是在计算受限场景下的速度。

Figure 1. Accuracy and speedups for QUIK at different model sizes. 图 7: LLaMA-2 模型的加速因子。

图 7 总结了这一胜利。对于庞大的 LLaMA-2 70B 模型,QUIK 实现了相对于 FP16 基线 3.4 倍的加速 。 即使对于 7B 和 13B 这样较小的模型,加速比也在 2.5 倍左右。

Figure 9. End-to-end inference speedups for QUIK-4B. 图 8: 吞吐量比较 (Token/秒) 。

图 8 提供了原始吞吐量数据。看图表 (a) 中的 OPT-66B:

  • 基线: 439 tokens/秒。
  • QUIK-4B: 1343 tokens/秒。

这有效地将单个 GPU 变成了一个动力室,能够在相同的时间内服务更多的用户或处理更长的文档。

案例研究: 敏感的 LLaMA-2 Down-Projection 层

研究过程中发现的一个迷人细节是 LLaMA-2 架构的特殊敏感性。LLaMA 在其前馈 (MLP) 块中使用了 SwiGLU 激活函数。

MLP 块由三个线性层组成: Up-Proj、Gate-Proj 和 Down-Proj。“Down-Proj”层采用前两层的逐元素乘积。

Figure 10. The variance of the inputs in different layers of LLaMA2-70B. 图 9: 各层的输入方差。注意红线 (Down-Proj) 。

图 9 所示,与其他层相比,Down-Proj 层 (红线) 的输入方差巨大。这意味着数值波动剧烈,使得 4-bit 量化对这一特定层极具破坏性。

为了解决这个问题,作者在模型内部采用了一种“混合精度”方法。

  • 大多数层: 4-bit (QUIK-4B)。
  • Down-Proj 层: 8-bit (QUIK-8B)。

Figure 11. FLOP/s analysis of the LLaMA2-70B linear layers with QUIK. 图 10: LLaMA-2 70B 的精度细分。

图 10 可视化了这种细分。绿色部分代表 INT4 计算,占总工作量的绝大部分 (约 70%) 。橙色部分代表为了保持精度而保留在 INT8 中的特定 Down-Proj 层。这种有针对性的精度允许模型在保持准确的同时,不牺牲 4-bit 处理带来的整体速度优势。

结论

QUIK 论文填补了大语言模型效率方面的一个重要空白。虽然以前的工作成功地压缩了模型权重以节省空间,但 QUIK 证明了我们也可以将计算 (激活) 压缩到 4-bit 以节省时间。

通过智能地识别离群值并将其保持在高精度——同时对其余部分使用高度优化的融合核函数——QUIK 实现了推理的“圣杯”:

  1. 高压缩率: 显存占用减少高达 3 倍。
  2. 高速度: 吞吐量提高高达 3.4 倍。
  3. 高精度: 将困惑度保持在未压缩模型的 0.5 点以内。

对于学生和从业者来说,这强调了一个重要的教训: 简单地应用理论算法 (如“将一切量化为 4-bit”) 在实践中很少奏效。现实世界的性能需要对数据分布 (离群值) 和硬件 (核融合) 有深刻的理解。QUIK 是算法-硬件协同设计释放现代 GPU 全部潜力的一个典型例子。