File size: 3,250 Bytes
64ddf8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import os
import time
import torch
from accelerate import (
infer_auto_device_map,
init_empty_weights,
load_checkpoint_in_model,
)
from awq.quantize.pre_quant import apply_awq, run_awq
from awq.quantize.quantizer import real_quantize_model_weight
from awq.utils.utils import simple_dispatch_model
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
def create_awq_model(model_id, quant_config, config_id, load_quantized, save_dir):
quantized = False
quant_path = f"{save_dir}/{model_id}-{config_id}-awq"
tokenizer = AutoTokenizer.from_pretrained(
model_id, use_fast=False, trust_remote_code=True
)
config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
# Note (Haotian): To avoid OOM after huggingface transformers 4.36.2
config.use_cache = False
if load_quantized and os.path.exists(f"{quant_path}/qmodel.pth"):
with init_empty_weights():
model = AutoModelForCausalLM.from_config(
config=config, torch_dtype=torch.float16, trust_remote_code=True
)
max_memory = {0: "20GiB", "cpu": "60GiB"}
# Infer device map
kwargs = {"max_memory": max_memory} if len(max_memory) else {}
device_map = infer_auto_device_map(
model,
no_split_module_classes=[
"OPTDecoderLayer",
"LlamaDecoderLayer",
"BloomBlock",
"MPTBlock",
"DecoderLayer",
],
**kwargs,
)
# Load checkpoint in the model
load_checkpoint_in_model(
model,
checkpoint=quant_path,
device_map=device_map,
offload_state_dict=False,
)
# Dispatch model
model = simple_dispatch_model(model, device_map=device_map)
quantized = True
model.eval()
else:
kwargs = {"torch_dtype": torch.float16, "low_cpu_mem_usage": True}
model = AutoModelForCausalLM.from_pretrained(
model_id, config=config, trust_remote_code=True, **kwargs
)
return model, tokenizer, quantized, 0
def quantize_awq_model(model, tokenizer, quant_config, model_id, config_id, save_dir):
t1 = time.time()
nbits = quant_config.pop("w_bit")
awq_results = run_awq(
model,
tokenizer,
w_bit=nbits,
q_config=quant_config,
n_samples=128,
seqlen=512,
)
intermediate_fp = f"{save_dir}/{model_id}-{config_id}-awq/intermediate.pth"
dirpath = os.path.dirname(intermediate_fp)
os.makedirs(dirpath, exist_ok=True)
torch.save(awq_results, intermediate_fp)
awq_results = torch.load(intermediate_fp, map_location="cpu")
apply_awq(model, awq_results)
real_quantize_model_weight(model, w_bit=nbits, q_config=quant_config)
t2 = time.time()
print("Took " + str(t2 - t1) + " seconds to quantize the model with AWQ")
quant_path = f"{save_dir}/{model_id}-{config_id}-awq"
quant_fp = os.path.join(quant_path, "qmodel.pth")
torch.save(model.cpu().state_dict(), quant_fp)
tokenizer.save_pretrained(quant_path)
model_file_size = os.path.getsize(quant_fp)
return model, t2 - t1, model_file_size
|