太长不看版

大语言模型 4-bit 权重后训练量化方案，以及专为 4-bit On-Device LLM/VLMs 定制的推理框架 TinyChat。

大语言模型 (LLM) 已经改变了许多 AI 应用。On-device LLM 也变得越来越重要：在边缘设备上本地运行 LLM 可以降低云计算成本，保护用户的隐私。然而，巨大的模型尺寸以及有限的硬件资源为部署带来了重大的挑战。本文提出了激活感知权重量化方法 (Activation-aware Weight Quantization, AWQ)，一种硬件友好的 LLM 低 bit 权重量化方法。

AWQ 发现，并非 LLM 中的所有权重都同等重要。仅保护 1% 的 salient 权重可以大大减少量化误差。为了识别 salient 的 weight channel，应该参考 activation 而不是 weight 的分布。为了避免硬件效率低下的混合精度量化，作者在数学上推导出扩大 salient channel 可以减少量化误差。AWQ 采用等效变换来缩放 salient 的 weight channel 以保护它们。AWQ 不依赖于任何反向传播或重建，因此它可以推广到不同的领域和模式，而不会过度拟合校准集。

AWQ 在各种语言建模和特定领域基准 (代码和数学) 上的表现优于现有工作。由于更好的泛化，它为指令调整的 LM 实现了出色的量化性能，并且首次实现了多模态 LM。除了 AWQ，作者还实现了 TinyChat，这是一个高效且灵活的推理框架，专为 4 位设备上的 LLM/VLM 量身定制。通过 Kernel Fusion 和平台感知的权重打包，TinyChat 在桌面和移动 GPU 上都比 Huggingface FP16 实现提供了 3 倍以上的加速。它还使 70B Llama-2 模型在移动 GPU 上的部署民主化。

下面是对本文的详细介绍。

专栏目录

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

本文目录

1 AWQ：适合端侧的 4-bit 大语言模型权重量化
(来自 MIT 韩松团队，NVIDIA)
1 AWQ 论文解读
1.1 AWQ 研究背景
1.2 只保留 1% 的 Salient Weight
1.3 通过激活感知缩放保护 Salinet Weight
1.4 TinyChat：将 AWQ 映射到边缘设备
1.5 与经典工作 SmoothQuant 的关系
1.6 实验设置
1.7 实验结果
1.8 实际加速效果

1 AWQ：适合端侧的 4-bit 大语言模型权重量化

论文名称：AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (MLSys 2024 Best Paper)

论文地址：

http://arxiv.org/pdf/2306.00978

代码链接：

http://github.com/mit-han-lab/llm-awq

Demo：

http://www.youtube.com/watch%3Fv%3Dz91a8DrfgEw

1.1 AWQ 研究背景

直接在边缘设备上部署大语言模型 (LLM) 至关重要。On-device 地使用 LLM 的好处有：

1. 隐私和安全性： 通过保持敏感信息本地来增强数据安全性，从而减少数据泄露的机会。
2. 实时性： 消除将数据发送到云服务器引起的延时，并使 LLM 离线操作，有利于虚拟助手、聊天机器人和自动驾驶汽车等实时应用。
3. 降低运营成本： 减少与维护和缩放集中式云基础设施相关的运营成本。

LLM 基于基于 Transformer 架构，因其在不同基准测试中令人印象深刻的性能而引起了极大的关注。但是，较大的模型尺寸会导致较高的服务成本。例如，GPT-3 有 175B 参数，FP16 为 350GB，而最新的 B200 GPU 只有 192GB 内存，更不用说边缘设备了。

低比特权重量化可以显著降低 On-device LLM 推理的内存占用。由于训练成本高，量化感知训练 (Quantization-aware training, QAT) 效率不高，而后训练量化 (Post-Training Quantization, PTQ) 在低比特设置下存在较大的精度下降。

LLM 在当时，其量化一般有2种设置：

1. W8A8量化，其中激活和权重都被量化为 INT8，比如 SmoothQuant [1] 。
2. 低比特 Weight-only 量化 (例如 W4A16)，其中只有 weight 被量化为低位整数，比如 GPTQ [2] 。

本文属于是第2种，因为这样不仅减少了硬件障碍 (需要更低的显存)，而且还可以加速 token 的生成 (补救了 显存受限负载)。

与本文最接近的工作是 GPTQ，它使用二阶信息来执行误差补偿。然而，它可能会在重建过程中过度拟合校准集，扭曲分布外域上学习到的特征，这是有问题的，因为 LLM 是通才模型。

