AutoAWQ（Automatic Aware Quantization）是一个基于 AWQ（Activation-aware Weight Quantization for LLM） 算法的自动化量化工具，通过智能选择量化参数（如位宽、零点等）来优化深度学习模型的存储和计算效率，同时最大限度地保留模型精度。

GitHub：https://github.com/casper-hansen/AutoAWQ
AWQ Paper：https://arxiv.org/pdf/2306.00978

运行环境：Ubuntu 22.04 + Python 3.10 + RTX 3060 Driver 550.144.03 + CUDA 12.4 + 12G 显存

# 创建虚拟环境
conda create -n quan-env python=3.10

# 安装需要的库
pip install torch --index-url https://download.pytorch.org/whl/cu126
pip install transformers==4.46.3
pip install autoawq==0.2.8
pip install evaluate==0.4.0
pip install datasets==2.21.0

1. 准备数据

import pandas as pd
import pickle
from transformers import Qwen2Tokenizer

# 数据集：https://huggingface.co/datasets/Orion-zhen/firefly-exl-calibration

def show_info(all_text, all_size):
    print('样本数量:', len(all_size))
    print('最大长度:', max(all_size))
    print('最小长度:', min(all_size))
    print('平均长度:', int(sum(all_size) / len(all_size)))

    for idx, line in enumerate(all_text):
        print(line)
        if idx == 3:
            break
        print('-' * 100)


def demo():
    data = pd.read_parquet('calib_data/firefly.parquet')
    print('数据信息:', data.shape, data.columns)

    tokenizer = Qwen2Tokenizer.from_pretrained('Qwen2.5-7B-Instruct')

    calib_data = data[:1000]
    all_text, all_size = [], []
    for prompt, output in zip(calib_data['input'].tolist(), calib_data['output'].tolist()):

        message = [{'role': 'user', 'content': prompt},
                   {'role':'assistant', 'content': output}]
        text = tokenizer.apply_chat_template(message, add_generation_prompt=False, tokenize=False)

        if len(text) > 1000 or len(text) < 20:
            continue

        all_text.append(text)
        all_size.append(len(text))

        if len(all_size) == 500:
            break

    # 样本数量: 500
    # 最大长度: 997
    # 最小长度: 171
    # 平均长度: 341
    # 数据格式：
    # <|im_start|>system
    # You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
    # <|im_start|>user
    # 在上海的苹果代工厂，较低的基本工资让工人们形成了“软强制”的加班默契。加班能多拿两三千，“自愿”加班成为常态。律师提示，加班后虽能获得一时不错的报酬，但过重的工作负荷会透支身体，可能对今后劳动权利造成不利影响。
    # 输出摘要：<|im_end|>
    # <|im_start|>assistant
    # 苹果代工厂员工调查：为何争着“自愿”加班
    #
    # <|im_end|>
    show_info(all_text, all_size)
    pickle.dump(all_text, open('calib_data/校准数据.pkl', 'wb'))



if __name__ == '__main__':
    demo()

2. 模型量化

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
import pickle


