一、简介

文本采用待量化模型自己生成数据集的方式，使用llmcompressor对智谱公司的GLM-4-9b-chat-hf模型进行AWQ量化。

二、环境安装

conda activate meeting

pip install vllm llmcompressor --index-url https://pypi.org/simple/
pip install transformers accelerate datasets safetensors --index-url https://pypi.org/simple/

三、准备数据集

import os
import json
import torch
from datasets import Dataset, DatasetDict
from transformers import AutoModelForCausalLM, AutoTokenizer
from tqdm import tqdm
import argparse

# ======================
# 配置参数
# ======================
def parse_args():
    parser = argparse.ArgumentParser(description="生成GLM-4-9b-chat-hf问答数据集")
    parser.add_argument("--model_path", type=str, default="/data2/model/zhipu/glm-4-9b-chat-hf",
                       help="原始模型路径")
    parser.add_argument("--output_dir", type=str, default="./glm-4-9b-chat-dataset",
                       help="数据集保存路径")
    parser.add_argument("--num_samples", type=int, default=128,
                       help="生成的样本数量")
    parser.add_argument("--max_new_tokens", type=int, default=200,
                       help="生成回答的最大token数")
    parser.add_argument("--temperature", type=float, default=0.7,
                       help="生成温度")
    parser.add_argument("--top_p", type=float, default=0.9,
                       help="top-p采样参数")
    parser.add_argument("--batch_size", type=int, default=8,
                       help="批量生成的大小")
    return parser.parse_args()

# ======================
# 问题列表
# ======================
QUESTION_SETS = {
    "general_knowledge": [
        "你是谁？",
        "中华人民共和国成立是哪一天？",
        "什么是人工智能？",
        "解释一下机器学习的基本概念",
        "Python是什么编程语言？",
        "太阳系有哪些行星？",
        "水的化学式是什么？",
        "中国的首都是哪里？",
        "什么是深度学习？",
        "解释一下牛顿第一定律",
        "如何保持健康的生活方式？",
        "什么是气候变化？",
        "解释一下光合作用的过程",
        "什么是区块链技术？",
        "中国的四大发明是什么？",
        "什么是量子计算？",
        "如何学习编程？",
        "什么是可再生能源？",
        "解释一下相对论的基本概念",
        "什么是大数据？"
    ],
    
    "science_tech": [
        "什么是神经网络？",
        "解释一下Transformer架构",
        "什么是自然语言处理？",
        "量子力学的基本原理是什么？",
        "什么是基因编辑技术？",
        "解释一下5G技术",
        "什么是物联网？",
        "人工智能有哪些应用场景？",
        "什么是虚拟现实？",
        "解释一下云计算",
        "什么是边缘计算？",
        "自动驾驶汽车是如何工作的？",
        "什么是机器学习中的过拟合？",
        "解释一下梯度下降算法",
        "什么是GPT模型？",
        "什么是强化学习？",
        "什么是计算机视觉？",
        "解释一下注意力机制",
        "什么是大语言模型？",
        "如何评估机器学习模型的性能？"
    ],
    
    "history_culture": [
        "中国历史上第一个封建王朝是哪个？",
        "文艺复兴时期有哪些重要人物？",
        "解释一下工业革命的影响",
        "中国古代的丝绸之路是什么？",
        "什么是唐宋八大家？",
        "解释一下儒家思想的核心",
        "第二次世界大战的主要事件有哪些？",
        "中国的春节有哪些传统习俗？",
        "什么是古希腊哲学？",
        "解释一下启蒙运动",
        "中国的长城是什么时候修建的？",
        "什么是佛教的基本教义？",
        "解释一下法国大革命",
        "中国古代的科举制度是什么？",
        "什么是印象派艺术？",
        "解释一下美国独立战争",
        "中国的茶文化有哪些特点？",
        "什么是罗马帝国的兴衰？",
        "解释一下中世纪欧洲的社会结构",
        "什么是中国的京剧？"
    ],
    
    "daily_life": [
        "如何做西红柿炒鸡蛋？",
        "早上如何快速高效地开始一天？",
        "如何提高学习效率？",
        "什么是健康饮食的基本原则？",
        "如何缓解工作压力？",
        "如何进行时间管理？",
        "如何开始健身锻炼？",
        "什么是有效的沟通技巧？",
        "如何养成良好的阅读习惯？",
        "如何规划个人财务？",
        "如何改善睡眠质量？",
        "什么是积极的心态？",
        "如何处理人际关系中的冲突？",
        "如何进行职业规划？",
        "什么是环保的生活方式？",
        "如何学习一门外语？",
        "如何培养创造力？",
        "什么是有效的学习方法？",
        "如何进行情绪管理？",
        "如何准备一次旅行？"
    ],
    
    "philosophy_ethics": [
        "什么是道德？",
        "解释一下功利主义",
        "什么是自由意志？",
        "什么是存在主义？",
        "解释一下康德的道德哲学",
        "什么是社会契约论？",
        "什么是幸福？如何获得幸福？",
        "解释一下柏拉图的理论",
        "什么是公平正义？",
        "什么是人工智能伦理？",
        "什么是环境保护的伦理基础？",
        "解释一下亚里士多德的伦理学",
        "什么是个人与集体的关系？",
        "什么是真理？",
        "解释一下道家思想",
        "什么是人生的意义？",
        "什么是权力与责任？",
        "解释一下尼采的哲学",
        "什么是美德伦理学？",
        "什么是科学与伦理的关系？"
    ]
}

