大模型量化之AWQ原理和应用

本文将介绍如下内容：

• 一、什么是 AWQ
• 二、AWQ 核心原理
• 三、搭建 AWQ 环境
• 四、大模型量化转换
• 五、量化后模型测试

一、什么是 AWQ

AWQ（Activation-aware Weight Quantization）量化是一种基于激活值分布(activation distribution)挑选显著权重(salient weight)进行量化的方法，其不依赖于任何反向传播或重建，因此可以很好地保持LLM在不同领域和模式上的泛化能力，而不会过拟合到校准集，属训练后量化(Post-Training Quantization, PTQ)大类。

• 论文地址： AWQ: ACTIVATION-AWARE WEIGHT QUANTIZATION FOR
  ON-DEVICE LLM COMPRESSION AND ACCELERATION
• github 地址：
  • llm-awq
  • AutoAWQ

二、AWQ 核心论述

1、核心观点1：权重并不同等重要，仅有小部分显著权重对推理结果影响较大

作者指出，模型的权重并不同等重要，仅有0.1%～1%的小部分显著权重对模型输出精度影响较大。 因此如果能有办法只对0.1%~1%这一小部分权重保持原来的精度(FP16)，对其他权重进行低比特量化，就可以在保持精度几乎不变的情况下，大幅降低模型内存占用，并提升推理速度。这就涉及到一个问题，如何鉴别显著权重，常用的方法有三种：

• 随机挑选：听天由命，随机选出0.1%~1%的权重作为显著权重，当然这种方法很不科学。
• 基于权重分布挑选：对权重矩阵(比如自注意力中的 $W_q,W_k,W_v$ )中的元素按绝对值大小由大到小排序，绝对值越大越显著，选择前0.1%~1%的元素作为显著权重。
• 基于激活值分布挑选: 作者用“激活值”一词很具有迷惑性，查阅了很多blog都没太确定其具体含义，阅读源码后确定，激活值就是与权重矩阵作matmul运算的输入值，比如自注意力机制中，计算 $Q=W_{q}X,K=W_{k}X,V=W_{v}X,O=W_o\left[\text{softmax}(K{Q}^{\top })V\right]$，$X$ 就是 $W_q,W_k,W_v$的激活值，$\left[\text{softmax}(K{Q}^{\top })V\right]$ 是 $W_o$ 的激活值。按激活值绝对值大小由大到小排序，绝对值越大越显著，选择前0.1%~1%的元素作为显著权重。

作者分别对三种方法进行了实验，将PPL(越小越好)作为评估指标，发现：

• 随机挑选显著权重(random列)的方式与对所有权重进行低比特量化(RTN列)的性能差不多，预料之内；
• 基于权重分布挑选(base on W列)的方式与随机挑选(random列)差不多，这是打破常规认知的，因为很多量化方法都是基于该方式挑选显著权重；
• 基于激活值分布挑选(based on act.列)的方式相比FP16几乎没有精度损失，是重大发现。
  

因此，基于激活值分布挑选显著权重是最为合理的方式。只要把这部分权重保持FP16精度，对其他权重进行低比特量化，就可以在保持精度几乎不变的情况下，大幅降低模型内存占用，并提升推理速度。

这里还有一个细节，作者为了避免方法在实现上过于复杂，在挑选显著权重时，并非在“元素”级别进行挑选，而是在“通道(channel)”级别进行挑选，即权重矩阵的一行作为一个单位。在计算时，首先将激活值对每一列求绝对值的平均值，然后把平均值较大的一列对应的通道视作显著通道，保留FP16精度。对其他通道进行低比特量化，如下图：

但另一个问题随之而来，如果权重矩阵中有的元素用FP16格式存储，有的用INT4格式存储，不仅存储时很麻烦， 计算时取数也很麻烦，kernel函数写起来会很抽象。于是，作者想了一个变通的方法——Scaling。

2、核心观点2：量化时对显著权重进行放大可以降低量化误差

