[论文解读] LLM-QAT: Data-Free Quantization Aware Training for Large Language Models

Zechun Liu, Barlas Oğuz+7 more|arXiv (Cornell University)|May 29, 2023

Topic Modeling被引用 15

一句话总结

LLM-QAT 引入面向大语言模型的无数据量化感知训练，通过生成数据蒸馏实现权重、激活值和 KV 缓存的 4 位量化，在跨 7B/13B/30B 的 LLaMA 模型中，优于后训练量化。

ABSTRACT

Several post-training quantization methods have been applied to large language models (LLMs), and have been shown to perform well down to 8-bits. We find that these methods break down at lower bit precision, and investigate quantization aware training for LLMs (LLM-QAT) to push quantization levels even further. We propose a data-free distillation method that leverages generations produced by the pre-trained model, which better preserves the original output distribution and allows quantizing any generative model independent of its training data, similar to post-training quantization methods. In addition to quantizing weights and activations, we also quantize the KV cache, which is critical for increasing throughput and support long sequence dependencies at current model sizes. We experiment with LLaMA models of sizes 7B, 13B, and 30B, at quantization levels down to 4-bits. We observe large improvements over training-free methods, especially in the low-bit settings.

研究动机与目标

• 在保留零样本和少样本能力的同时，推动大语言模型超过 8 位的量化。
• 提出一种无数据蒸馏的方法，从预训练模型中生成用于量化感知训练（QAT）的训练信号，而无需原始数据。
• 使用针对 LLM 分布量身定制的量化器，对权重、激活值和 KV 缓存进行量化。
• 证明带有无数据蒸馏的 QAT 能保持输出分布，在不同模型规模上优于 PTQ。
• 评估在实际部署中模型大小、位宽和吞吐量之间的权衡。

提出的方法

• 通过从预训练模型生成下一个词数据引入无数据蒸馏，使用混合采样策略（前 3–5 个令牌采用 top-1，其后采用随机采样）。
• 使用来自教师（全精度）的交叉熵 logits 蒸馏来训练量化学生；软标签在训练中优于硬标签。
• 对权重和激活采用对称的 MinMax 量化，权重量化按通道、激活按令牌进行量化。
• 对 KV 缓存按令牌进行量化，使用存储的缩放因子，并将量化集成到 QAT 过程的梯度计算中。
• 将 Transformer 的全连接层（包括 KV 缓存）量化为 4 位权重/4 位激活（某些设置中为 6 位激活）。
• 在 LLaMA-7B、-13B、-30B 上展示无数据蒸馏的有效性，达到首批准确的 4 位 LLM，并实现 KV 缓存量化以提升吞吐量。

实验结果

研究问题

• RQ1 无数据知识蒸馏在量化感知训练中是否能保持大语言模型的原始输出分布？
• RQ2 在保持零样本和少样本能力的同时，LLMs 的权重、激活值和 KV 缓存的位宽可以降到多低？
• RQ3 不同量化器、采样策略和蒸馏目标对 LLM 量化性能有何影响？
• RQ4 同时进行 KV 缓存量化是否在不牺牲准确性的情况下显著提高吞吐量？

主要发现

• LLM-QAT 使 LLaMA-7B/13B/30B 的权重、激活值和 KV 缓存实现 4 位量化，并在与 8 位 PTQ 基线的对比中保持显著的准确性。
• 在 8-8-8 和 4-8-4/4-8-8 设置下，LLM-QAT 通常优于最先进的 PTQ 方法，例如 30B 的 8-8-4 在零样本准确性方面高于 PTQ 基线。
• 4 位权重和 4 位激活配置在若干零样本任务上可接近全精度性能（平均损失约在 1.5 点内），在某些情况下 4 位激活显示出 6 位的可行性。
• 由模型通过下一个令牌采样生成的数据在零样本任务上的泛化性优于真实数据子集或仅 WikiText 微调。
• 来自教师的基于 logits 的知识蒸馏在消融实验中始终优于仅标签或注意力/隐藏状态蒸馏选项。
• 将 LLM-QAT 与 SmoothQuant 结合在某些 4/4/16 与 4/8/16 设置中取得增益，但收益取决于权重与激活的位分配。