本文提出了激活感知权重量化 (Activation-aware Weight Quantization, AWQ)，这是一种硬件友好的 LLM 低比特权重量化方法。本文的方法基于这样一个观察：Weight 对 LLM 的性能并不同等的重要。Salient weight 只有一小部分 (0.1%-1%)，不量化这些 salient weight 将显着降低量化损失 (图2)。

为了找到这些 salient weight channel，本文给出的见解是：应该参考 activation 分布而不是 weight 分布，尽管是在进行 weight-only 的量化：对应更大激活幅值的 weight channel 更加 salient，因为它们处理更重要的特征。为了避免硬件效率低下的混合精度实现，作者从权重量化中分析误差，并推导出 scaling salient channel 可以减少其相对量化误差。根据直觉，作者设计了一种 per-channel scaling 方法自动搜索最优缩放，在量化全部权重的情况下最小化量化误差。AWQ 不依赖任何反向传播或者重建，因此它可以很好地保留 LLM 在各种域和模态的泛化能力，而不会过度拟合校准集。

为了实现 AWQ，作者设计了 TinyChat，一个高效推理框架，可以将 4-bit LLM 的理论显存节省转换为可以测量的实际加速。本框架通过动态去量化显着加快了线性层。作者还利用高效的 4-bit 权重打包和 Kernel Fusion 来最小化推理开销 (比如中间 DRAM 访问和 Kernel 启动的开销)，以便更好地实现从量化权重到 4-bit 的速度，尽管计算机是字节对齐的。

在桌面、笔记本电脑和移动 GPU 上，可以始终观察到：与各种 LLM 的 Huggingface 的 FP16 实现相比，平均加速比为 3.2-3.3 倍。此外，它可以在单个 NVIDIA Jetson Orin 上轻松部署内存为 64GB 的 Llama-2-70B。它还以 30 tokens/second 的交互式速度在只有 8GB 内存的笔记本电脑 RTX 4070 GPU 上运行了 1.3B 参数的 LLM。

AWQ 已被行业和开源社区广泛采用：HuggingFace Transformers、NVIDIA TensorRT-LLM、Microsfot DirectML、Google Vertex AI、Intel Neural Compressor、Amazon Sagemaker、AMD、FastChat、vLLM、LMDeploy，并使 Falcon180B 可部署在单个 H200 GPU 上。

1.2 只保留 1% 的 Salient Weight

量化将浮点数映射到低比特整数，是一种降低 LLM 大小和推理成本的有效方法。在本节中，作者首先提出了一种 Weight-only 的量化方法，通过保护更多的 "重要" 权重来提高准确性，而无需训练/回归。然后开发了一种数据驱动的方法来搜索减少量化误差的最佳缩放。

观察到 LLM 的权重并非同等重要：与其他模型相比，一小部分 salient weight 对于 LLM 的性能更重要。不量化这些 salient 的量化因为量化损失而可以帮助弥补性能下降，同时无需任何训练或者回归，如图3所示。为了验证这个想法，作者在图2中对量化 LLM 的性能进行了基准测试，这里就跳过了一些 weight channel 的量化。作者测量了 INT3 量化模型的性能，同时将一定比例的权重通道保持为 FP16。确定权重重要性的一种广泛使用的方法是查看其大小或 L_2 -Norm [3] 。但是作者发现跳过具有 Large Norm 的 weight channel (based on \text{W} ) 并不能显着提高量化性能，与随机选择结果是类似的。但是有趣的是，基于激活幅度选择权重可以显著提高性能，尽管只保留了 0.1%-1% 的通道为 FP16。作者假设具有较大幅度的输入特征通常更重要。保持 FP16 中相应的权重可以保留这些特征，从而获得更好的模型性能。

局限性： 尽管保留 0.1% 的权重为 FP16 可以提高量化性能，而不会显著增加模型大小 (使用 total bits 测量)，但是这种混合精度数据类型会使得系统实现变得困难。需要提出一种保护重要权重的方法，而不实际将其保留为 FP16。

1.3 通过激活感知缩放保护 Salinet Weight

本文提出了一种替代方法，通过逐通道缩放 (Scaling) 来减少 Salinet Weight 的量化误差，不会受到硬件效率低下问题的影响。

分析量化误差

作者首先分析了 Weight-only 的误差。考虑一个权重 的组块；线性运算可以写成 ,量化对应为 。具体来说, 量化函数定义为:

其中 是量化比特数， 是由绝对最大值确定的量化 scaler。现在考虑一个权重 ，如果将 与 和 inverse scaler 相乘, 则将有： 。上式即可以写成：