考虑一个权重矩阵 $w$，线性运算可以写作 $y=wx$。对权重矩阵进行量化后，可以写作 $y=Q(w)x$，量化函数${}^{\star }$ $Q(\cdot )$ 定义为公式1：

• $Q(w)=\Delta \cdot \text{Round}\left(\frac{w}{\Delta }\right),\Delta =\frac{\max (|w|)}{2^{N-1}},$ (1)

其中 $N$ 是量化后的比特数，$\Delta$ 是量化因子${}^{\star }$ (scaler)。 $w^{\prime }=\text{Round}\left(\frac{w}{\Delta }\right)$ 是量化过程，$\Delta \cdot w^{\prime }$ 是反量化过程。
原始的 $w$、$\Delta$ 和输入 $x$ 都是 FP16 格式，不会带来精度损失。 整个过程的精度损失全部来源于量化过程中的 Round 取整函数${}^{\star }$，其误差近似成 $[0,0.5]$ 的均匀分布，期望为 $0.25$，可以写作 $\text{RoundErr}(\cdot )\sim 0.25$。

考虑对于权重矩阵 $w$ 中的单个元素 $w$，引入一个 缩放因子${}^{\star }$ $s>1$，量化过程将 $w$ 与该因子相乘，写作 $w^{\prime }=\text{Round}\left(\frac{ws}{{\Delta }^{\prime }}\right)$，相应地将反量化过程写作 $\frac{{\Delta }^{\prime }\cdot w^{\prime }}{s}$，这样在计算过程中是“等价”的，如公式2：

• $Q(w\cdot s)\cdot \frac{x}{s}={\Delta }^{\prime }\cdot \text{Round}\left(\frac{ws}{{\Delta }^{\prime }}\right)\cdot x\cdot \frac{1}{s},$ (2)

虽然公式1和公式2在计算过程上是“等价”的，但是带来的精度损失是不一样的。两种计算方法的误差可以写作公式3：

• $\text{Err}(Q(w)x)=\Delta \cdot \text{RoundErr}\left(\frac{w}{\Delta }\right)\cdot x$ （3）
• $\text{Err}\left(Q(w\cdot s)\left(\frac{x}{s}\right)\right)={\Delta }^{\prime }\cdot \text{RoundErr}\left(\frac{ws}{{\Delta }^{\prime }}\right)\cdot x\cdot \frac{1}{s}$ （3）

$\text{RoundErr}(\cdot )$ 视作常数 0.25，公式 2 的误差与公式 1 的误差的比值可以表示为 $\frac{{\Delta }^{\prime }}{\Delta }\cdot \frac{1}{s}$。

通常情况下，我们对权重矩阵 $w$ 中的某几个元素 $w$，乘以缩放因子 $s$ 后，大概率是不会影响权重矩阵 $w$ 的元素最大值的，
除非乘的 $w$ 本身就是最大值，但这种概率很小。进而，根据公式 1 中 $\Delta$ 的计算方法，我们认为 ${\Delta }^{\prime }\approx \Delta$。
加上 $s>1$，公式 2 与公式 1 的误差比值 $<1$，于是作者提出一个观点： 量化时对显著权重进行放大，可以降低量化误差。

这种理论是否正确？作者进行了实验，如下表，随着 $s$ 的增大，${\Delta }^{\prime }\ne \Delta$ 的概率由 0 不断升高，
但在 $s<2$ 之前，概率还是很低的(<5%)；同时，在一定范围内，随着 $s$ 的增大，误差比值越来越小，这是完全支持作者观点的。

因此，作者改变了思路：为了更加hardware-friendly，我们对所有权重均进行低比特量化，但是，在量化时，对于显著权重乘以较大的
$s$ ，相当于降低其量化误差；同时，对于非显著权重，乘以较小的$s$，相当于给予更少的关注。这便是上一节提到的缩放(Scaling)方法。

3、算法：自动计算缩放(scaling)系数

