w8a16 per-group量化场景

# 导入相关依赖
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()

# 准备校准数据，请根据实际情况修改，HQQ和MinMax的W8A16 Data-Free模式下请忽略此步骤
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)  #校准数据获取

"""
# AWQ算法量化场景下，需要配置此步骤
# 执行离群值抑制的操作
from msmodelslim.pytorch.llm_ptq.anti_outlier import AntiOutlier, AntiOutlierConfig
anti_config = AntiOutlierConfig(
    w_bit=8, 
    a_bit=16, 
    anti_method='m3', 
    dev_id=model.device.index,
    dev_type='npu',   # 在cpu进行量化时，需配置参数dev_type='cpu'，并取消参数dev_id=model.device.index的配置
    w_sym=True
)
anti_outlier = AntiOutlier(model, calib_data=dataset_calib, cfg=anti_config)
anti_outlier.process()
# 可选配置结束
"""   

# 量化配置，请根据实际情况修改
from msmodelslim.pytorch.llm_ptq.llm_ptq_tools import Calibrator, QuantConfig    # 导入量化配置接口
# 使用QuantConfig接口，配置量化参数，并返回量化配置实例
quant_config = QuantConfig(
    w_bit=8,     
    a_bit=16,         
    disable_names=[], 
    dev_id=model.device.index,
    dev_type='npu',  # 在cpu进行量化时，需配置参数dev_type='cpu'，并取消参数dev_id=model.device.index的配置
    w_sym=True,
    mm_tensor=False,     
    w_method='MinMax',   # MinMax和AWQ算法量化场景使用默认值'MinMax'，GPTQ算法量化场景需配置为'GPTQ',HQQ算法量化场景需配置为'HQQ'      
    is_lowbit=True,
    open_outlier=False,
    group_size=64
)  
#使用Calibrator接口，输入加载的原模型、量化配置和校准数据，定义校准
calibrator = Calibrator(model, quant_config, disable_level='L0')
calibrator.run()     #使用run()执行量化
calibrator.save('./quant_weight', save_type=['numpy', 'safe_tensor'])      #使用save()保存模型量化参数，请根据实际情况修改路径及保存的格式
print('Save quant weight success!')