其中 是应用 后的新量化 scaler。

我们实验性地发现：

1. Round (表示为 RoundErr ) 的 Expected Error 不会改变：由于 round 函数将浮点数映射到整数, 误差大致均匀分布在 中, 平均误差为 0.25 。即 。
2. 扩展单个权重 通常不会改变组 的最大值。因此有 。
3. 由于 和 使用 FP16 表示, 因此它们没有量化误差。

因此，式 1 和 2 的量化误差可以表示为：

因此，新误差与原始误差之比为 。给定 ，因此 Salient Weight 的相对误差较小。

为了验证这个想法，作者将 OPT-6.7B 模型 1% 的 Salient Channel 与 相乘，并测量每组 的变化，如图5所示。作者发现 Scaling Salient channels 非常有效：对于 ，困惑度相比 (简单 RTN) 时从 23.54 提高到 11.92。随着 的增大, 变化 的百分比变大, 但对于 时百分比变化仍然非常小（低于5%）；随着 的增加，显著通道的相对误差继续变小。尽管如此，最好的 PPL 实际上出现在 。这是因为如果我们使用非常大的 ，当 增加时，它将增加 Non-salient channels 的相对误差（Non-salient channels 的误差将被放大 ，并且 下 21.2 的 channels 的比率大于 1, 这可能会损害模型的整体精度。因此, 在保护 Salient Channels 时，还需要考虑 Non-salient channels 的错误。

搜索 Scaling

为了同时考虑 Salient Weight 和 Non-salient Weight，作者选择自动搜索最优 (每个输入通道) 的 scaling，使某一层量化后的输出差异最小化。形式上，希望优化以下目标：

这里 表示权重量化函数（比如，组大小为 128 的 INT3/INT4 量化）， 是 FP16 中的原始权重, 是从一个小的校准集缓存的输入特征 (从预训练数据集中获取一个小的校准集, 以便不过拟合特定任务)。s 是 input channel 的缩放因子。对于 , 它通常可以融合到前一个算子中 。因量化函数是不可微的, 故无法用反向传播直接优化问题。一些技术依赖于近似梯度 , 作者发现它仍然存在收敛不稳定的问题。

为了使整个过程更加稳定，作者通过分析影响比例因子选择的因素来定义最优 Scaling 的搜索空间。如上一节所示，权重通道的显著性实际上是由 activation 的尺度决定的 (因此称之为 activation-awareness)。因此，作者简单地使用非常简单的搜索空间：

其中, 是 activation 的平均幅值 (每个 channel), 作者使用超参数 来平衡 Salient channels和 Non-salient channels 的保护。可以通过在 区间上的快速 grid search 找到最佳 (0 表示不缩放；1 对应于搜索空间中最激进的缩放)。进一步应用权重裁剪来最小化量化的 MSE 误差。作者在图6中提供了 INT3-g128 量化下 OPT 模型的消融实验结果；AWQ 始终优于 Round-ToNearest Quantization (RTN)，并实现了与混合精度 (1% FP16) 相当的性能，同时更硬件友好。

优势

AWQ 不依赖于许多 QAT 方法所需要的回归或者反向传播操作。它对校准集的依赖最小，因为只测量每个 channel 的平均大小，从而防止过拟合。因此，AWQ 需要更少的量化过程得数据，并且可以保留 LLM 在校准集分布之外的知识。

1.4 TinyChat：将 AWQ 映射到边缘设备

AWQ 可以大大减少 LLM 的尺寸。然而，将 W4A16 (4-bit weight、16-bit activation) 量化的理论内存节省转换为可以实测的加速并非易事。其替代的 W8A8 量化方法，如 SmoothQuant [1] ，在存储和计算方面保持相同的数据精度。这允许去量化过程无缝集成到计算 Kernel 的目录中。另一方面，W4A16 量化采用不同的数据类型进行显存访问和计算。因此，必须将其去量化合并到主要计算循环中以获得最佳性能，带来了实施的挑战。为了解决这个问题，作者引入了 TinyChat：一个用于 AWQ 模型推理的系统。它拥有 PyTorch 前端和后端，利用特定于设备的指令集 (例如 CUDA/PTX、Neon、AVX)。

为什么 AWQ 有助于 On-device LLM？

为了了解量化 LLM 在边缘设备的加速机会，作者首先在 RTX 4090 GPU 上分析 LLaMA-7B 模型的延时分解。作者采用 Batch Size 为 1，满足边缘用例，并使用 NVIDIA FasterTransformer 以 FP16 精度实现模型。

