太长不看版

采用 QAT 量化 LLM 的首个成功实践。

量化方案：

Weight： Per-channel，Activation： Per-token，KV cache： Per-token

一些后训练量化 (Post-Training Quantization, PTQ) 方法已应用于大语言模型 (LLM)，并在 8-bit 设置下表现良好。本文发现这些方法在更低比特的设置中会出现严重损坏。因此，本文研究使用量化感知训练 (Quantization-Aware Training, QAT) 方法，以进一步推动量化水平。

本文还提出了一种 Data-free 的知识蒸馏方案，利用预训练模型的输出，更好地保留了原始输出的分布，并允许量化任何模型，与训练数据无关。除了量化 weight 和 activation 之外，这个工作还量化了 KV cache，这对提高吞吐量以及支持当前模型的长序列依赖关系至关重要。作者在 7B，13B，30B 的 LLaMA 模型上进行了实验，量化低至 4-bit。结论是观察到 LLM-QAT 相比于 Training-free 方法有显著改进，尤其是在低比特设置中。

专栏目录

https://zhuanlan.zhihu.com/p/687092760

本文目录

1 LLM-QAT：无需数据的 LLM 量化感知训练
(来自 Meta)
1 LLM-QAT 论文解读
1.1 LLM-QAT 研究背景
1.2 LLM-QAT 方法概述
1.3 无数据知识蒸馏
1.4 量化感知训练
1.5 实验设置
1.6 主要实验结果
1.7 数据选择
1.8 量化函数
1.9 知识蒸馏

1 LLM-QAT：无需数据的 LLM 量化感知训练

论文名称：LLM-QAT: Data-Free Quantization Aware Training for Large Language Models

论文地址：https://arxiv.org/pdf/2305.17888

1.1 LLM-QAT 研究背景

GPT-3 之后，一些大语言模型 (LLM) 家族，例如 OPT、PALM、BLOOM、Chinchilla 和 LLaMA 已经证实，增加模型大小可以提高模型能力。因此，今天人工智能领域的基本规范逐渐演变为了数十亿甚至数百亿参数规模的语言模型。尽管 LLM 的能力值得兴奋，但由于计算成本和环境代价大，为数十亿用户的利益提供这样的模型面临着重大障碍。

幸运的是，最近的一些工作，比如 SmoothQuant[1]，ZeroQuant[2]等专注于 weight 和 activation 的 8-bit 后训练量化，并且几乎精度无损。然而，65 亿参数 LLaMA 模型光是权重就占用 65GB 的 GPU 显存。此外，注意力层的 KV cache 可以很容易地产生数十 GB，是当前一些应用中长序列机制下的吞吐量瓶颈。上述工作只考虑了 weight 和 activation 的量化，没有考虑 KV cache 的量化。不幸的是，当进行 8-bit 以下的量化时，SoTA 的 PTQ 方法的质量出现显著下降。对于更高量化要求，作者发现有必要求助于量化感知训练 (QAT)。

在本文公开的时间节点，LLM 的 QAT 方法还没有见到过，作者认为原因有二。首先，LLM 训练技术上困难，且是资源密集型的。其次，QAT 需要训练数据，这对于 LLM 来讲很难获得。预训练数据的绝对规模和多样性本身就是一个障碍。预处理简直令人望而却步，更糟糕的是，由于法律限制，某些数据可能根本不能用。分多个阶段训练 LLM 也很常见，包括指令调优和强化学习，这在 QAT 期间很难复制。这项工作中，作者通过使用 LLM 本身生成的数据进行知识蒸馏来规避这个问题。这种简单的解决方法，作者称之为 Data-free 的知识蒸馏，适用于任何独立于原始训练数据是否可用的生成模型。与在原始训练集的子集上进行训练相比，本文方法能够更好地保留原始模型的输出分布。此外，可以只使用一小部分 (100K) 的采样数据成功地提取量化模型，使计算成本也更合理。所有实验都是使用单个 8 GPU 训练节点进行的。

基于上述做法，本文可以蒸馏出 7B、13B 和 30B 的 LLaMA 模型，weight 和 KV cache 成功量化为 4-bit。在这方面，与 PTQ 相比，本文的 QAT 方法显著提升了质量。此外，本文成功地将 activation 量化为 6-bit 精度，超过了现有方法的。

总之，本文可以看做是 QAT 对 LLM 的第一个成功实践，得到准确的 4-bit 量化 LLM。作者还同时量化了 KV cache，weight 和 activation，这对于缓解长序列生成的吞吐量瓶颈至关重要。这些是通过无数据蒸馏方法实现的。

1.2 LLM-QAT 方法概述

使用量化感知训练 (QAT) 量化 LLM 是一个不平凡的任务，有两方面的挑战：

其一，预训练的 LLM 在 Zero-Shot 泛化能力方面表现出色，那么在量化后保持这种能力至关重要。因此，选择合适的微调数据集很重要。如果 QAT 数据域太窄或与原始预训练的分布有显著的不同，则可能会损害模型性能。

