基于重排序的新量化方法RPTQ：实现大型语言模型的 3 比特量化

大型语言模型（LLMs）在各种任务上表现出色，但由于其庞大的模型规模，部署方面存在挑战。

在这篇论文中，来自后摩智能、腾讯 AI Lab、华中科技大学、北京大学、伊利诺伊理工学院的研究人员发现量化大型语言模型的主要挑战来自于通道之间不同的激活范围，而不仅仅是离群值问题。

作者提出了一种新颖的基于重排序的量化方法 RPTQ，解决了量化大型语言模型激活的问题。RPTQ 通过重新排列激活中的通道，然后将它们分簇进行量化，从而降低通道范围差异的影响。

此外，作者通过避免显式重新排序来减少存储和计算开销。该工作首次将 LLM 模型的激活量化推向 3 比特。激活值的 3 比特量化同时也能与权重的 3 比特量化共同使用，大幅度降低了大型欲言模型的开销。

论文地址：https://arxiv.org/abs/2304.01089

开源地址：https://github.com/hahnyuan/RPTQ4LLM

大规模语言模型（LLMs）在各种任务中展现出了卓越的性能，但由于它们的巨大模型大小，它们的部署面临着挑战。例如一个 66B 模型包含 660 亿个参数，需要大量内存进行存储。另外随着序列长度和批处理大小的增加，内存消耗问题变得更加严重，因为激活也需要消耗大量的内存，例如 Self-attention 中的键和值缓存（Key/Value Cache）。然而，单个 GPU 或服务器并没有足够的内存容量来存储如此庞大的模型权重和中间激活。因此，LLM 需要被分割成多个块并存储在不同的设备中。由于权重和激活存储在不同的设备上，在计算过程中需要在这些设备之间传输数据，导致显著的带宽和能量消耗。

为了克服 LLMs 在模型规模上的挑战，模型量化已成为一种有前途的方法。该技术涉及使用低位整数对 LLMs 的权重和激活进行量化，从而显著降低存储和运行这些模型的成本。具体而言，量化有助于减少保存张量的内存要求，并加速计算密集型操作，例如矩阵乘法和线性层。通过量化权重和激活，减少了存储和通信开销，提高了效率和推理速度。

图 1 不同通道的数值分布情况。每个点都代表一个激活通道中的 (最大值，最小值)。

最近，GPTQ 使用 PTQ 成功地将 LLMs 的权重量化为 4 bit 或 3 bit。然而，对 LLMs 中的激活进行量化仍然是一项具有挑战性的任务。

从图 1 中可以得出两个观察结果：

1）一些通道具有显著的异常值，其最大或最小值比其他通道大几百倍。以前的研究也已经发现了这个问题并针对异常值提出了特殊处理。

2）不同的通道在数值范围上有显著差异。使用相同的量化参数来量化不同的通道可能会导致显著的量化误差。即使两个通道具有相同的异常值绝对值，它们在数值范围上也可能存在很大的变化。

将整个张量使用相同的量化参数进行量化的每张量量化技术可能不太有效。原因是将量化范围设置为覆盖大范围的数值范围可能导致数值较小的通道产生较大的量化误差，而将其设置为覆盖小范围的数值范围可能导致对异常值的显著截断，并导致显著的量化误差。例如，一个通道可能具有 - 100 到 - 50 的值范围，而另一个通道可能具有 80 到 100 的值范围。试图通过将范围从 - 100 到 100 量化来覆盖它们的范围将导致两个通道产生显著的量化误差。这是以前的研究中尚未有效解决的一个挑战。

图 2 解决激活量化问题的不同方法对比。

为了解决激活中通道差异的问题，作者提出了一种新颖的基于重新排序的训练后量化方法 RPTQ。

如图 2 所示，该方法将数值范围相似的通道聚类在一起，然后对每个聚类中的值使用相同的量化参数进行量化。为了实施基于重排序的量化方法，作者首先使用校准数据集作为推理输入，从中导出每个激活通道的最大值和最小值。

接下来，他们采用 KMeans 算法将不同通道划分为 g 个簇，基于每个通道的最大值和最小值形成的点。一旦建立了簇，他们就进行通道重排序，将来自同一簇的通道放置在相近的位置。这样，具有类似最大值和最小值的通道被组合在一起，并共享一组通用的量化参数。在重排序过程之后，作者对每个簇内的激活进行量化。这种方法为每个簇单独计算量化参数（缩放因子 s 和零点 z），确保这些参数专门针对相应的通道。

最终，该方法确保了每个簇的量化都得到优化，减少了量化误差。

图 3：一重排权重和激活的量化 Transformer 层的推断过程失意图。重排索引用符号 R1 到 R5 表示。

显式重排是一种运行时重新排列激活中通道的操作，需要将不同通道的数据从一个内存位置物理移动到另一个位置，因此对于具有大量通道的大型模型，重排过程可能非常耗时。此外，存储源和目标激活张量会增加内存开销。因此，作者提出了避免显式重新排序的策略，从而降低了计算开销，提高了推理效率。

如图 3 所示，作者重新排序线性层的权重，使它们能够直接按照排序顺序产生激活。此外，作者修改了 LayerNorm 的方法，使其能够直接产生重新排序的激活，从而省去了在推理过程中进行显式通道调整。

作者在三种不同的位宽配置下评估了 OPT 的性能：W4A16、W4A8 和 W4A4。此外，作者还开发了一种新的量化方案，W4A4KV、W4A3KV 和 W3A3KV，只对主要内存的消耗 —— 键缓存和值缓存（Key/Value Cache）进行量化。根据表格，可以得出以下观察结果：一般而言，随着激活量化位宽的减少，模型的性能会下降。但在大多数情况下，这种退化并不显著。对于某些任务，性能下降不到 1%，甚至有所提高。

LLM 的内存开销来自三个方面：首先，设备内存中权重的存储，可以通过量化来显著减少。其次，在网络执行期间生成的临时激活需要分配设备内存。这些临时激活可以在使用后释放，因此不会对整体内存使用造成显著影响。最后，键和值缓存（Key/Value Cache）对预测后续单词是必要的。当批次大小和序列长度很大时，键和值缓存占用了大部分内存。

表中展示了不同设置下的内存使用情况，可以发现较低比特的激活可以显著减少内存使用，特别是在批次大小和序列长度较大的情况下。具体而言，当批次大小为 64，序列长度为 8192 时，OPT-66b 下的 W4A4 配置相比 W4A16 配置节省了超过 900GB 的内存。因此，采用 RPTQ 对激活量化可以极大地减少在长文本任务或大批量场景中的内存压力。