def demo():
    # 1. 模型加载
    # torch_dtype: 设置加载模型时的 torch 数据类型。可以指定为 float16、float32 等类型，或者使用 "auto" 让框架自动选择合适的数据类型。
    # trust_remote_code: 指示是否信任来自 Hugging Face Hub 上的远程代码。如果设置为 True，则允许执行来自远程模型的代码。如果为 False，则会禁用远程执行代码的功能。
    # device_map: 如果设置为 None，则模型将加载到默认设备上(通常是 CPU)。如果为 auto，则加载到 GPU, 如果可用。也可以指定不同层加载到不同设备。注意：meta 设置是一个占位设备，表示数据实际并未加载到内存或显存中，需要的时候再进行加载。
    # low_cpu_mem_usage: 减少模型加载过程中对 CPU 内存的占用，该参数为 True 时，device_map 参数可以设置为 auto，参数延迟加载，而不是一次性加载
    # use_cache: 模型计算时是否使用缓存。
    estimator = AutoAWQForCausalLM.from_pretrained(model_path='Qwen2.5-7B-Instruct',
                                                   torch_dtype='auto',
                                                   trust_remote_code=True,
                                                   device_map=None,
                                                   low_cpu_mem_usage=True,
                                                   use_cache=False)
    # 打印参数信息: 参数名 参数类型 计算设备
    # for name, param in estimator.named_parameters():
    #     print(f"Layer: {name}, dtype: {param.dtype}, device: {param.device}")
    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path='Qwen2.5-7B-Instruct',
                                              trust_remote_code=True)

    # 2. 模型量化
    # 监控显存：watch -n 1 nvidia-smi
    # calib_data: 校准数据路径
    # max_calib_samples: 如果你的校准数据集非常大，使用全部数据进行量化可能会导致计算开销过大，或者内存消耗过高。因此，使用 max_calib_samples 参数可以限制使用的样本数量。
    # n_parallel_calib_samples: None 表示所有的校准样本会一次性送入模型进行计算，可能导致内存溢出。
    # max_calib_seq_len: 输入文本长度，超出则丢弃
    quant_config = {'q_group_size': 128, 'w_bit': 4}
    # w_bit: 权重的量化位宽，决定了每个权重值的存储大小。较低的位宽可以大幅减少模型的内存占用，但可能会降低精度。
    # q_group_size: 指定量化时每组的大小。通常在量化时，参数会分组处理，这个参数控制每个组包含多少元素。

    # 读取校准数据
    calib_data = pickle.load(open('calib_data/校准数据.pkl', 'rb'))
    estimator.quantize(tokenizer=tokenizer,
                       calib_data=calib_data,
                       max_calib_samples=128,
                       n_parallel_calib_samples=None,
                       split='validation',
                       max_calib_seq_len=1024,
                       text_column='text',
                       quant_config=quant_config)

    # 3. 模型存储
    # AWQ: 100%|██████████████████████████████████████| 28/28 [21:53<00:00, 46.91s/it]
    save_path = 'Qwen2.5-7B-Instruct-AWQ'
    estimator.save_quantized(save_path)
    tokenizer.save_pretrained(save_path)
    print(f'模型量化结束，并存储在 "{save_path}"')


if __name__ == '__main__':
    demo()

3. 模型推理

from awq import AutoAWQForCausalLM
from awq.utils.utils import get_best_device
from transformers import AutoTokenizer
from transformers import TextStreamer


# from transformers import Qwen2Tokenizer
# from transformers import Qwen2ForCausalLM
# Qwen2ForCausalLM().generate()


def demo():
    device = get_best_device()
    # 1. 模型加载
    # quant_path: 模型文件所在的路径
    # max_seq_len: 模型能够处理的最大序列长度，默认值为2048。
    # fuse_layers: 在加载时是否尝试合并某些层以提高性能，默认值为True。
    # use_exllama: ExLlama是一个优化的Transformer推理库，默认值为False。
    # use_ipex: 该选项用于启用 Intel 优化的处理，通常是为加速推理而使用，默认值为False。
    # batch_size: 指定模型每次推理时处理的样本数量，默认值为1。
    # safetensors: Safetensors是一种模型权重格式，提供了比常规格式更高的安全性，默认为True。
    # device_map: 指定如何分配模型到不同的设备上，常见值包括 balanced（默认值,均衡分配) 或 auto(自动选择)
    # max_memory: 指定可以使用的最大内存量，常用于多GPU的分配策略。如果为None，则会使用所有可用内存。
    estimator = AutoAWQForCausalLM.from_quantized(quant_path='Qwen2.5-7B-Instruct-AWQ')
    tokenizer = AutoTokenizer.from_pretrained('Qwen2.5-7B-Instruct-AWQ', trust_remote_code=True)

    # 2. 模型推理
    # 2.1 构建输入
    prompt = [{'role': 'system', 'content': '你是一个私人智能助手，你的名字叫阿亮。'}, {'role': 'user', 'content': '我想知道宋江是谁？'}]
    # tokenize=False 表示只构造输入，并不会转换为 ids
    # tokens = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
    # print(tokens)
    inputs = tokenizer.apply_chat_template(conversation=prompt,
                                           tokenize=True,
                                           add_generation_prompt=True,
                                           padding=True,
                                           return_tensors='pt',
                                           return_dict=True).to(device)

    # 2.2
    # 2.2.1 流式输出，streamer 会自动将数据推送到标准输出流，用户不需要 print 打印输出
    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    outputs = estimator.generate(**inputs,
                                 do_sample=True,
                                 top_k=5,
                                 top_p=0.5,
                                 max_new_tokens=256,
                                 streamer=streamer)

    # 2.2.2 一次性输出

    # 计算 prompt 长度
    # propmt_length = len(inputs['input_ids'][0])
    # # 模型进行内容生成
    # outputs = estimator.generate(**inputs, do_sample=True, max_new_tokens=256)
    # # 去掉输出前面的 prompt
    # outputs = outputs[0][propmt_length:]
    # outputs = tokenizer.decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    # print(outputs)



if __name__ == '__main__':
    demo()