使用llm-compressor 对 Qwen3-14B 做 AWQ + INT4 量化

一、环境准备

1. 创建与激活 Conda 环境

bash 复制代码

conda create -n qwen3-awq python=3.10 -y
conda activate qwen3-awq

2. 安装 PyTorch（带 CUDA）

注意：下面的 CUDA 版本 cu129 要和机器上的 CUDA/驱动匹配。

bash 复制代码

pip install --upgrade pip

pip install torch==2.9.0 torchvision torchaudio \
  --extra-index-url https://download.pytorch.org/whl/cu129

3. 安装 llmcompressor、vLLM 和相关依赖

bash 复制代码

pip install "vllm>=0.6.0" "llmcompressor>=0.8.0"
pip install "transformers>=4.44.0" "accelerate" "datasets" "safetensors" "huggingface_hub"
pip install "bitsandbytes"

二、确认 GPU 是否可用

在 Python 里执行：

python 复制代码

import torch
print("cuda:", torch.cuda.is_available(), "device:", torch.cuda.get_device_name(0))

如果输出类似：

bash 复制代码

cuda: True device: NVIDIA A100-SXM-80GB

说明 GPU 可以正常使用。

三、编写量化脚本（AWQ + INT4）

下面给出一个完整可运行脚本，保存为 awq_int4_qwen3.py 之类的名字即可。

该脚本主要实现以下功能：

加载 Qwen3-14B 原始模型和分词器
加载校准数据集 ultrachat_200k
配置 AWQ 的 INT4 量化（W4A16）
调用 oneshot 完成量化
简单生成一段文本测试
将量化后模型保存到本地

python 复制代码

import os

# 可选：使用 HuggingFace 镜像（例如中国大陆环境）
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier

# ======================
# 1. 模型路径配置
# ======================
# 如果你是本地已经下好的模型，保持文档中的写法，例如：
MODEL_ID = "/mnt/e/models/Qwen3-14B-2"
# 如果是从 HuggingFace 在线拉取，可以写成：
# MODEL_ID = "Qwen/Qwen2-7B-Instruct"  # 仅示例，换成真正的 Qwen3-14B 路径

# ======================
# 2. 加载模型和分词器
# ======================
print(f"Loading model from {MODEL_ID} ...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype="auto",        # 自动选择精度，一般是 bfloat16 / float16
    device_map="auto",   # 自动将模型切到 GPU / 多卡
    offload_folder="offload"  # 注意：这里参数名是 offload_folder
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# ======================
# 3. 准备校准数据集
# ======================
# 使用 ultrachat_200k 的 train_sft split
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"

# 采样 256 条做校准（可以适当增加，精度会更好但更耗时）
NUM_CALIBRATION_SAMPLES = 256
MAX_SEQUENCE_LENGTH = 512

print("Loading calibration dataset...")
ds = load_dataset(
    DATASET_ID,
    split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]"
)
ds = ds.shuffle(seed=42)

# 将聊天数据转成纯文本，使用 chat_template
def preprocess(example):
    return {
        "text": tokenizer.apply_chat_template(
            example["messages"],
            tokenize=False,
        )
    }

ds = ds.map(preprocess)

# 如需显式 tokenization，可以再执行一次 map：
def tokenize(sample):
    return tokenizer(
        sample["text"],
        padding=False,
        max_length=MAX_SEQUENCE_LENGTH,
        truncation=True,
        add_special_tokens=False,
    )

# 可选：如果你希望在 oneshot 前就把 input_ids 准备好，可以取消下面注释
# ds = ds.map(tokenize, remove_columns=ds.column_names)

# ======================
# 4. 配置 AWQ INT4 量化
# ======================
# 说明：
#   - ignore：跳过 lm_head 以及一些 MoE 门控层不做量化
#   - scheme：这里是重点！INT4 用 W4A16
#   - targets：只量化 Linear 层
#   - group_size：分组量化的组大小，128 是个常用设置
recipe = [
    AWQModifier(
        ignore=["lm_head", "re:.*mlp.gate$", "re:.*mlp.shared_expert_gate$"],
        scheme="W4A16",          # ✅ INT4 量化：W4A16（文档注释：INT4量化用W4A16）
        targets=["Linear"],
        group_size=128,
    ),
]

# ======================
# 5. 运行 oneshot 量化
# ======================
if __name__ == "__main__":
    print("Start AWQ INT4 quantization ...")
    oneshot(
        model=model,
        dataset=ds,
        recipe=recipe,
        max_seq_length=MAX_SEQUENCE_LENGTH,
        num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    )

    # ======================
    # 6. 简单生成测试
    # ======================
    print("\n\n========== SAMPLE GENERATION ==============")
    input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to(
        model.device
    )
    output = model.generate(input_ids, max_new_tokens=100)
    print(tokenizer.decode(output[0]))
    print("==========================================\n\n")

    # ======================
    # 7. 保存量化后模型
    # ======================
    # 以原模型名 + 后缀的方式命名输出目录
    SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-awq-int4-sym"
    print(f"Saving quantized model to {SAVE_DIR} ...")

    model.save_pretrained(SAVE_DIR, save_compressed=True)
    tokenizer.save_pretrained(SAVE_DIR)

    print("Done.")

四、运行脚本

在终端中，确保已经激活 qwen3-awq 环境，然后执行：

bash 复制代码

python awq_int4_qwen3.py

过程大概会经历：

加载 Qwen3-14B 原始模型（会占较多显存/内存）
下载并加载 ultrachat_200k 子集（如本地无缓存）
运行 AWQ INT4 量化
打印一段示例生成结果
在当前目录下生成一个新文件夹，如：Qwen3-14B-2-awq-int4-sym，里面包含 config.json, model.safetensors 等压缩权重