Qwen3-Coder-30B-A3B-Instruct AWQ 量化

import os
import gc
import torch
import loggingfrom datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier
from transformers import AutoModelForCausalLM, AutoTokenizer# 清理缓存
gc.collect()
torch.cuda.empty_cache()# 配置日志以显示进度
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)model_path = "./models/Qwen3-Coder-30B-A3B-Instruct"
quant_path = "./Qwen3-Coder-30B-A3B-Instruct-AWQ"# 创建 offload 目录
# os.makedirs("offload_folder", exist_ok=True)
# 配置内存限制
max_memory = {0: "23GB", 1: "23GB", "cpu": "100GB"}
logger.info(f"设定最大内存分配: {max_memory}")logger.info("使用 accelerate 的 'auto' device_map 并结合 max_memory 进行加载...")
torch_dtype = torch.float16
logger.info(f"设定模型加载精度 (torch_dtype): {torch_dtype}")
logger.info("正在加载模型和tokenizer...")
try:model = AutoModelForCausalLM.from_pretrained(model_path,torch_dtype=torch_dtype,trust_remote_code=True,device_map="auto",max_memory=max_memory,# offload_folder="offload_folder",attn_implementation="eager",low_cpu_mem_usage=True)logger.info(f"模型已加载，分配的设备映射 (model.hf_device_map): {model.hf_device_map}")
except Exception as e:logger.error(f"加载模型时出错: {e}")raisetokenizer = AutoTokenizer.from_pretrained(model_path,trust_remote_code=True
)# 配置校准数据集
logger.info("正在加载校准数据集...")
NUM_CALIBRATION_SAMPLES = 256
MAX_SEQUENCE_LENGTH = 512# 加载本地 parquet 数据集
logger.info(f"使用前 {NUM_CALIBRATION_SAMPLES} 个样本进行校准...")
ds = load_dataset("parquet",data_files="./train-00000-of-00206.parquet",split=f"train[:{NUM_CALIBRATION_SAMPLES}]"
)
ds = ds.shuffle(seed=42)# 确认数据集结构
logger.info(f"数据集列名: {ds.column_names}")
logger.info(f"样本数量: {len(ds)}")# # 预处理函数: 适配 Qwen3 的 chat template
# def preprocess(example):
#     return {
#         "text": tokenizer.apply_chat_template(
#             [{"role": "user", "content": example["content"]}],
#             tokenize=False,
#             add_generation_prompt=True
#         )
#     }# 修正的预处理函数
def preprocess(examples):conversations = [[{"role": "user", "content": text}]for text in examples["content"]]texts = tokenizer.apply_chat_template(conversations,tokenize=False,)return {"text": texts}logger.info("正在预处理校准数据集（多进程加速）...")
ds = ds.map(preprocess, batched=True, batch_size=128)# 不需要额外的 tokenize 函数，oneshot 内部会处理
# Tokenize 函数
# def tokenize(sample):
#     return tokenizer(
#         sample["text"],
#         padding=False,    # AWQ 通常处理变长序列
#         max_length=MAX_SEQUENCE_LENGTH,
#         truncation=True,
#         add_special_tokens=False,
#     )# logger.info("正在 tokenize 校准数据集...")
# 批处理加速
# ds = ds.map(tokenize, remove_columns=ds.column_names, batched=True, batch_size=128)# 配置 AWQ 量化
logger.info("配置 AWQ 量化参数...")
recipe = [AWQModifier(# MoE 模型需要忽略这些层ignore=["lm_head","re:.*mlp.gate$","re:.*mlp.shared_experts.gate$"],scheme="W4A16",            # 4-bit 权重，16-bit 激活targets=["Linear"],),
]# 执行量化
logger.info("开始 AWQ 量化过程...")
# logging.getLogger("llmcompressor").setLevel(logging.DEBUG)
oneshot(model=model,dataset=ds,recipe=recipe,max_seq_length=MAX_SEQUENCE_LENGTH,num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)# 验证量化模型
logger.info("\n\n")
logger.info("========== 样本生成测试 ==============")
input_text = "写一个Python函数，实现快速排序"
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=100)
logger.info(tokenizer.decode(output[0], skip_special_tokens=True))
logger.info("==========================================\n\n")# 保存量化模型（4GB 分片）
logger.info("正在保存量化模型（4GB分片）...")
model.save_pretrained(quant_path,save_compressed=True,max_shard_size="4GB",safe_serialization=True
)
tokenizer.save_pretrained(quant_path)
logger.info(f"AWQ 量化完成！模型已保存至 {quant_path}")# 清理缓存
gc.collect()
torch.cuda.empty_cache()#######################################
# 验证分片情况
import globshards = glob.glob(os.path.join(quant_path, "model-*.safetensors"))
logger.info(f"\n模型已成功分片保存，共 {len(shards)} 个分片:")
for shard in shards:size_gb = os.path.getsize(shard) / (1024**3)logger.info(f"- {os.path.basename(shard)}: {size_gb:.2f} GB")logger.info(f"\n量化完成！模型已保存至 {quant_path}")
logger.info("注意：模型现在使用标准命名格式 model-00001-of-0000X.safetensors")