其二，由于 LLM 训练的规模和复杂性，很难精确复制原始的训练设置。因此本文对 LLM 进行无数据量化感知训练 (QAT)，使用下一个 token 数据生成生成 QAT 数据。与使用原始预训练数据的子集相比，该方法表现出更好的性能。其次，LLM 相比于小模型，表现出独特的 weight 和 activation 的分布，其特征是存在大量异常值。因此，最先进的量化裁剪方法对于 LLM 适配得并不好。因此，本文为 LLM 确定了一个合适的量化器。

1.3 无数据知识蒸馏

为了以有限的微调数据密切合成预训练数据的分布，本文提出了从原始预训练模型生成的下一个 token 数据。如图 1 (a) 所示，从词汇表中随机化第 1 个 token <> ，并让预训练模型生成下一个 token <> ，然后将生成的标记附加到起始标记以生成新输出 <> 。重复这个迭代过程，直到到达句子的末尾或最大生成长度。

作者测试了 3 种不同的采样策略。最直接的方法是选择 top-1 candidate 作为下一个 token。但是，生成的句子缺乏多样性，会循环重复几个 token。为了解决这个问题，改为使用预训练模型的 SoftMax 输出概率分布，从分布中采样下一个 token。这种采样策略产生了更多样的句子，大大提高了微调学生模型的精度。

此外，作者发现最初的少数一些 token 在确定预测趋势方面起着至关重要的作用。因此，它们具有更高的置信度是很重要的。为此，本文的生成过程采用了一种混合采样策略，该策略为前 3~5 个 token 确定性地选择 top-1 预测，并对剩余的 token 随机采样。

1.4 量化感知训练

本文研究线性量化，即均匀量化。线性量化可以根据实际值是否被裁剪分为两类：MinMax 量化，它保留所有值范围，以及基于裁剪 (clipping) 的量化。

在 MinMax 量化中，量化过程可以表述为：

这里 和 分别表示量化和全精度变量。 指的是张量中的第 个元素。 是 scaling factor, 是零点。对于对称量化， 。对于非对称量化， 。与 MinMax 量化相比, 裁剪异常值有助于提高精度并将更多的比特分配给中间值。因此，最近的工作[3]对语言模型采用基于裁剪的量化。

在基于 clipping 的量化中，量化过程可以表述为：

其中 scale 和零点 可以通过梯度统计或学习。

量化函数

下图 2 中为本文 Transformer 模型的量化。与 SmoothQuant[1]的发现一致，作者观察到大语言模型 (LLM) 的 weight 和 activation 中存在异常值。这些异常值对量化过程有显著的影响，因为会使量化步长增加，降低中间值的精度。然而，在量化过程中裁剪这些异常值被证明不利于 LLM 的性能。在训练的初始阶段，任何基于裁剪的方法都会导致极高的困惑度分数，导致大量信息丢失，而且难以通过微调恢复。因此，作者选择保留这些异常值。此外，作者发现在具有门控线性单元 (GLU) 的模型中，activation 大多是对称分布的。作者为 weight 和 activation 选择对称 MinMax 量化：

式中, 表示量化的 weight 或 activation, 表示实值 weight 或 activation。对 activation使用 per-token 量化, 对 weight 采样 per-channel 量化, 如图 3 (a) 所示。

KV cache 的量化感知训练

除了 weight 和 activation 量化外，LLM 中的 KV cache 也消耗不可忽略的显存。在这项研究中，作者在本文中证明了用于 activation 量化的类似 QAT 方法也可以用于 KV cache 的量化。如图 3 所示，在式 3 中采用了 per-token 量化，考虑到 Key 和 Value 是逐 token 生成的。在生成过程中，对当前 Key 和 Value 进行量化，并存储它们对应的 scale factor。在 QAT 的训练过程中，将量化应用于 Key 和 Value 的整个 activation 张量，如图 2 所示。通过将量化函数集成到梯度计算中，确保使用量化 Key-Value 对进行有效训练。

知识蒸馏

作者使用基于交叉熵的 logits 蒸馏从全精度预训练的教师网络训练量化学生网络：

式中, 表示当前批次中具有 个句子的第 个样本, 表示类的数量, 它等于词汇表的大小。 和 分别是教师网络和学生网络。作者利用预训练模型的预测作为软标签, 为指导学生模型的训练提供了更多信息。

1.5 实验设置

在量化网络训练过程中，使用预训练的模型初始化模型，并将其作为教师进行知识蒸馏。为了优化模型，使用了权重衰减为零的 AdamW 优化器。每个 GPU 分配的 Batch Size 为 1，学习率设置为 2e-5，遵循余弦学习率衰减策略。数据生成使用 LLaMA-7B 模型，生成序列的最大长度设置为 1024。

1.6 主要实验结果

比较 3 种后训练量化 (PTQ) 方法：RTN、GPT-Q 和 SmoothQuant。作者在几种不同的设置中与它们进行比较，其中 weight、activation 和 KV cache 被量化为不同级别 (表示为 W-A-KV)。不同的PTQ方法在不同的设置下表现良好，本文将 LLM-QAT 与每个设置中的最佳 PTQ 结果进行比较。

