Upload batch_quantize_w4a16.py with huggingface_hub

batch_quantize_w4a16.py

@@ -0,0 +1,175 @@
+import os
+import torch
+from llmcompressor import oneshot
+from llmcompressor.modifiers.awq import AWQModifier, AWQMapping
+from datasets import Dataset
+
+# ============================================================
+# 配置区（批量量化，仅需修改此处全局配置，无需逐个调整模型）
+# ============================================================
+ROOT_MODEL_DIR = "./"
+QUANT_SUFFIX = "_awq_w4a16"
+
+# 校准数据
+CALIB_DATA = [
+    """You are a helpful assistant.
+User: 帮我写一份关于全球气候变化的报告大纲。
+Assistant: 当然，这是一个关于全球气候变化报告的大纲建议：
+I. 引言
+   A. 什么是全球气候变化
+   B. 报告的目的和范围
+II. 气候变化的原因
+   A. 自然原因（太阳活动、火山喷发）
+   B. 人为原因（温室气体排放、土地利用变化）
+"""
+]
+
+# AWQ 映射表（适配Qwen系列模型）
+LLAMA_MAPPINGS = [
+    AWQMapping(
+        "re:.*input_layernorm",
+        ["re:.*q_proj", "re:.*k_proj", "re:.*v_proj"],
+    ),
+    AWQMapping("re:.*v_proj", ["re:.*o_proj"]),
+    AWQMapping(
+        "re:.*post_attention_layernorm",
+        ["re:.*gate_proj", "re:.*up_proj"],
+    ),
+    AWQMapping(
+        "re:.*up_proj",
+        ["re:.*down_proj"],
+    ),
+]
+
+# ============================================================
+# 工具函数：获取所有待量化的sft_qwenxxx模型目录
+# ============================================================
+def get_target_model_dirs():
+    """
+    遍历ROOT_MODEL_DIR，筛选出所有sft_qwen开头的目录（待量化模型）
+    排除已经量化过的模型（包含_awq_的目录）
+    """
+    target_dirs = []
+    skipped_dirs = []
+    
+    for item in os.listdir(ROOT_MODEL_DIR):
+        item_path = os.path.abspath(os.path.join(ROOT_MODEL_DIR, item))
+        
+        # 筛选条件：是目录 + 以sft_qwen开头
+        if os.path.isdir(item_path) and item.startswith("sft_qwen"):
+            # 【修复】排除已经量化过的模型（包含_awq_的目录）
+            if "_awq_" in item:
+                skipped_dirs.append(item)
+                print(f"[跳过已量化模型] {item}")
+            else:
+                target_dirs.append(item)
+                print(f"[发现待量化模型] {item}")
+    
+    if skipped_dirs:
+        print(f"\n⏭️  跳过 {len(skipped_dirs)} 个已量化模型")
+    
+    if not target_dirs:
+        print("⚠️  未发现任何待量化的sft_qwen模型目录")
+    else:
+        print(f"\n✅ 共发现 {len(target_dirs)} 个待量化模型\n")
+    
+    return target_dirs
+
+# ============================================================
+# 核心函数：单个模型量化
+# ============================================================
+def quantize_single_model(model_name):
+    """
+    量化单个模型
+    :param model_name: 模型目录名（如sft_qwen3_4b）
+    """
+    MODEL_PATH = os.path.join(ROOT_MODEL_DIR, model_name)
+    QUANT_PATH = os.path.join(ROOT_MODEL_DIR, f"{model_name}{QUANT_SUFFIX}")
+
+    print(f"\n" + "="*100)
+    print(f"开始量化模型: {model_name}")
+    print(f"模型输入路径: {MODEL_PATH}")
+    print(f"量化输出路径: {QUANT_PATH}")
+    print("="*100 + "\n")
+
+    if not torch.cuda.is_available():
+        print("❌ 错误：此过程需要GPU支持，无法继续量化")
+        return False
+    
+    try:
+        calib_dataset = Dataset.from_dict({"text": CALIB_DATA})
+    except Exception as e:
+        print(f"❌ 构建校准数据集失败，错误：{e}")
+        return False
+
+    # 每次量化创建全新的AWQModifier实例
+    recipe = [
+        AWQModifier(
+            scheme="W4A16_ASYM",
+            mappings=LLAMA_MAPPINGS,
+            ignore=["lm_head"],
+            targets=["Linear"]
+        ),
+    ]
+
+    try:
+        print("[步骤 1/2] 正在执行AWQ W4A16 oneshot量化...")
+        print("   此过程会进行权重缩放和低比特量化，耗时较长，请耐心等待...")
+        
+        # 【修复】移除不支持的 tokenizer_kwargs 参数
+        oneshot(
+            model=MODEL_PATH,
+            dataset=calib_dataset,
+            recipe=recipe,
+            output_dir=QUANT_PATH,
+            num_calibration_samples=len(CALIB_DATA),
+            max_seq_length=4096,
+        )
+        
+        print("\n[步骤 2/2] oneshot量化完成！")
+
+    except Exception as e:
+        print(f"\n❌ 量化模型 {model_name} 过程中发生错误")
+        print(f"错误详情: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+    finally:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+
+    print("\n" + "="*80)
+    print(f"🎉 模型 {model_name} 量化成功！")
+    print(f"4-bit AWQ模型已保存到: {QUANT_PATH}")
+    print("="*80 + "\n")
+    return True
+
+# ============================================================
+# 主函数：批量执行所��模型量化
+# ============================================================
+def run_batch_quantization():
+    print("🚀 启动Qwen系列模型批量W4A16量化任务")
+    print(f"工作目录: {os.path.abspath(ROOT_MODEL_DIR)}\n")
+
+    target_models = get_target_model_dirs()
+    if not target_models:
+        return
+
+    success_count = 0
+    fail_count = 0
+    for idx, model_name in enumerate(target_models, 1):
+        print(f"\n========== 批量量化 {idx}/{len(target_models)} ==========")
+        if quantize_single_model(model_name):
+            success_count += 1
+        else:
+            fail_count += 1
+
+    print("\n" + "="*100)
+    print("📊 批量量化任务全部结束")
+    print(f"✅ 成功量化：{success_count} 个模型")
+    print(f"❌ 量化失败：{fail_count} 个模型")
+    print("="*100)
+
+if __name__ == "__main__":
+    run_batch_quantization()