AWQ 激活感知权重量化详解:4-Bit 量化与 vLLM 部署

讲解 AWQ(Activation-aware Weight Quantization):依据激活分布保护关键权重通道的硬件友好低位纯权重量化,相比 FP16 提速 3 倍、显存降至三分之一,并给出 AutoAWQ 量化与 vLLM 部署的代码示例。

| **[GitHub](https://github.com/casper-hansen/AutoAWQ) | [Paper ](https://arxiv.org/abs/2306.00978)| [AutoAWQ](https://casper-hansen.github.io/AutoAWQ/) |**

简介

AWQ 指的是 “激活感知权重量化”(Activation-aware Weight Quantization),这是一种对硬件友好的 LLM 低位纯权重量化方法。AutoAWQ 是一个易于使用的 4 位量化模型软件包。与 FP16 相比,AutoAWQ 可将模型速度提高 3 倍,内存需求减少 3 倍。AutoAWQ 实现了用于量化 LLM 的激活感知权重量化 (AWQ) 算法。

AWQ 有两个版本:GEMM 和 GEMV,其中:

  • GEMV(量化):在batch size = 1 时,比 GEMM 快 20%(不适合大上下文)。
  • GEMM(量化):在 batch size 低于 8 时比 FP16 快得多(适用于大型上下文)。
  • FP16(非量化):推荐用于最高吞吐量:vLLM

计算限制与内存限制

在使用 7B 小模型的小批量情况下,会受到内存限制。这意味着受限于 GPU 在内存中推送权重的带宽,而这正是限制每秒能生成多少 tokens 的根本原因。受内存限制使得量化模型的速度更快,因为权重小了 3 倍,因此在内存中的推送速度也更快。这与计算绑定不同,在计算绑定中,生成模型所花费的主要时间是进行矩阵乘法运算。

在计算受限的情况下,也就是在批量较大的情况下,使用 W4A16 量化模型并不会提高速度,因为去量化的开销会减慢整体生成速度。这是因为 AWQ 量化模型仅以 INT4 保存权重,但在推理过程中会执行 FP16 操作,因此我们在推理过程中基本上是将 INT4 转换为 FP16。

Fused 模块

融合模块是 AutoAWQ 加速的重要组成部分,将多个层组合成一个操作,从而变得更加高效。

Fused 模块代表一组与 Huggingface 模型分开工作的自定义模块。与 model.generate() 和其他

Huggingface 方法兼容,如果激活融合模块的一些限制如下:

  • 当使用 fuse_layers=True 时,融合模块将被激活。
  • 为了实现了自定义缓存。它根据 batch size 和 sequence 长度进行预分配。
  • 融合模块中的主要加速器来自FasterTransformer,仅兼容Linux。
  • model.generate() 中的 past_key_values 只是 dummy values,因此生成后无法使用。

    AWQ 量化模型与 Transformer 的结合使用

使用fuse_layers=True

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, TextStreamer

quant_path = "TheBloke/zephyr-7B-beta-AWQ"

# Load model
model = AutoAWQForCausalLM.from_quantized(quant_path, fuse_layers=True)
tokenizer = AutoTokenizer.from_pretrained(quant_path, trust_remote_code=True)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# Convert prompt to tokens
prompt_template = """\
<|system|>
</s>
<|user|>
{prompt}</s>
<|assistant|>"""

prompt = "You're standing on the surface of the Earth. "\
"You walk one mile south, one mile west and one mile north. "\
"You end up exactly where you started. Where are you?"

tokens = tokenizer(
prompt_template.format(prompt=prompt),
return_tensors='pt'
).input_ids.cuda()

# Generate output
generation_output = model.generate(
tokens,
streamer=streamer,
max_seq_len=512
)

不使用fuse_layers=True

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen1.5-7B-Chat-AWQ", # the quantized model
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B-Chat-AWQ")

prompt = "Give me a short introduction to large language model."
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(device)

generated_ids = model.generate(
model_inputs.input_ids,
max_new_tokens=512
)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

将 AWQ 量化模型与 vLLM 结合使用

vLLM 支持 AWQ,这意味着可以直接使用提供的 AWQ 模型或使用 AutoAWQ 与 vLLM 训练的模型。下面展示如何使用 vLLM 和 Qwen1.5-7B-Chat-AWQ 启动 OpenAI-API 兼容 API:

1
2
3
4
5
6
7
8
9
python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen1.5-7B-Chat-AWQ

curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "Qwen/Qwen1.5-7B-Chat-AWQ",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me something about large language models."}
],
}'

或者使用python client调用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from openai import OpenAI
# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"

client = OpenAI(
api_key=openai_api_key,
base_url=openai_api_base,
)

chat_response = client.chat.completions.create(
model="Qwen/Qwen1.5-7B-Chat-AWQ",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me something about large language models."},
]
)
print("Chat response:", chat_response)

使用AutoAWQ量化自定义模型

对于较小的 7B 模型,预计需要 10-15 分钟,对于 70B 模型,大约需要 1 小时。推荐通过安装源代码来获取并安装该工具包的最新版本:

1
2
3
git clone https://github.com/casper-hansen/AutoAWQ.git
cd AutoAWQ
pip install -e .

构建自己的AWQ量化模型,并使用训练数据进行校准:

1
2
3
4
5
6
7
8
9
10
11
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

# Specify paths and hyperparameters for quantization
model_path = "your_model_path"
quant_path = "your_quantized_model_path"
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)

需要准备数据以进行校准。将样本放入一个列表中,其中每个样本都是一段文本。由于直接使用微调数据来进行校准,所以需要使用ChatML模板对其进行格式化:

1
2
3
4
5
data = []
for msg in messages:
msg = c['messages']
text = tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=False)
data.append(text.strip())

其中每个 msg 是一个典型的聊天消息,如下所示:

1
2
3
4
5
[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me who you are."},
{"role": "assistant", "content": "I am a large language model named Qwen..."}
]

然后只需通过一行代码运行校准过程:

1
model.quantize(tokenizer, quant_config=quant_config, calib_data=data)

最后,保存量化模型,得到一个可以用于部署的AWQ量化模型:

1
2
model.save_quantized(quant_path, safetensors=True, shard_size="4GB")
tokenizer.save_pretrained(quant_path)
本文结束 感谢您的阅读