图 4、图 5 和图 7 分别比较了所提出的 QAT 方法与 SOTA PTQ 方法在 Common Sense Reasoning 任务的 Zero-Shot 结果，对 Wiki2 和 C4 的困惑度评估以及 MMLU 和 TriviaQA 基准上的 Few-Shot 结果。困惑度评估验证了量化模型是否能够在其训练域的不同样本上保留模型的输出分布。如果保留模型在下游任务上的能力，零样本和少样本评估措施。

每个表的趋势是相似的。所有方法在所有模型大小的 8-bit 设置中都倾向于表现良好。即使 KV cache 也被量化为8位 (weight 和 activation 也被量化为 8-bit 也成立)。然而，当这 3 个中的任何一个被量化为小于 8 位时，PTQ 方法会出现精度损失，而 LLM-QAT 精度保持得更好。在 8-8-4 设置中，30B LLM-QAT 的平均 Zero-Shot 精度为 69.7，而 SmoothQuant 的平均 Zero-Shot 精度为 50.7。在 4-8-8 设置中差异较小，但是 LLM-QAT 仍然优于最佳 PTQ 方法 (本例中为 RTN) 1.4。在 4-8-4 设置中，weight 和 KV cache 都量化为 4-bit，所有 PTQ 方法都产生较差结果，而 LLM-QAT 达到 69.9，平均仅落后全精度模型 1.5 分。LLM-QAT 在 6-bit activation 量化方面也很有效。

从业者的一个重要问题是是否以全精度使用小模型，或者使用更大的量化模型类似的推理成本。虽然确切的权衡可以根据几个因素而有所不同，但可以根据本文结果提几个建议。首先，8-bit 量化应该比较小的全精度模型更受欢迎，PTQ 方法足以满足这种情况。8-8-8 30B 量化模型优于 13B 大小相似的模型，在实践中应该具有较低的延迟和更高的吞吐量。与 16 位 7B 模型相比，这也适用于 8 位 13B 模型。此外，使用 LLM-QAT 量化的 4-bit 模型应该优于类似大小的 8-bit 模型。例如，4-8-4 LLM-QAT 30B 优于 8 位 LLaMA-13B，4-8-8 LLM-QAT 13B 优于 8 位 LLaMA-7B。因此，作者建议 4-bit LLM-QAT 模型以获得最佳的效率和精度权衡。

1.7 数据选择

在图 7 中，作者观察到使用从 Wikipedia 中提取的文本构建的 WikiText，不包括预训练期间使用的所有信息。因此，仅在 WikiText 上微调的模型往往会过拟合这个特定的数据集，难以很好地推广到其他数据集。另一方面，Crawled Corpus (C4) 数据集是由从网络中收集的数百 GB 干净的英文文本组成。在对 WikiText 数据集进行评估时，在 C4 上微调模型会产生合理的传输精度。然而，当任务 Zero-Shot 推理任务时，它表现出较差的精度。

作者还测试了 3 种数据生成策略。

• Generated data1 是指不采样，只选择 top-1 candidate token。
• Generated data2 是指从分布中采样下一个 token。
• Generated data 3 指的是前面 3~5 个 token 是通过确定性选择生成的，而其余的 token 是从分布中采样得到的。

与现有的数据相比，在生成的数据上微调的模型表现出卓越的泛化能力，尤其是在 Zero-Shot 任务中。此外，与没有采样生成的数据相比，从分布中采样生成的数据表现出更大的多样性。这种增强的多样性可以显着提高所有任务的性能。

1.8 量化函数

作者在图 8 中将 no-clipping 量化方法与基于 clipping 的方法进行了比较。按照 BiT[4]的做法，使用 StatsQ[5]，一种基于 clipping 的 weight 量化的统计计算 scale factor，和 LSQ[6]，基于 clipping 的 activation 量化的可学习比例因子。

然而，本文的研究结果表明，这两种最先进的基于 clipping 的量化方法并没有超过 MinMax 非裁剪方法的性能。这一观察结果强化了保留异常值对大型语言模型性能至关重要的论点。

1.9 知识蒸馏

图 9 显示，不同的知识蒸馏方法对微调模型的最终准确性有显著的影响。值得注意的是，由于在生成过程中从候选者的分布中采样引入的固有随机性和噪声，单独使用下一个 token 作为标签的结果是次优的。相比之下，与基于 label 的训练方法相比，利用教师模型的完整 logit 分布的 logit 蒸馏可以显著提高微调模型的性能。有趣的是，作者还观察到结合 Attention 蒸馏或 Hidden 层的蒸馏会阻碍性能。因此，作者在所有实验中使用 logit 蒸馏。

参考

1. ^abSmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
2. ^ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers
3. ^Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT
4. ^BiT: Robustly Binarized Multi-distilled Transformer
5. ^Nonuniform-to-Uniform Quantization: Towards Accurate Quantization via Generalized Straight-Through Estimation
6. ^Learned Step Size Quantization

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT