basic_quant.py 757 Bytes
Newer Older
Casper's avatar
Casper committed
1
2
3
4
5
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'lmsys/vicuna-7b-v1.5'
quant_path = 'vicuna-7b-v1.5-awq'
Casper Hansen's avatar
Casper Hansen committed
6
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
Casper's avatar
Casper committed
7
8

# Load model
Casper Hansen's avatar
Casper Hansen committed
9
# NOTE: pass safetensors=True to load safetensors
10
model = AutoAWQForCausalLM.from_pretrained(model_path, **{"low_cpu_mem_usage": True})
Casper's avatar
Casper committed
11
12
13
14
15
16
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
Casper Hansen's avatar
Casper Hansen committed
17
# NOTE: pass safetensors=True to save quantized model weights as safetensors
Casper's avatar
Casper committed
18
19
20
21
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

print(f'Model is quantized and saved at "{quant_path}"')