AutoAWQ: 基于AWQ算法的4位量化推理加速工具

AutoAWQ

AutoAWQ简介

AutoAWQ是一个基于AWQ(Activation-aware Weight Quantization)算法的模型量化工具,可以将大型语言模型(LLM)量化到4位精度,同时在推理过程中获得2倍的加速。该工具由Casper Hansen开发,目前已在GitHub上开源。

AutoAWQ的主要特点包括:

实现了AWQ算法,可以将模型量化到4位精度
推理速度提升2倍,内存占用减少3倍(相比FP16)
支持多种主流大语言模型,如LLaMA、Vicuna、MPT等
使用简单,提供了易用的Python API
支持CPU和GPU推理
与Hugging Face的Transformers库兼容

工作原理

AutoAWQ基于MIT开发的原始AWQ算法进行了改进和优化。AWQ算法的核心思想是在量化过程中考虑激活值的分布,从而在降低位宽的同时尽可能保留模型的精度。

具体来说,AutoAWQ通过以下步骤实现4位量化:

分析模型权重和激活值的分布特征
根据激活值分布确定最优的量化参数
对权重进行4位量化,同时保留关键信息
针对量化后的模型进行推理优化

通过这种方法,AutoAWQ可以在大幅压缩模型大小的同时,将精度损失控制在可接受的范围内。

支持的模型

目前AutoAWQ支持多种主流的大语言模型,包括但不限于:

LLaMA/LLaMA-2 (7B/13B/70B)
Vicuna (7B/13B)
MPT (7B/30B)
Falcon (7B/40B)
OPT (125M-30B)
BLOOM (560M-7B)
Mistral (7B)
Mixtral (8x7B)
Qwen (1.8B-72B)

用户可以方便地对这些模型进行量化和加速。AutoAWQ团队也在持续增加对新模型的支持。

安装与使用

安装

AutoAWQ可以通过pip进行安装:

pip install autoawq

安装时需要注意:

要求CUDA Toolkit 11.8及以上版本
建议使用conda创建虚拟环境进行安装
torch版本需要与AutoAWQ的构建版本匹配

基本使用流程

使用AutoAWQ的一般流程如下:

量化模型

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'mistralai/Mistral-7B-Instruct-v0.2'
quant_path = 'mistral-7b-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)

# 量化
model.quantize(tokenizer, quant_config=quant_config)

# 保存量化后的模型
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

加载量化模型进行推理

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, TextStreamer

quant_path = "mistral-7b-awq"

# 加载量化模型
model = AutoAWQForCausalLM.from_quantized(quant_path, fuse_layers=True)
tokenizer = AutoTokenizer.from_pretrained(quant_path)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# 生成文本
prompt = "你好,请介绍一下自己。"
inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
outputs = model.generate(**inputs, streamer=streamer, max_new_tokens=512)

性能表现

根据官方提供的基准测试结果,AutoAWQ在多个模型上都实现了显著的加速:

Vicuna 7B (GEMV版本):
- 生成速度: 153.632 tokens/s
- 内存占用: 4.66 GB
MPT 7B (GEMM版本):
- 生成速度: 119.52 tokens/s
- 内存占用: 3.70 GB
Falcon 7B (GEMM版本):
- 生成速度: 94.793 tokens/s
- 内存占用: 4.48 GB

这些结果显示,AutoAWQ在保持较低内存占用的同时,可以实现非常快速的文本生成。

高级功能

除了基本的量化和推理功能外,AutoAWQ还提供了一些高级特性:

自定义校准数据

用户可以提供自定义的校准数据来优化量化效果:

def load_calib_data():
    # 加载自定义校准数据
    return [...] 

model.quantize(tokenizer, quant_config=quant_config, calib_data=load_calib_data())

长文本处理优化

对于长文本场景,可以调整相关参数来优化性能:

model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=load_long_text_data(),
    n_parallel_calib_samples=32,
    max_calib_samples=128,
    max_calib_seq_len=4096
)

多模态模型支持

AutoAWQ也支持LLaVa等视觉-语言模型:

from transformers import AutoProcessor

model = AutoAWQForCausalLM.from_quantized(quant_path)
processor = AutoProcessor.from_pretrained(quant_path)

# 处理图像和文本输入
inputs = processor(text, image, return_tensors='pt').to('cuda')
outputs = model.generate(**inputs, max_new_tokens=512)