按照上文的分析，我们需要找到权重矩阵每个通道的缩放系数$s$，使得量化误差最小，即最小化公式4：

• $s^{\ast }=\arg \underset{s}{\min }\mathcal{L}(s)$ （4）

• $\mathcal{L}(s)=\Vert Q(\mathbf{W}\cdot \text{diag}(s))(\text{diag}(s{)}^{-1}\cdot \mathbf{X})-\mathbf{WX}\Vert$ （4）

但是，量化函数 $Q(\cdot )$ 是不可微的，我们无法使用梯度下降法优化求解上式。同时，上式是非凸的，使得该问题的求解更加困难。

按照作者的观点，激活值越大，对应通道越显著，就应该分配更大的缩放系数降低其量化误差。因此，作者统计了各通道的平均激活值（计算输入矩阵各列绝对值的平均值） $s_x$，并直接将此作为各通道的缩放系数。同时引入一个变量 $\alpha$ 用于平衡显著通道和非显著通道的系数，由此，问题转化为：

• $\mathbf{s}=\mathbf{s}{\mathbf{x}}^{\alpha },{\alpha }^{\ast }=\arg \underset{\alpha }{\min }\mathcal{L}(\mathbf{s}{\mathbf{x}}^{\alpha })(5)$

这里有个细节作者没有提，通过阅读源码发现，为了防止 $s$ 过大或者过小，作者还进行了一步数据标准化：

(图：llm-awq/awq/quantize/auto_scale.py第129-130行)
注意在求 $\alpha$ 时，会参照超参数group_size对通道进行分组，每组共享一个 $\alpha$ 。

现在的问题就转换成了找最合适的 $\alpha$ ，只优化一个数肯定比公式4中优化一堆缩放系数容易。作者在论文中提到，采用一种称为“Fast Grid Search”的方法在[0, 1]区间快速找出最合适的 $\alpha$ ，但这种方法到底是怎么做的，文中并没有提。

通过阅读源码，发现该方法实际上就是在[0,1]区间平均取20个数，0, 0.05, 0.10, 0.15 …… 然后逐个计算不同 $\alpha$ 下的MSE损失，损失最小的就是最佳的 $\alpha$ 。得到最佳 $\alpha$ 后，最佳缩放系数 $s$ 也随之确定。.

(图：llm-awq/awq/quantize/auto_scale.py第123-130行, for循环就是fast grid search算法)
将所有权重矩阵的 $s$ 保存下来，供推理阶段反量化使用。

4、算法超参数

• group-size，分组大小，请结合“算法步骤”部分理解。对于激活值 $X\in [\text{batch},\text{seq\_len},\text{in\_channel}]$ 和权重 $W\in [\text{out\_channel},\text{in\_channel}]$（$Y=X{W}^{\top }$），对 $\text{in\_channel}$ 维度进行分组。
  • group-size 较大时，每组权重较多，量化参数较少，可能会降低量化后的模型精度，但计算和存储成本低。
  • group-size 较小时，每组权重较少，量化参数较多，模型精度可能更高，但计算和存储成本增加。
  • 业界常用默认值 128。

三、搭建 AWQ 环境

这里推荐使用 https://github.com/xorbitsai/inference 开源项目的环境。

1、匹配 xprobe/xinference 与 Cuda 对应版本

在 Docker Hub 中 找到 Cuda 对应版本。

2、搭建 Docker Container 环境

#!/bin/bash

# script parameters
docker_run_name=$1
docker_image_name=$2
docker_image_version=$3
ssh_port=$4
jupyter_port=$5
tensorboard=$6

# docker run --gpus all -itd \
docker run --gpus all --cpus 48 --shm-size 16G --memory 500gb -itd \
  -p ${ssh_port}:22 \
  -p ${jupyter_port}:8888 \
  -p ${tensorboard}:6006 \
  -v /etc/localtime:/etc/localtime \
  -v /data:/data \
  -v /nasdata:/nasdata \
  --name ${docker_run_name} \
  -e JUPYTER_PASSWORD="123456" \
  -e JUPYTER_PORT=8888 \
  -e ROOT_PASS="123456" \
  ${docker_image_name}:${docker_image_version}

