68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
|
# Copyright Huawei Technologies Co., Ltd. 2024. All rights reserved.
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from msmodelslim.pytorch.llm_ptq.llm_ptq_tools import QuantConfig
|
|||
|
|
from msmodelslim.pytorch.llm_ptq.anti_outlier import AntiOutlierConfig
|
|||
|
|
|
|||
|
|
from atb_llm.models.llama.modeling_llama import LlamaConfig
|
|||
|
|
from examples.convert.model_slim.get_calibration_dataset import load_jsonl
|
|||
|
|
from examples.convert.model_slim.quantifier import parse_arguments, Quantifier
|
|||
|
|
from examples.convert.convert_utils import copy_tokenizer_files, modify_config
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
args = parse_arguments()
|
|||
|
|
|
|||
|
|
rank = int(os.getenv("RANK", "0"))
|
|||
|
|
|
|||
|
|
config = LlamaConfig.from_pretrained(args.model_path)
|
|||
|
|
|
|||
|
|
disable_names = []
|
|||
|
|
if args.a_bit != 16:
|
|||
|
|
# W8A16没有回退层
|
|||
|
|
num_layers = config.num_hidden_layers
|
|||
|
|
disable_names = [f"model.layers.{layer}.mlp.down_proj" for layer in range(num_layers)]
|
|||
|
|
disable_names.append("lm_head")
|
|||
|
|
|
|||
|
|
anti_outlier_config = None
|
|||
|
|
if args.anti_method:
|
|||
|
|
anti_outlier_config = AntiOutlierConfig(anti_method=args.anti_method, dev_type="npu")
|
|||
|
|
|
|||
|
|
quant_config = QuantConfig(
|
|||
|
|
a_bit=args.a_bit,
|
|||
|
|
w_bit=args.w_bit,
|
|||
|
|
disable_names=disable_names,
|
|||
|
|
act_method=args.act_method,
|
|||
|
|
mm_tensor=False,
|
|||
|
|
dev_type=args.device_type,
|
|||
|
|
dev_id=rank,
|
|||
|
|
pr=1.0,
|
|||
|
|
fraction=args.fraction,
|
|||
|
|
co_sparse=args.co_sparse,
|
|||
|
|
do_smooth=args.do_smooth,
|
|||
|
|
use_sigma=args.use_sigma,
|
|||
|
|
sigma_factor=args.sigma_factor,
|
|||
|
|
is_lowbit=args.is_lowbit,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 默认无校准数据集
|
|||
|
|
calibration_dataset = None
|
|||
|
|
# 若存在calib_file，则使用calib_file作为校准数据集
|
|||
|
|
if args.calib_file:
|
|||
|
|
calibration_dataset = load_jsonl(args.calib_file)
|
|||
|
|
calibration_dataset = calibration_dataset
|
|||
|
|
quant_weight_generator = Quantifier(args.model_path, quant_config, anti_outlier_config, args.device_type)
|
|||
|
|
quant_weight_generator.tokenizer.pad_token_id = 0
|
|||
|
|
|
|||
|
|
tokenized_data = None
|
|||
|
|
if calibration_dataset is not None:
|
|||
|
|
tokenized_data = quant_weight_generator.get_tokenized_data(calibration_dataset)
|
|||
|
|
|
|||
|
|
quant_weight_generator.convert(tokenized_data, args.save_directory, args.disable_level)
|
|||
|
|
modify_config(
|
|||
|
|
args.model_path, args.save_directory, config.torch_dtype,
|
|||
|
|
f"w{args.w_bit}a{args.a_bit}" + ("s" if args.co_sparse else "")
|
|||
|
|
)
|
|||
|
|
copy_tokenizer_files(args.model_path, args.save_directory)
|