深度解析：大模型量化技术原理——AWQ与AutoAWQ

猿类崛起@

于 2024-08-31 21:50:27 发布

阅读量363

点赞数 11

文章标签：人工智能大模型教程 AI大模型大模型学习大模型大模型入门

本文链接：https://blog.csdn.net/2401_84495872/article/details/141759326

版权

近年来，随着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规模参数，从而导致模型变得越来越大，因此，我们需要一些大模型压缩技术来降低模型部署的成本，并提升模型的推理性能。模型压缩主要分为如下几类：

剪枝（Pruning）
知识蒸馏（Knowledge Distillation）
量化

本系列将针对大模型的一些常见训练后量化方案（GPTQ、LLM.int8()、SmoothQuant、AWQ等）进行讲述。

之前讲述了MIT的大模型量化方案SmoothQuant，本文针对MIT另一篇大模型量化方案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 应用于包含大部分参数的线性层）。

AWQ 技术原理

AWQ（AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration）是一种对大模型仅权重量化方法。通过保护更“重要”的权重不进行量化，从而在不进行训练的情况下提高准确率。

通过保留1%的显著权重来改进LLM量化

由于 LLM 的权重并非同等重要，与其他权重相比，有一小部分显著权重对 LLM 的性能更为重要。因此，作者认为跳过这些显著权重不进行量化，可以在不进行任何训练的情况下，弥补量化损失造成的性能下降。

为了验证这个想法，作者测量了 INT3 量化模型的性能，同时保留了一定权重通道的比例为 FP16。确定了权重重要性的一种广泛使用的方法是查看其大小（magnitude）或 L2-norm。

但是发现跳过具有较大 Norm（基于 W）的权重通道并不能显著提高量化性能，与随机选择效果类似仅少量改进。而根据激活幅度（magnitude）选择权重可以显著提高性能，通过只保留 0.1%-1% 的较大激活对应权重通道就能显著提高量化性能，甚至能与基于重构的 GPTQ 相媲美。因此，我们认为幅度较大的输入特征通常更为重要，通过保留相应的权重为 FP16 可以保留这些特征，从而提高模型性能。

尽管将 0.1% 的权重保留为 FP16 可以在不明显增加模型大小的情况下提高量化性能，但这种混合精度数据类型会给系统实现带来困难（硬件效率低下）。因此，我们需要想出一种方法来保护重要的权重，而不将其实际保留为 FP16。

通过激活感知缩放保护显著权重

作者提出了另一种方法，通过按逐通道（per-channel）缩放来减少显著权重的量化误差，这种方法不存在硬件效率低下的问题。

首先分析仅权重量化产生的误差。考虑一个权重为 w 的组/块；线性运算可写成 y=wxy=wxy=wx，量化后的对应运算为 y=Q(w)xy=Q(\mathbf{w})\mathbf{x}y=Q(w)x. 量化函数定义如下：

Q(w)=Δ⋅Round(wΔ),Δ=max⁡(∣w∣)2N−1Q(\mathbf{w}) = \Delta \cdot \text{Round}(\frac{\mathbf{w}}{\Delta}), \quad \Delta = \frac{\max(|\mathbf{w}|)}{2^{N-1}}Q(w)=Δ⋅Round(Δw),Δ=2N−1max(∣w∣)

其中，N 是量化比特数，Δ\DeltaΔ 是由绝对最大值(abs value)确定的量化缩放比例。

现在考虑一个权重元素 w∈ww\in \mathbf{w}w∈w，如果我们将 w 与 s （s>1）相乘，然后，再用 x 除以 s（其idea来源于之前的工作SmoothQuant），我们将得到 y=Q(w)x=Q(w⋅s)(x/s)y=Q(\mathbf{w})\mathbf{x}=Q(w\cdot s)(x/s)y=Q(w)x=Q(w⋅s)(x/s)，即：

