大模型量化技术原理:AWQ 与 AutoAWQ
近年来,随着 Transformer、MOE 架构的提出,深度学习模型轻松突破上万亿规模参数,导致模型变得越来越大。为了降低模型部署成本并提升推理性能,我们需要一些大模型压缩技术。模型压缩主要分为以下几类:
- 剪枝(Pruning)
- 知识蒸馏(Knowledge Distillation)
- 量化
本系列将针对大模型的一些常见训练后量化方案(GPTQ、LLM.int8()、SmoothQuant、AWQ 等)进行讲述。
背景
将 LLM 进行低比特权重量化可以节省内存,但却很难实现。量化感知训练(QAT)由于训练成本较高并不实用,而训练后量化(PTQ)在低比特场景下面临较大的精度下降。
最接近的工作是 GPTQ,它使用二阶信息来进行误差补偿,但它可能在重建过程中过拟合校准集,从而扭曲分布之外领域上的学习特征,这可能会出现问题,因为 LLM 是通才模型。
因此,作者提出了一种'激活感知权重量化(Activation-aware Weight Quantization,AWQ)'方法,这是一种对硬件友好的低比特 LLM 仅权重化方法。该方法源于'权重对于 LLM 的性能并不同等重要'的观察,存在约 0.1%-1% 显著权重对大模型性能影响太大,通过跳过这 1% 的显著权重(salient weight)不进行量化,可以大大减少量化误差。
尽管我们只做了权重量化,但要找到显著的权重通道,我们应该根据激活分布而不是权重分布。与较大激活幅度 (activation magnitudes) 相对应的权重通道更加突出,因为它们处理了更重要的特征。
为了避免硬件效率低下的混合精度实现,我们分析了权重量化产生的误差,并推导出放大显著通道(salient channels)可以减少其相对量化误差。根据这一直觉,我们设计了一种按通道缩放的方法,以自动搜索最优缩放(scaling),使全部权重下的量化误差最小。AWQ 不依赖于任何反向传播或重构,因此可以很好地保持 LLM 在不同领域和模态上的泛化能力,而不会过拟合校准集。此外,我们还实现了一个高效的服务框架,将 AWQ 理论上节省的内存转换为实际的加速。我们的框架利用 kernel 融合的优势,最大限度地减少推理开销(例如,中间 DRAM 访问和 kernel 启动开销),以便我们可以更好地实现量化线性层的加速(AWQ 应用于包含大部分参数的线性层)。
为什么需要大模型量化
在大模型落地过程中,显存占用往往是最大的瓶颈之一。一个 7B 参数的 FP16 模型大约需要 14GB 显存,加上 KV Cache 和激活值,推理时可能需要 20GB 以上。这对于消费级显卡来说门槛很高。量化通过将权重从 16 位浮点数转换为 4 位整数,可以将模型体积缩小至原来的 1/4 左右。这不仅降低了存储需求,还减少了内存带宽压力,从而显著提升推理速度,特别是在 Memory-bound 的场景下。
然而,直接量化会导致精度严重损失。这是因为某些关键权重(即显著权重)对模型输出影响巨大,如果将它们强行压缩到 INT4,会引入巨大的噪声。AWQ 的核心思想就是保护这些关键权重,同时允许非关键权重被大幅压缩。
AWQ 技术原理
AWQ(AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration)是一种对大模型仅权重量化方法。通过保护更'重要'的权重不进行量化,从而在不进行训练的情况下提高准确率。
1. 通过保留 1% 的显著权重来改进 LLM 量化
由于 LLM 的权重并非同等重要,与其他权重相比,有一小部分显著权重对 LLM 的性能更为重要。因此,作者认为跳过这些显著权重不进行量化,可以在不进行任何训练的情况下,弥补量化损失造成的性能下降。
为了验证这个想法,作者测量了 INT3 量化模型的性能,同时保留了一定权重通道的比例为 FP16。确定了权重重要性的一种广泛使用的方法是查看其大小(magnitude)或 L2-norm。
但是发现跳过具有较大 Norm(基于 W)的权重通道并不能显著提高量化性能,与随机选择效果类似仅少量改进。而根据激活幅度(magnitude)选择权重可以显著提高性能,通过只保留 0.1%-1% 的较大激活对应权重通道就能显著提高量化性能,甚至能与基于重构的 GPTQ 相媲美。因此,我们认为幅度较大的输入特征通常更为重要,通过保留相应的权重为 FP16 可以保留这些特征,从而提高模型性能。
尽管将 0.1% 的权重保留为 FP16 可以在不明显增加模型大小的情况下提高量化性能,但这种混合精度数据类型会给系统实现带来困难(硬件效率低下)。因此,我们需要想出一种方法来保护重要的权重,而不将其实际保留为 FP16。
2. 通过激活感知缩放保护显著权重
作者提出了另一种方法,通过按逐通道(per-channel)缩放来减少显著权重的量化误差,这种方法不存在硬件效率低下的问题。
首先分析仅权重量化产生的误差。考虑一个权重为 w 的组/块;线性运算可写成 y = Wx,量化后的对应运算为 y = Q(W)x。量化函数定义如下:
Q(w) = Δ · Round(w / Δ)
Δ = max(|w|) / 2^(N-1)
其中,N 是量化比特数,Δ 是由绝对最大值 (abs value) 确定的量化缩放比例。
现在考虑一个权重元素 w ∈ W,如果我们将 w 与 s(s > 1)相乘,然后,再用 x 除以 s(其 idea 来源于之前的工作 SmoothQuant),我们将得到 y = Q(W)x = Q(W·s)(x/s),即:
Q(w·s) · (x / s) = Δ' · Round(ws / Δ) · x · (1 / s)
其中,Δ' 是应用 s 后的新的量化缩放(scaler)因子。
根据经验发现:
- Round(·) 的预期误差(记为 RoundErr)没有变化:由于舍入函数将浮点数映射为整数,误差大致均匀分布在 0-0.5 之间,平均误差为 0.25;
- 放大单个元素 w 通常不会改变组 w 的极值,因此 Δ' ≈ Δ;
- Q(w·s)·(x/s) 的误差可表示为 Err' = Δ' · RoundErr · (1/s),与原始误差 RoundErr 之比为 (Δ'/Δ) · (1/s)。给定 Δ' ≈ Δ 和 s > 1, 显著权重 w 的相对误差较小。
为了验证这个想法,作者将 OPT-6.7B 模型的 1% 显著通道乘以 s > 1,并测量下表中每组的 Δ 变化。发现缩放显著通道非常有效:困惑度从 s = 1(即 RTN)的 23.54 提高到 s = 2 的 11.92。
随着 s 变大,Δ 变化的百分比通常会变大,但当 s < 2 时,Δ 变化的百分比仍然很小;随着 s 的增加,显著通道的相对误差继续变小。尽管如此,最佳 PPL 实际上出现在 s = 2 时。这是因为如果我们使用非常大的 s,当 Δ 增加时,非显著通道的相对误差将会增加(非显著通道的误差将被放大)。并且在 s = 4 下,21.2% 的通道的比率大于 1,这可能会损害模型的整体精度。因此,我们在保护显著通道的同时还需要考虑非显著通道的误差。这就需要自动获取缩放比的方法,使得减少显著权重量化损失的同时也不能增加其它权重的量化损失。
为了同时考虑显著权重和非显著权重,作者选择自动搜索最佳(每个输入通道)缩放因子,使某一层量化后的输出差最小。从形式上看,我们希望优化以下目标:
s* = argmin_s L(s), L(s) = ||Q(W·s)(s^-1·X) - WX||
其中,Q 表示权重量化函数(例如,组大小为 128 的 INT3/INT4 量化),W 表示 FP16 中的原始权重,X 是从小校准集中的输入特征(我们从预训练数据集中获取小校准集,以免过拟合特定任务)。由于量化函数不可微,我们无法直接用梯度反向传播来优化问题。有一些技术依赖于近似梯度,但我们发现它仍然存在收敛不稳定的问题。
为了使这一过程更加稳定,我们通过分析影响缩放因子选择的因数,为最佳缩放比例定义了一个搜索空间。如前所示权重通道的显著性实际上是由激活比例(scale)决定的(即'激活感知')。因此,我们只需使用一个非常简单的搜索空间:
s = s_X^α, α* = argmin_α L(s_X^α)
其中,s 仅与激活 s_X 的大小有关,作者使用单个超参数 α 来平衡显著通道和非显著通道的保护。我们可以通过在 [0, 1] 区间内进行快速网格搜索(grid search)来找到最佳的 α(0 表示我们不进行缩放;1 对应于最激进的缩放)。
作者还通过最小化 MSE 误差来进一步应用权重剪裁,因为剪裁权重可以进一步帮助减少 Q(w·s)·(x/s) 中的 Δ';从而减少量化误差。
下表中提供了 INT3-g128 量化下 OPT 模型的消融研究;AWQ 始终优于 RTN 量化,并实现了与混合精度 (1% FP16) 相当的性能,同时更加硬件友好。
该方法不依赖于任何回归或反向传播,而这是许多量化感知训练方法所必需的。它对校准集的依赖最小,因为我们只测量每个通道的平均幅度(magnitude),从而防止过拟合。因此,该方法在量化过程中需要更少的数据,并且可以将 LLM 的知识保留在校准集分布之外。
AWQ 实验细节
AWQ 在不同模型家族(如:LLaMA、OPT 等)和模型大小的各种任务上优于现有工作。
同时,由于具有更好的泛化能力,它还在指令精调的 LM(如:Vicuna)和多模态 LM(如:OpenFlamingo)上实现了良好的量化性能。
AWQ 生态:AutoAWQ
AutoAWQ 是一个易于使用的 4 比特量化模型包。与 FP16 相比,AutoAWQ 将模型速度提高了 3 倍,并将对内存需求降低了 3 倍。AutoAWQ 实现激活感知权重量化 (AWQ) 算法来量化 LLM。AutoAWQ 是在 MIT 的 [LLM-AWQ] 基础上创建和改进的。
LLM 推理的 Compute-bound 与 Memory-bound
Roofline 模型是一个面向吞吐量的性能模型。计算密度为横坐标,FLOP/s(可达到的浮点性能)为纵坐标,可得出 roofline 模型图像(因图像长得像屋顶所以叫 roofline 模型)。蓝色段中,性能受限于理论带宽(即斜率,Peak GB/s)即 Memory-bound,在粉色段中,性能受限于浮点计算峰值性能(Peak GFLOP/s),即 Compute-bound。
对于小 batch sizes 的 7B 模型,我们会受到 Memory-bound。这意味着我们受到 GPU 内存带宽限制(移动内存中权重到计算核心),这本质上限制了我们每秒可以生成的 Token 数量。受 Memory-bound 使得量化模型更快,因为权重小了 3 倍,因此权重可以更快地在内存中移动。这与 Compute-bound 不同,在 Compute-bound 中,生成期间花费的主要时间是进行矩阵乘法。
在 Compute-bound 的情况下(batch sizes 较大时发生),使用 W4A16 量化模型不会获得加速,因为反量化的开销会减慢整体生成速度。发生这种情况是因为 AWQ 量化模型仅将权重存储在 INT4 中,但在推理过程中执行 FP16 操作,因此我们本质上是在推理过程中转换 INT4 -> FP16。
AutoAWQ 支持的模型
目前支持的主流模型包括:
- LLaMA-2 (7B/13B/70B)
- LLaMA (7B/13B/30B/65B)
- Mistral (7B)
- Vicuna (7B/13B)
- MPT (7B/30B)
- Falcon (7B/40B)
- OPT (125m/1.3B/2.7B/6.7B/13B/30B)
- Bloom (560m/3B/7B)
- GPTJ (6.7B)
- Aquila (7B)
- Aquila2 (7B/34B)
- Yi (6B/34B)
- Qwen (1.8B/7B/14B/72B)
- BigCode (1B/7B/15B)
- GPT NeoX (20B)
- GPT-J (6B)
- LLaVa (7B/13B)
- Mixtral (8x7B)
- Baichuan (7B/13B)
使用 AutoAWQ 进行量化
目前,Transformers 已经集成了 AutoAWQ,具体示例如下:
1. 安装环境
首先需要安装 AutoAWQ 库及其依赖项:
pip install autoawq
pip install transformers accelerate
确保你的 CUDA 版本与 PyTorch 版本兼容,AWQ 主要支持 NVIDIA GPU。
2. 使用 AutoAWQ 量化 LLM
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
model_path = "facebook/opt-125m"
quant_path = "opt-125m-awq"
quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version":"GEMM"}
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model.quantize(tokenizer, quant_config=quant_config)
在此步骤中,quant_config 定义了量化配置。w_bit 设置为 4 表示 4-bit 量化,q_group_size 为 128 表示每组 128 个权重共享一个缩放因子。version 指定了内核实现方式,GEMM 通常性能更好。
3. 修改配置文件以兼容 Transformers
为了使量化后的模型与 Transformers 兼容,我们需要修改配置文件。
from transformers import AwqConfig, AutoConfig
from huggingface_hub import HfApi
quantization_config = AwqConfig(
bits=quant_config["w_bit"],
group_size=quant_config["q_group_size"],
zero_point=quant_config["zero_point"],
version=quant_config["version"].lower(),
).to_dict()
model.model.config.quantization_config = quantization_config
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
这一步至关重要,它确保了加载模型时能够正确识别量化配置并调用相应的推理后端。
4. 加载量化后的模型
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("ybelkada/opt-125m-awq")
model = AutoModelForCausalLM.from_pretrained("ybelkada/opt-125m-awq").to(0)
text = "Hello my name is"
inputs = tokenizer(text, return_tensors="pt").to(0)
out = model.generate(**inputs, max_new_tokens=5)
print(tokenizer.decode(out[0], skip_special_tokens=True))
加载时,Transformers 会自动检测 quantization_config 并使用 AWQ 优化的推理路径,无需手动转换数据类型。
总结
本文简要介绍了诞生的 AWQ 背景和技术原理,同时以 AutoAWQ 为例进行了简单的说明。AWQ 通过激活感知的方式保护显著权重,在保持精度的同时实现了高效的 4-bit 量化。AutoAWQ 工具链的完善使得开发者可以便捷地将大模型部署到资源受限的环境中。未来,随着更多硬件对低比特算力的支持,AWQ 类技术将在边缘计算和端侧大模型应用中发挥更大作用。
在实际应用中,建议根据具体的硬件环境和业务需求选择合适的量化策略。对于显存敏感的场景,AWQ 是极佳的选择;而对于计算密集型的场景,则需要权衡反量化带来的开销。希望本文能为大模型量化技术的理解与实践提供有价值的参考。