w8a8_kvcache量化场景导入的样例代码如下:
# 导入相关依赖
import torch
import torch_npu # 若需要在cpu上进行量化,可忽略此步骤
from transformers import AutoTokenizer, AutoModel
# for local path
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path='./chatglm2', trust_remote_code=True)
model = AutoModel.from_pretrained(
pretrained_model_name_or_path='./chatglm2',
trust_remote_code=True,
).npu() #若在npu上进行多卡量化时,需要先参考前提条件进行配置,并配置device_map='auto',创建model时需去掉.npu();若在cpu上进行量化时,需要配置torch_dtype=torch.float32,创建model时需去掉.npu()
# 准备校准数据,请根据实际情况修改
calib_list = ["中国的首都在哪里?",
"请做一首诗歌:",
"我想要学习python,该怎么学习?",
"请帮我写一篇关于大模型推理优化的任职报告:",
"中国最值得去的几个景点"]
#获取校准数据函数定义
def get_calib_dataset(tokenizer, calib_list):
calib_dataset = []
for calib_data in calib_list:
inputs = tokenizer([calib_data], return_tensors='pt').to(model.device)
print(inputs)
calib_dataset.append([inputs.data['input_ids'], inputs.data['attention_mask']])
return calib_dataset
dataset_calib = get_calib_dataset(tokenizer, calib_list) #校准数据获取
# 量化配置,请根据实际情况修改
from msmodelslim.pytorch.llm_ptq.llm_ptq_tools import Calibrator, QuantConfig # 导入量化配置接口
# 使用QuantConfig接口,配置量化参数,并返回量化配置实例
quant_config = QuantConfig(
a_bit=8,
w_bit=8,
disable_names=['transformer.encoder.layers.0.self_attention.query_key_value','transformer.encoder.layers.0.self_attention.dense', 'transformer.encoder.layers.0.mlp.dense_h_to_4h'],
dev_type='npu', # 在cpu进行量化时,需要配置参数dev_type='cpu',并取消参数dev_id=model.device.index的配置
dev_id=model.device.index,
act_method=3,
pr=1.0,
mm_tensor=False,
use_kvcache_quant=True
)
#使用Calibrator接口,输入加载的原模型、量化配置和校准数据,定义校准
calibrator = Calibrator(model, quant_config, calib_data=dataset_calib, disable_level='L0')
calibrator.run() #使用run()执行量化
calibrator.save('./quant_weight', save_type=[ 'numpy', 'safe_tensor']) #使用save()保存模型量化参数,请根据实际情况修改路径及保存的格式
print('Save quant weight success!')