AutoGPTQ: 一个用户友好的大型语言模型量化工具包

AutoGPTQ简介

AutoGPTQ是一个基于GPTQ算法的大型语言模型(LLM)量化工具包,旨在为用户提供简单易用的API来对LLM进行量化和推理。该项目的主要目标是让LLM的量化过程变得更加简单和高效,从而使更多人能够利用量化技术来降低LLM的资源需求。

主要特点

1. 用户友好的API:AutoGPTQ提供了简洁明了的API,使用户能够轻松地对模型进行量化和推理。

2. 基于GPTQ算法:采用了高效的GPTQ(Generative Pre-trained Transformer Quantization)算法来进行模型量化。

3. 支持多种模型架构:可以对各种主流的Transformer模型进行量化,如BERT、GPT、OPT等。

4. 高效的推理:量化后的模型可以显著降低内存占用和计算需求,同时保持较高的推理性能。

5. 与Hugging Face生态系统集成:可以无缝地与Transformers库和Hugging Face Hub集成使用。

安装方法

AutoGPTQ支持在Linux和Windows平台上安装。用户可以通过pip使用预构建的轮子来安装最新的稳定版本:

# CUDA 11.8
pip install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/

# CUDA 12.1
pip install auto-gptq --no-build-isolation

# ROCm 5.7
pip install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/rocm573/

# Intel Gaudi 2 AI加速器
BUILD_CUDA_EXT=0 pip install auto-gptq --no-build-isolation

如果需要使用Triton后端,可以通过以下命令安装(目前仅支持Linux,不支持3位量化):

pip install auto-gptq[triton] --no-build-isolation

快速上手

以下是使用AutoGPTQ进行模型量化和推理的简单示例:

from transformers import AutoTokenizer, TextGenerationPipeline
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
import logging

logging.basicConfig(
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s", 
    level=logging.INFO, 
    datefmt="%Y-%m-%d %H:%M:%S"
)

pretrained_model_dir = "facebook/opt-125m"
quantized_model_dir = "opt-125m-4bit"

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)

# 准备量化配置
quantize_config = BaseQuantizeConfig(
    bits=4,  # 量化为4-bit
    group_size=128,  # 推荐设置为128
    desc_act=False,  # 设置为False可以显著加快推理速度,但可能略微影响困惑度
)

# 加载未量化的模型
model = AutoGPTQForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)

# 量化模型
examples = [
    tokenizer(
        "auto-gptq is an easy-to-use model quantization library with user-friendly apis, based on GPTQ algorithm."
    )
]
model.quantize(examples)

# 保存量化后的模型
model.save_quantized(quantized_model_dir)

# 加载量化后的模型进行推理
model = AutoGPTQForCausalLM.from_quantized(quantized_model_dir, device="cuda:0")

# 使用模型生成文本
print(tokenizer.decode(model.generate(**tokenizer("auto_gptq is", return_tensors="pt").to(model.device))[0]))

# 或者使用pipeline
pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
print(pipeline("auto-gptq is")[0]["generated_text"])

性能对比

AutoGPTQ在推理速度和内存占用方面都表现出色。以下是一些基准测试结果:

这些结果显示,使用GPTQ量化后的模型在推理速度上通常优于FP16模型,同时还能显著降低内存占用。

支持的模型

AutoGPTQ支持多种主流的Transformer模型架构,包括但不限于:

• BLOOM
• GPT-2
• GPT-NeoX
• GPT-J
• LLaMA
• OPT
• MOSS
• Falcon

结语

AutoGPTQ为大型语言模型的量化提供了一个简单而强大的解决方案。通过使用这个工具包,研究人员和开发者可以更容易地部署和使用大型语言模型,即使在计算资源有限的环境中也能获得良好的性能。随着项目的不断发展,我们期待看到更多创新应用和改进。

欢迎访问AutoGPTQ的GitHub仓库了解更多详情,参与项目开发,或提出宝贵的建议和反馈。