radna/eval_llm

1from awq import AutoAWQForCausalLM2from transformers import AutoTokenizer3 4# take in a model path and quantization args5import argparse6 7parser = argparse.ArgumentParser()8parser.add_argument(9    "--model_path", type=str, default="deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"10)11parser.add_argument("--quant_path", type=str, default="r1-14b-awq-max-ptb")12args = parser.parse_args()13 14model_path = args.model_path15quant_path = args.quant_path16quant_config = {17    "zero_point": True,18    "q_group_size": 128,19    "w_bit": 4,20    "version": "GEMM",21}22 23 24# Load model25model = AutoAWQForCausalLM.from_pretrained(model_path)26tokenizer = AutoTokenizer.from_pretrained(27    model_path,28    trust_remote_code=True,29)30 31 32# Quantize33model.quantize(34    tokenizer,35    quant_config=quant_config,36    # calib_data="neuralmagic/LLM_compression_calibration",37    # calib_data=get_long_dataset(),38    # calib_data="ptb",39    # max_calib_samples=128,40    # max_calib_seq_len=12288,41    # n_parallel_calib_samples=128,42)43 44# Save quantized model45model.save_quantized(quant_path)46tokenizer.save_pretrained(quant_path)47 48print(f'Model is quantized and saved at "{quant_path}"')49