# bash create_vllm_quantization_container.sh njh_vllm_quantization_medical xprobe/xinference  v1.2.0  20004 20005 20006

四、大模型量化实现

1、大模型量化转换

vim AWQ.py

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
import json

# Specify paths and hyperparameters for quantization
model_path = "/nasdata/zhanjie/models/Qwen2.5-32B-Instruct"
quant_path = "Outputs-Qwen2.5-32B-Instruct"
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load your tokenizer and model with AutoAWQ
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto", safetensors=True)

# Load Dataset
with open('dataset.json', 'r') as file:
    datasets = json.load(file)

data = []
for dataset in datasets:
    msg = dataset[0]["content"] + "\n" + dataset[1]["content"]
    text = tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=False)
    data.append(text.strip())
##只需通过一行代码运行校准过程：
model.quantize(tokenizer, quant_config=quant_config, calib_data=data)

##保存量化模型：
model.save_quantized(quant_path, safetensors=True, shard_size="4GB")
tokenizer.save_pretrained(quant_path)

dataset.json 文件数据 format 如下：

[
	[  
	    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},  
	    {"role": "user", "content": "Tell me who you are."},  
	    {"role": "assistant", "content": "I am a large language model named Qwen..."}  
    ],
    [  
	    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},  
	    {"role": "user", "content": "Tell me who you are."},  
	    {"role": "assistant", "content": "I am a large language model named Qwen..."}  
    ],
    ...
]  

执行脚本

CUDA_VISIBLE_DEVICES=2,3,4,5,6,7 python3 AWQ.py

2、vllm 开启量化模型推理

#!/bin/bash

BASE_MODEL="/home/xxx"
SERVED_MODEL_NAME="qwen2.5-32b"

CUDA_VISIBLE_DEVICES=1,2 python3 -m vllm.entrypoints.openai.api_server \
    --model $BASE_MODEL \
    --served-model-name $SERVED_MODEL_NAME \
    --trust-remote-code \
    --tensor-parallel-size 2 \
    --max-model-len 24576 \
    --quantization awq \
    --host 0.0.0.0 \
    --port 6006

五、量化后模型测试

1、推理验证代码

from openai import OpenAI

api_key = "EMPTY"
api_base = "http://127.0.0.1:xxxx/v1"
client = OpenAI(api_key=api_key, base_url=api_base)

# 创建聊天完成请求，并启用流式输出
completion = client.chat.completions.create(
    model="qwen2.5-32b",
    # extra_body={"lora_name": "adapter1"},  # 指定 LoRA 适配器名称
    messages=[
        {"role": "system", "content": """你是一个乐于助人、尊重他人、诚实的INTP-T AI助理，名叫巴迪。你在和一个人类用户说话。
在安全的情况下，总是尽可能有帮助地、合乎逻辑地回答。你的回答不应包括任何有害、政治、宗教、不道德、种族主义、性别歧视、有毒、危险或非法的内容。请确保你的回应是社会公正和积极的。
如果一个问题没有任何意义，或者事实上不连贯，解释原因，而不是回答不正确的问题。如果你不知道问题的答案，请不要分享虚假信息。
你可以流利地说多种语言，例如：英语、汉语。"""},
        {"role": "user", "content": "给我写一首李白的行路难，谢谢"},
    ],
    max_tokens=2048,
    temperature=0.6,
    top_p=0.9,
    n=1,
    stream=True  # 启用流式输出
)

# 逐步处理并打印生成的内容
for chunk in completion:
    content = chunk.choices[0].delta.content
    if content:
        print(content, end='', flush=True)

2、性能验证

Todo

参考：

• 【LLM模型量化】-官方教程-qwen官方教程v2.0
• AWQ量化 – LlaMA3.1-8B 在 MLP 的最佳实践
• 深入理解AWQ量化技术