# ======================
# 生成回答
# ======================
def generate_answer(model, tokenizer, question, generation_config):
    """生成单个问题的回答"""
    messages = [{"role": "user", "content": question}]
    
    # 使用chat_template格式化输入
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            **generation_config
        )
    
    # 解码输出，跳过输入部分
    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 提取生成的回答部分
    answer = full_output[len(tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)):].strip()
    
    return answer

def batch_generate_answers(model, tokenizer, questions, generation_config, batch_size=8):
    """批量生成回答"""
    answers = []
    
    for i in tqdm(range(0, len(questions), batch_size), desc="生成回答"):
        batch_questions = questions[i:i+batch_size]
        batch_answers = []
        
        for question in batch_questions:
            answer = generate_answer(model, tokenizer, question, generation_config)
            batch_answers.append(answer)
        
        answers.extend(batch_answers)
    
    return answers

# ======================
# 创建数据集
# ======================
def create_dataset(args):
    """创建问答数据集"""
    
    print("="*60)
    print(f"加载模型: {args.model_path}")
    print("="*60)
    
    # 加载模型和tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )
    
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True
    )
    
    # 确保pad_token设置正确
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # 设置生成参数
    generation_config = {
        "max_new_tokens": args.max_new_tokens,
        "temperature": args.temperature,
        "do_sample": True,
        "top_p": args.top_p,
        "pad_token_id": tokenizer.pad_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }
    
    # 准备问题列表
    all_questions = []
    for category, questions in QUESTION_SETS.items():
        all_questions.extend(questions)
    
    # 如果问题数量不够，重复使用
    if len(all_questions) < args.num_samples:
        print(f"问题库只有{len(all_questions)}个，需要{args.num_samples}个，将重复使用问题...")
        repeated_questions = []
        while len(repeated_questions) < args.num_samples:
            repeated_questions.extend(all_questions)
        all_questions = repeated_questions[:args.num_samples]
    else:
        # 随机选择指定数量的样本
        import random
        random.seed(42)
        all_questions = random.sample(all_questions, args.num_samples)
    
    print(f"总共使用 {len(all_questions)} 个问题")
    
    # 批量生成回答
    answers = batch_generate_answers(
        model, 
        tokenizer, 
        all_questions, 
        generation_config,
        batch_size=args.batch_size
    )
    
    # 构建数据集样本
    samples = []
    for i, (question, answer) in enumerate(zip(all_questions, answers)):
        sample = {
            "messages": [
                {"role": "user", "content": question},
                {"role": "assistant", "content": answer}
            ]
        }
        samples.append(sample)
        
        # 打印前几个样本
        if i < 3:
            print(f"\n样本 {i+1}:")
            print(f"问题: {question}")
            print(f"回答: {answer[:100]}...")
    
    # 创建数据集
    dataset = Dataset.from_list(samples)
    
    # 创建DatasetDict（可选）
    dataset_dict = DatasetDict({
        "train": dataset
    })
    
    # 保存数据集
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 保存为多种格式
    dataset.save_to_disk(args.output_dir)
    
    # 保存为JSON文件（便于查看）
    json_path = os.path.join(args.output_dir, "dataset.json")
    with open(json_path, 'w', encoding='utf-8') as f:
        json_data = []
        for sample in samples:
            json_data.append(sample)
        json.dump(json_data, f, ensure_ascii=False, indent=2)
    
    # 保存为parquet格式
    parquet_path = os.path.join(args.output_dir, "dataset.parquet")
    dataset.to_parquet(parquet_path)
    
    print("\n" + "="*60)
    print("数据集生成完成!")
    print(f"样本数量: {len(dataset)}")
    print(f"保存路径: {args.output_dir}")
    print("="*60)
    
    # 验证数据集可以正常加载
    print("\n验证数据集加载...")
    try:
        from datasets import load_from_disk
        loaded_dataset = load_from_disk(args.output_dir)
        print(f"成功加载数据集，包含 {len(loaded_dataset)} 个样本")
        print(f"第一个样本: {loaded_dataset[0]}")
    except Exception as e:
        print(f"加载数据集时出错: {e}")
    
    # 清理内存
    del model
    torch.cuda.empty_cache()
    
    return dataset

