AWQ 激活感知权重量化¶
尝试使用这个 笔记本 来体验 AWQ 量化!
激活感知权重量化(AWQ) 不会量化模型中的所有权重,而是保留一小部分对大语言模型性能至关重要的权重。这显著减少了量化损失,使你可以在 4 位精度下运行模型而不影响性能。
有几个库支持使用 AWQ 算法量化模型,例如 llm-awq、autoawq 或 optimum-intel。Transformers 支持加载使用 llm-awq 和 autoawq 库量化的模型。本指南将向你展示如何加载使用 autoawq 量化的模型,但加载 llm-awq 量化的模型过程类似。
确保你已经安装了 autoawq:
pip install autoawq
可以通过检查模型的 config.json 文件中的 quantization_config 属性来识别 AWQ 量化的模型:
使用 [from_pretrained()](/docs/transformers/v4.46.3/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) 方法加载量化的模型。如果你在 CPU 上加载了模型,请确保将其移动到 GPU 设备上。使用 device_map 参数指定模型放置的位置:
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "TheBloke/zephyr-7B-alpha-AWQ"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda:0")
加载 AWQ 量化的模型时,默认会将其他权重设置为 fp16 以提高性能。如果你希望以不同的格式加载这些权重,请使用 torch_dtype 参数:
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "TheBloke/zephyr-7B-alpha-AWQ"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
AWQ 量化还可以与 FlashAttention-2 结合使用,以进一步加速推理:
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("TheBloke/zephyr-7B-alpha-AWQ", attn_implementation="flash_attention_2", device_map="cuda:0")
融合模块¶
融合模块可以提高准确性和性能,并且在 Llama 和 Mistral 架构中默认支持 AWQ 模块,但你也可以为不支持的架构融合 AWQ 模块。
融合模块不能与其他优化技术(如 FlashAttention-2)结合使用。
支持的架构¶
不支持的架构¶
要为支持的架构启用融合模块,创建一个 AwqConfig,并设置参数 fuse_max_seq_len 和 do_fuse=True。fuse_max_seq_len 参数是总序列长度,应包括上下文长度和预期生成长度。为了安全起见,你可以将其设置为较大的值。
例如,要融合 TheBloke/Mistral-7B-OpenOrca-AWQ 模型的 AWQ 模块:
import torch
from transformers import AwqConfig, AutoModelForCausalLM
model_id = "TheBloke/Mistral-7B-OpenOrca-AWQ"
quantization_config = AwqConfig(
bits=4,
fuse_max_seq_len=512,
do_fuse=True,
)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config).to(0)
TheBloke/Mistral-7B-OpenOrca-AWQ 模型使用 batch_size=1 进行了基准测试,对比了有无融合模块的情况。
未融合的模块¶
| 批量大小 | 预填充长度 | 解码长度 | 预填充 token/s | 解码 token/s | 内存 (VRAM) |
|---|---|---|---|---|---|
| 1 | 32 | 32 | 60.0984 | 38.4537 | 4.50 GB (5.68%) |
| 1 | 64 | 64 | 1333.67 | 31.6604 | 4.50 GB (5.68%) |
| 1 | 128 | 128 | 2434.06 | 31.6272 | 4.50 GB (5.68%) |
| 1 | 256 | 256 | 3072.26 | 38.1731 | 4.50 GB (5.68%) |
| 1 | 512 | 512 | 3184.74 | 31.6819 | 4.59 GB (5.80%) |
| 1 | 1024 | 1024 | 3148.18 | 36.8031 | 4.81 GB (6.07%) |
| 1 | 2048 | 2048 | 2927.33 | 35.2676 | 5.73 GB (7.23%) |
融合的模块¶
| 批量大小 | 预填充长度 | 解码长度 | 预填充 token/s | 解码 token/s | 内存 (VRAM) |
|---|---|---|---|---|---|
| 1 | 32 | 32 | 81.4899 | 80.2569 | 4.00 GB (5.05%) |
| 1 | 64 | 64 | 1756.1 | 106.26 | 4.00 GB (5.05%) |
| 1 | 128 | 128 | 2479.32 | 105.631 | 4.00 GB (5.06%) |
| 1 | 256 | 256 | 1813.6 | 85.7485 | 4.01 GB (5.06%) |
| 1 | 512 | 512 | 2848.9 | 97.701 | 4.11 GB (5.19%) |
| 1 | 1024 | 1024 | 3044.35 | 87.7323 | 4.41 GB (5.57%) |
| 1 | 2048 | 2048 | 2715.11 | 89.4709 | 5.57 GB (7.04%) |
融合和未融合模块的速度和吞吐量也通过 optimum-benchmark 库进行了测试。

前向峰值内存/批量大小

生成吞吐量/批量大小
ExLlama-v2 支持¶
最近版本的 autoawq 支持 ExLlama-v2 内核,以加快预填充和解码速度。要开始使用,请首先安装最新版本的 autoawq:
pip install git+https://github.com/casper-hansen/AutoAWQ.git
通过传递带有 version="exllama" 的 AwqConfig() 开始使用:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig
quantization_config = AwqConfig(version="exllama")
model = AutoModelForCausalLM.from_pretrained(
"TheBloke/Mistral-7B-Instruct-v0.1-AWQ",
quantization_config=quantization_config,
device_map="auto",
)
input_ids = torch.randint(0, 100, (1, 128), dtype=torch.long, device="cuda")
output = model(input_ids)
print(output.logits)
tokenizer = AutoTokenizer.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.1-AWQ")
input_ids = tokenizer.encode("How to make a cake", return_tensors="pt").to(model.device)
output = model.generate(input_ids, do_sample=True, max_length=50, pad_token_id=50256)
print(tokenizer.decode(output[0], skip_special_tokens=True))
pip install intel-extension-for-pytorch
pip install git+https://github.com/casper-hansen/AutoAWQ.git
通过传递带有 version="ipex" 的 AwqConfig() 开始使用:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig
quantization_config = AwqConfig(version="ipex")
model = AutoModelForCausalLM.from_pretrained(
"TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ",
quantization_config=quantization_config,
device_map="cpu",
)
input_ids = torch.randint(0, 100, (1, 128), dtype=torch.long, device="cpu")
output = model(input_ids)
print(output.logits)
tokenizer = AutoTokenizer.from_pretrained("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ")
input_ids = tokenizer.encode("How to make a cake", return_tensors="pt")
pad_token_id = tokenizer.eos_token_id
output = model.generate(input_ids, do_sample=True, max_length=50, pad_token_id=pad_token_id)
print(tokenizer.decode(output[0], skip_special_tokens=True))
请注意,此功能支持 Intel CPU。