Q(w⋅s)⋅xs=Δ′⋅Round(wsΔ)⋅x⋅1sQ(w\cdot s)\cdot \frac{x}{s} = \Delta^{'} \cdot \text{Round}(\frac{ws}{\Delta}) \cdot x \cdot \frac{1}{s} Q(w⋅s)⋅sx=Δ′⋅Round(Δws)⋅x⋅s1

其中，Δ′\Delta^{'}Δ′ 是应用 s 后的新的量化缩放（scaler）因子。

根据经验发现:

Round(⋅)\text{Round}(\cdot)Round(⋅) 的预期误差（记为 RoundErrRoundErrRoundErr）没有变化：由于舍入函数将浮点数映射为整数，误差大致均匀分布在 0-0.5 之间，平均误差为 0.25；
放大单个元素 w 通常不会改变组 w 的极值，因此 Δ′≈Δ\Delta^{'}\approx\DeltaΔ′≈Δ;
Q(w⋅s)⋅xs=Δ′⋅Round(wsΔ)⋅x⋅1sQ(w\cdot s)\cdot \frac{x}{s} = \Delta^{‘} \cdot \text{Round}(\frac{ws}{\Delta}) \cdot x \cdot \frac{1}{s} Q(w⋅s)⋅sx=Δ′⋅Round(Δws)⋅x⋅s1 的误差可表示为 Err′=Δ′⋅RoundErr⋅1sErr^{'}=\Delta{’} \cdot RoundErr \cdot \frac{1}{s}Err′=Δ′⋅RoundErr⋅s1，与原始误差 RoundErrRoundErrRoundErr 之比为 Δ′Δ⋅1s\frac{\Delta^{‘}}{\Delta} \cdot \frac{1}{s}ΔΔ′⋅s1. 给定 Δ′≈Δ\Delta^{’}\approx\DeltaΔ′≈Δ 和 s>1s>1s>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，当 Δ 增加时，非显著通道的相对误差将会增加（非显著通道的误差将被放大） Δ′Δ\frac{\Delta^{'}}{\Delta}ΔΔ′。并且在 s = 4 下，21.2% 的通道的比率大于 1，这可能会损害模型的整体精度。因此，我们在保护显著通道的同时还需要考虑非显著通道的误差。这就需要自动获取缩放比的方法，使得减少显著权重量化损失的同时也不能增加其它权重的量化损失。

为了同时考虑显著权重和非显著权重，作者选择自动搜索最佳（每个输入通道）缩放因子，使某一层量化后的输出差最小。从形式上看，我们希望优化以下目标：

s∗=arg min⁡sL(s),L(s)=∥Q(W⋅s)(s−1⋅X)−WX∥\mathbf{s}^* = \argmin_{\mathbf{s}} \mathcal{L}(\mathbf{s}), \quad \mathcal{L}(\mathbf{s})=\lVert Q(\mathbf{W}\cdot \mathbf{s}) (\mathbf{s^{-1}} \cdot \mathbf{X}) - \mathbf{W}\mathbf{X} \lVerts∗=sargminL(s),L(s)=∥Q(W⋅s)(s−1⋅X)−WX∥

其中，Q 表示权重量化函数（例如，组大小为 128 的 INT3/INT4 量化），W 表示 FP16 中的原始权重，X 是从小校准集中的输入特征（我们从预训练数据集中获取小校准集，以免过拟合特定任务）。由于量化函数不可微，我们无法直接用梯度反向传播来优化问题。有一些技术依赖于近似梯度，但我们发现它仍然存在收敛不稳定的问题。

为了使这一过程更加稳定，我们通过分析影响缩放因子选择的因数，为最佳缩放比例定义了一个搜索空间。如前所示权重通道的显著性实际上是由激活比例（scale）决定的（即 “激活感知”）。因此，我们只需使用一个非常简单的搜索空间：

s=sXα,α∗=arg min⁡αL(sXα) \mathbf{s}=\mathbf{s_X}^{\alpha}, \quad \alpha^{*=\argmin_{\alpha}\mathcal{L}(\mathbf{s_X}}{\alpha})s=sXα,α∗=αargminL(sXα)

其中，s仅与激活sXs_XsX的大小有关，作者使用单个超参数α来平衡显著通道和非显著通道的保护。我们可以通过在 [0, 1] 区间内进行快速网格搜索（grid search）来找到最佳的 α（0 表示我们不进行缩放；1 对应于最激进的缩放）。

作者还通过最小化 MSE 误差来进一步应用权重剪裁，因为剪裁权重可以进一步帮助减少 Q(w⋅s)⋅xs=Δ′⋅Round(wsΔ)⋅x⋅1sQ(w\cdot s)\cdot \frac{x}{s} = \Delta^{'} \cdot \text{Round}(\frac{ws}{\Delta}) \cdot x \cdot \frac{1}{s} Q(w⋅s)⋅sx=Δ′⋅Round(Δws)⋅x⋅s1 中的 Δ′；从而减少量化误差。

下表中提供了 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 模型

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 支持的模型如下表所示：

Models	Sizes
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
QWen	1.8B/7B/14/72B

目前，Transformers 已经集成了 AutoAWQ ，具体示例如下：

使用 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"}

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

为了使量化后的模型与Transformers兼容，我们需要修改配置文件。

from transformers import AwqConfig, AutoConfig
from huggingface_hub import HfApi

# modify the config file so that it is compatible with transformers integration
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()

# the pretrained transformers model is stored in the model attribute + we need to pass a dict
model.model.config.quantization_config = quantization_config

# save model weights
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

加载量化后的模型：

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))