# ======================
# 使用示例
# ======================
def main():
    args = parse_args()
    
    print("GLM-4-9b-chat-hf问答数据集生成器")
    print(f"模型路径: {args.model_path}")
    print(f"输出目录: {args.output_dir}")
    print(f"样本数量: {args.num_samples}")
    print(f"批量大小: {args.batch_size}")
    
    dataset = create_dataset(args)
    
    # 生成使用说明
    readme_content = f"""
# GLM-4-9b-chat-hf 问答数据集

## 基本信息
- 模型: {args.model_path}
- 样本数量: {args.num_samples}
- 生成时间: {os.path.basename(args.output_dir)}
- 格式: HuggingFace Dataset

## 数据集结构
每个样本的格式为:
```json
{{
  "messages": [
    {{"role": "user", "content": "问题文本"}},
    {{"role": "assistant", "content": "回答文本"}}
  ]
}}
"""


if __name__ == "__main__":
    main()

执行：

python generate_dataset.py \
  --model_path /data2/model/zhipu/glm-4-9b-chat-hf \
  --output_dir ./glm-4-9b-chat-dataset \
  --num_samples 128

执行完后，就会在当前目录下面多出一个glm-4-9b-chat-dataset目录，里面就是生成的数据集。

五、AWQ量化

import os
import torch
import numpy as np
from datasets import load_from_disk

from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier, AWQMapping

# ======================
# 1. 内存优化设置
# ======================
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"


# ======================
# 2. 模型配置
# ======================
MODEL_PATH = "/data2/model/zhipu/glm-4-9b-chat-hf"
SAVE_DIR = "./glm-4-9b-chat-awq-int4-self"

# ======================
# 3. 加载模型
# ======================
print(f"Loading model from {MODEL_PATH} ...")
torch.cuda.empty_cache()

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True,
    max_memory={0: "20GB"}
)

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

# ======================
# 4. 创建校准数据集
# ======================
NUM_CALIBRATION_SAMPLES = 128
MAX_SEQUENCE_LENGTH = 512


print("Loading calibration dataset...")
ds = load_from_disk("./glm-4-9b-chat-dataset")
ds = ds.shuffle(seed=42)


print("********************************")
print("原始数据样本结构:")
print(ds[0])

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

ds = ds.map(preprocess)
print("================================")
print("预处理后的数据样本:")
print(ds[0])
print("********************************")

# 测试 chat_template 转换
test_example = {
    "messages": [
        {"role": "user", "content": "What's AI?"},
        {"role": "assistant", "content": "AI means Artificial Intelligence."},
        {"role": "user", "content": "Give examples."},
        {"role": "assistant", "content": "Examples: self-driving cars, chatbots."}
    ]
}

result = preprocess(test_example)
print("测试 chat_template 转换:")
print(result)
#输出: {'text': "[gMASK]<sop><|user|>\nWhat's AI?<|assistant|>\nAI means Artificial Intelligence.<|user|>\nGive examples.<|assistant|>\nExamples: self-driving cars, chatbots."}
print("********************************")

# ======================
# 5. 配置AWQ量化
# ======================
print("Configuring AWQ quantization...")

recipe = [
    AWQModifier(
        scheme="W4A16",
        targets=["Linear"],
        ignore=["lm_head", "re:.*mlp\.gate_up_proj$", "re:.*mlp\.down_proj$"],
        mappings=[
            AWQMapping(
                smooth_layer='re:.*layers\.\d+\.input_layernorm$',
                balance_layers=[
                    're:.*layers\.\d+\.self_attn\.q_proj$',
                    're:.*layers\.\d+\.self_attn\.k_proj$',
                    're:.*layers\.\d+\.self_attn\.v_proj$'
                ]
            ),
        ]
    ),
]