上下文与生成时延

如图7所示，在总结提示时，生成 20 个 token 需要 310ms，但是总结 200 个 token 仅需 10ms。因此，生成阶段比上下文阶段慢得多，特别是对于 On-device 的交互应用程序。

生成阶段是 Memory-bound 的

为了加速生成，作者在图8中进行了屋顶线分析。4090 GPU 的峰值计算吞吐量为 165 TFLOPS，内存带宽为 1TB/s。因此，任何算术强度 (FLOPs 与内存访问的比率) 小于 165 的工作负载对于 4090 GPU 来讲都是 Memory-bound 的。值得注意的是，以 FP16 运行时，On-device LLM 的生成阶段的算术强度约为1。这强调了工作负载是 Memory-bound 的。由于给定模型的 FLOPs 是固定的，提高峰值性能的唯一方法是减少内存流量的总数。AWQ 将权重内存减少了 4 倍。

权重访问主导了 Memory Traffic

作者进一步分解了 weight 和 activation 的 Memory Access，如图9所示。显然，weight access 主导了 On-device LLM 的 Memory Traffic。将模型权重量化为 4-bit 整数将近似将算术强度增加到 4 个 FLOP/Byte，导致 4TFLOPS 峰值性能。由于仅权重量化导致权重位宽较低 (因此理论性能上限较高)，因此 AWQ 很自然地遵循 On-device LLM 应用程序。

使用 TinyChat 部署 AWQ

作者证明了 4-bit 权重量化可以带来 4 倍的理论峰值性能。作者进一步设计了 TinyChat 来实现这种加速。在 GPU 上，只关注实现基本组件，包括注意力、Layer Norm 和线性投影 Kernel。灵活的前端可以轻松定制和快速支持新模型。与 Huggingface 的 FP16 实现相比，4-bit AWQ 的 TinyChat 实现了3倍以上的加速。在 CPU 上，将整个计算图降低到 C++ 以最小化开销。

对于量化层，由于硬件不提供 INT4 和 FP16 之间的乘法指令，因此需要在执行矩阵计算之前把整数 dequantize 为 FP16。作者通过将 Dequantize Kernel 与矩阵乘法 Kernel 融合来避免将 dequantized 的权重写入 DRAM。请注意，这种融合用于 matrix-matrix (MM) 乘积和 matrix-vector (MV) 乘积的 Kernel。

Kernel Fusion

作者还应用 Kernel Fusion 来优化 On-deviece 的 LLM 推理。对于 LN，将所有运算符 (例如乘法、除法和平方根) 融合到单个 Kernel 中。对于 Attention 层，将 QKV 投影融合到单个 Kernel 中，并执行动态位置编码计算。作者还预分配 KV cache 并在 Attention Kernel 中执行缓存更新。Kernel Fusion 对于前向传播实现效率低下的模型特别有用，例如 Falcon 和 StarCoder。值得注意的是，每个 FP16 Kernel 的计算时间在 4090 GPU 上约为 0.01ms，与 GPU Kernel 启动开销相当。因此，通过 Kernel Fusion 减少 Kernel 调用的数量会导致直接加速。

1.5 与经典工作 SmoothQuant 的关系

与经典工作 SmoothQuant 的相同点：

1. 都是后训练量化 (Post-Training Quantization PTQ)。
2. 都有对一些 weight (及其对应的 input activation) 做 scaling，即 weight 乘以一个 scaling factor，对应的 input activation 除以这个 scaling factor。
3. 都需要校准集确定 scaling factor 的值 (无需额外的训练)。
4. 都来自著名的 MIT HAN LAB 团队，很经典，很厉害。

与经典工作 SmoothQuant 的区别：

1. 量化精度不同： SmoothQuant 量化精度为 W8A8；AWQ 量化精度为 W4A16。
2. Scaling factor 的确定方法不同: SmoothQuant 的 scaling factor 是算出来的: ; AWQ 的 scaling factor 是搜出来的 , 其中 是 activation 的平均幅值。
3. Scaling factor 施加的 weight 不同： SmoothQuant 平等地对每个 weight (及其对应的 input activation) 做 scaling；AWQ 只对少量 (约 0.1%) 的 salient weight (及其对应的 input activation) 做 scaling。
4. 实验不同： AWQ 还开发了推理框架 TinyChat，作为后来的工作，更加完备了。

1.6 实验设置

量化设置

AWQ 专注于 Weight-only 的 Group Quantization。如之前的工作 [6]