From 08514cd9ec4628a9d7bd457ad1a7c920ac33e512 Mon Sep 17 00:00:00 2001
From: root <403644786@qq.com>
Date: Mon, 15 Jul 2024 14:49:32 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=BA=86bnb=E7=9A=84?=
 =?UTF-8?q?=E9=87=8F=E5=8C=96demo?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 quantize/bnb_quantize.py | 57 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 quantize/bnb_quantize.py

diff --git a/quantize/bnb_quantize.py b/quantize/bnb_quantize.py
new file mode 100644
index 0000000..c32cea6
--- /dev/null
+++ b/quantize/bnb_quantize.py
@@ -0,0 +1,57 @@
+"""
+the script will use bitandbytes to quantize the MiniCPM language model.
+the be quantized model can be finetuned by MiniCPM or not.
+you only need to set the model_path 、save_path and run bash code 
+
+cd MiniCPM
+python quantize/bnb_quantize.py
+
+you will get the quantized model in save_path、quantized_model test time and gpu usage
+"""
+
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import time
+import torch
+import GPUtil
+import os
+
+model_path = "/root/ld/ld_model_pretrain/MiniCPM-1B-sft-bf16"  # 模型下载地址
+save_path = "/root/ld/ld_model_pretrain/MiniCPM-1B-sft-bf16_int4"  # 量化模型保存地址
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 创建一个配置对象来指定量化参数
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,  # 是否进行4bit量化
+    load_in_8bit=False,  # 是否进行8bit量化
+    bnb_4bit_compute_dtype=torch.float16,  # 计算精度设置
+    bnb_4bit_quant_storage=torch.uint8,  # 量化权重的储存格式
+    bnb_4bit_quant_type="nf4",  # 量化格式，这里用的是正太分布的int4
+    bnb_4bit_use_double_quant=True,  # 是否采用双量化，即对zeropoint和scaling参数进行量化
+    llm_int8_enable_fp32_cpu_offload=False,  # 是否llm使用int8，cpu上保存的参数使用fp32
+    llm_int8_has_fp16_weight=False,  # 是否启用混合精度
+    #llm_int8_skip_modules=["out_proj", "kv_proj", "lm_head"],  # 不进行量化的模块
+    llm_int8_threshold=6.0,  # llm.int8()算法中的离群值，根据这个值区分是否进行量化
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    device_map=device,  # 分配模型到device
+    quantization_config=quantization_config,
+    trust_remote_code=True,
+)
+
+gpu_usage = GPUtil.getGPUs()[0].memoryUsed
+start = time.time()
+response =  model.chat(tokenizer, "<用户>给我讲一个故事<AI>",history=[], temperature=0.5, top_p=0.8, repetition_penalty=1.02)  # 模型推理
+print("量化后输出", response)
+print("量化后推理用时", time.time() - start)
+print(f"量化后显存占用: {round(gpu_usage/1024,2)}GB")
+
+
+# 保存模型和分词器
+os.makedirs(save_path, exist_ok=True)
+model.save_pretrained(save_path, safe_serialization=True)
+tokenizer.save_pretrained(save_path)