# ======================
# 6. 运行oneshot量化
# ======================
if __name__ == "__main__":
    print("\n" + "="*60)
    print("Starting AWQ Quantization")
    print("="*60)
    print(f"Model: {MODEL_PATH}")
    print(f"Output: {SAVE_DIR}")
    print(f"Calibration samples: {len(ds)}")
    print(f"Sequence length: {MAX_SEQUENCE_LENGTH}")
    print("="*60)
    
    torch.cuda.empty_cache()
    
    try:
        oneshot(
            model=model,
            dataset=ds,
            recipe=recipe,
            max_seq_length=MAX_SEQUENCE_LENGTH,
            num_calibration_samples=len(ds),
            output_dir=SAVE_DIR,
        )
        
        print("\n✅ AWQ量化成功!")
        print(f"模型已保存到: {SAVE_DIR}")
        
    except Exception as e:
        print(f"\n❌ AWQ量化失败: {e}")
        print("\n建议：")
        print("1. 减少NUM_CALIBRATION_SAMPLES到64")
        print("2. 减少MAX_SEQUENCE_LENGTH到256")
        print("3. 确保GPU有足够内存")
        exit(1)
    
    # ======================
    # 7. 测试量化模型（修复版本）
    # ======================
    print("\n" + "="*60)
    print("Testing Quantized Model")
    print("="*60)
    
    # 清理内存
    del model
    torch.cuda.empty_cache()
    
    # 重新加载量化模型
    print(f"Loading quantized model from {SAVE_DIR} ...")
    try:
        quantized_model = AutoModelForCausalLM.from_pretrained(
            SAVE_DIR,
            device_map="auto",
            trust_remote_code=True,
            torch_dtype=torch.float16,
        )
        
        # 设置生成参数，避免attention mask警告
        generation_config = {
            "max_new_tokens": 100,
            "temperature": 0.7,
            "do_sample": True,
            "top_p": 0.9,
            "pad_token_id": tokenizer.pad_token_id,
            "eos_token_id": tokenizer.eos_token_id,
        }
        
        # 测试多个问题
        test_questions = [
            "中华人民共和国成立日期是哪一天?",
            "什么是机器学习?"
        ]
        
        for i, question in enumerate(test_questions):
            print(f"\n[问题 {i+1}]: {question}")
            
            # 编码输入，包括attention_mask
            inputs = tokenizer(
                question, 
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512
            ).to(quantized_model.device)
            
            with torch.no_grad():
                output = quantized_model.generate(
                    **inputs,
                    **generation_config
                )
            
            # 解码输出，跳过输入部分
            full_output = tokenizer.decode(output[0], skip_special_tokens=True)
            # 提取生成的回答部分
            answer = full_output[len(question):].strip()
            print(f"[回答]: {answer}")
        
        print("\n" + "="*60)
        print("✅ 量化完成！模型已成功保存并通过测试")
        print("="*60)
        
        # ======================
        # 8. 额外：检查模型大小和性能
        # ======================
        print("\n" + "="*60)
        print("模型信息")
        print("="*60)
        
        import os
        import glob
        
        # 检查量化模型大小
        model_files = glob.glob(os.path.join(SAVE_DIR, "*.safetensors")) + \
                     glob.glob(os.path.join(SAVE_DIR, "*.bin"))
        
        total_size = sum(os.path.getsize(f) for f in model_files) / (1024**3)
        print(f"量化模型大小: {total_size:.2f} GB")
        
        # 原始模型大小（估计）
        print(f"原始模型大小（估计）: ~18.0 GB (FP16)")
        print(f"压缩比例: {(1 - total_size/18.0)*100:.1f}%")
        
        # 检查量化后的参数
        print("\n量化后模型参数示例:")
        for name, param in quantized_model.named_parameters():
            if "q_proj" in name and "weight" in name:
                print(f"{name}: {param.shape}, dtype: {param.dtype}")
                break
        
    except Exception as e:
        print(f"\n❌ 测试量化模型失败: {e}")
        import traceback
        traceback.print_exc()
        print("量化模型可能已保存，但加载或测试时出错")

执行：

CUDA_VISIBLE_DEVICES=0,1 python awq_self_datasets_quantization.py

代码里面要注意的地方是：数据集要用apply_chat_template生成模型适配的格式，比如：{‘text’: “[gMASK]<|user|>\nWhat’s AI?<|assistant|>\nAI means Artificial Intelligence.<|user|>\nGive examples.<|assistant|>\nExamples: self-driving cars, chatbots.”}，不同模型生成的格式可能会不一样。还有一点，oneshot函数，有一个参数output_dir直接就可以指定量化后模型的保存路径了。