模型量化与保存
模型量化是将神经网络中的浮点计算(如32位/16位)转换为低比特定点计算(如8位整型)的技术,通过减少数值精度来压缩模型体积并加速计算。其本质是一种针对模型参数的有损信息压缩,在保证模型性能的前提下优化部署效率。
一、模型量化的目标
- 压缩模型体积
将浮点权重转换为低位宽整数存储,模型体积可减少至原大小的1/4。 - 加速推理计算
硬件对低位宽计算(如INT8)的支持效率通常比FP32快2-4倍
二、transfomers 进行模型量化与保存
- 前提环境
pip install --upgrade transformers accelerate bitsandbytes
1、4bit
- 4 位量化模型可将内存使用量减少 4 倍,对于大型模型,设置“device_map=“auto””可有效地在所有可用 GPU 上分配权重。
from transformers import AutoModelForCausalLM, BitsAndBytesConfig#配置量化参数
quantization_config = BitsAndBytesConfig(load_in_4bit=True)model_4bit = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b7",device_map="auto",quantization_config=quantization_config
)
2、8bit
- 8 位量化模型可将内存使用量减半,对于大型模型,设置“device_map=“auto””以有效地将权重分配到所有可用的 GPU 上。
from transformers import AutoModelForCausalLM, BitsAndBytesConfigquantization_config = BitsAndBytesConfig(load_in_8bit=True)model_8bit = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b7", device_map="auto",quantization_config=quantization_config
)
3、量化模型保存与加载
- 保存
# 保存量化模型
model.save_pretrained("./quant_llama2")
tokenizer.save_pretrained("./quant_llama2")
- 加载
from transformers import AutoModelForCausalLM, AutoTokenizermodel = AutoModelForCausalLM.from_pretrained("本地模型路径", device_map="auto")
4、模型反(去除)量化
- 量化后,您可以通过反量化(dequantize())将模型恢复到原始精度,但这可能会导致一些质量损失。请确保您拥有足够的 GPU 内存来容纳反量化后的模型。
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizermodel = AutoModelForCausalLM.from_pretrained("facebook/opt-125m", BitsAndBytesConfig(load_in_4bit=True))
model.dequantize()
三、torch 进行模型量化与保存
import torch.quantization# 动态量化示例
model = torch.nn.Linear(100, 50)
quantized_model = torch.quantization.quantize_dynamic(model,{torch.nn.Linear}, # 指定量化层dtype=torch.qint8
)# 保存量化模型
torch.save(quantized_model.state_dict(), "dynamic_quant.pth")[1,14](@ref)