rakesh9177/Quantization

1import gradio as gr2import tqdm3import torch4from torch import nn5from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline6from functools import partial7import gc8 9def get_model_size(model: nn.Module, data_width=16, group_size=-1):10 11    if group_size != -1:12        data_width += (16 + 4) / group_size13 14    num_elements = 015    for param in model.parameters():16        num_elements += param.numel()17    return num_elements * data_width18 19Byte = 820KiB = 1024 * Byte21MiB = 1024 * KiB22GiB = 1024 * MiB23 24# core quantization method (simulated quantization)25def pseudo_quantize_tensor(w, n_bit=4, q_group_size=-1):26    org_w_shape = w.shape27    if q_group_size > 0:28        assert org_w_shape[-1] % q_group_size == 029        w = w.reshape(-1, q_group_size)30 31    assert w.dim() == 232 33    # Calculate the maximum (\alpha) and minimum values (\beta) in the tensor.34    max_val = w.amax(dim=1, keepdim=True)35    assert max_val.dim() == 2 and max_val.size(0) == w.size(0) and max_val.size(1) == 136    min_val = w.amin(dim=1, keepdim=True)37    assert min_val.dim() == 2 and min_val.size(0) == w.size(0) and min_val.size(1) == 138 39    # Calculate the scale factor and zero point.  (Formula 1 & 2)40    max_int = 2 ** n_bit - 141    scales = (max_val - min_val).clamp(min=1e-5) / max_int42    assert scales.shape == max_val.shape43    zeros = (-torch.round(min_val / scales)).clamp_(0, max_int)44    assert scales.shape == min_val.shape45 46    assert torch.isnan(scales).sum() == 047    assert torch.isnan(w).sum() == 048 49    # Quantize W: Map values in the range [\beta, \alpha] to lie within [0, 2^b - 1] (Formula 3)50    w = torch.clamp(torch.round(w / scales) + zeros, 0, max_int)51    assert w.dim() == 2 and w.size(0) == scales.size(0) and w.size(1) == q_group_size52 53    # Dequantize W (pseudo quantization, the inverse transformation of Formula 3)54    w = (w - zeros) * scales55    assert w.dim() == 2 and w.size(0) == scales.size(0) and w.size(1) == q_group_size56 57    assert torch.isnan(w).sum() == 058 59    w = w.reshape(org_w_shape)60    return w61 62@torch.no_grad()63def pseudo_quantize_model_weight(64    model, w_bit, q_group_size,65):66    for n, m in model.named_modules():67        if isinstance(m, nn.Linear):68            m.weight.data = pseudo_quantize_tensor(m.weight.data, n_bit=w_bit, q_group_size=q_group_size)69            70            71    72            73# Load the tokenizer and model74model_path = "facebook/opt-125m"75 76model_q_path = "facebook/opt-125m_3bit"77offload_folder = "offload"78model = AutoModelForCausalLM.from_pretrained(model_q_path, device_map="auto", offload_folder=offload_folder)79generator = pipeline('text-generation', model="facebook/opt-125m_3bit")80#generator_q = pipeline('text-generation', model="facebook/opt-125m-awq")81 82 83def generate_text_pip(prompt):84    generated_text = generator(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']85    return generated_text86'''87def generate_text_pip_q(prompt):88    generated_text = generator_q(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']89    return generated_text90'''91print(generator("I went to boston and"))92#print("quantized model",generator_q("I went to boston and"))93 94'''95def generate_text(prompt):96    inputs = tokenizer(prompt, return_tensors="pt")97    output = model(**inputs)98    logits = output.logits99    predicted_ids = logits.argmax(-1)100    #generated_text = tokenizer.decode(predicted_ids[0], skip_special_tokens=True)101    generated_text = tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)[0]102    return generated_text103 104def generate_text_from_quantized(prompt):105    inputs = tokenizer(prompt, return_tensors="pt")106    output = model_q(**inputs)107    logits = output.logits108    predicted_ids = logits.argmax(-1)109    #generated_text = tokenizer.decode(predicted_ids[0], skip_special_tokens=True)110    generated_text = tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)[0]111    return generated_text112'''113# Create a Gradio interface114model_size = get_model_size(model, data_width=32, group_size=128)115description  = f"Model Name :              OPT-1.3b <br>Original Model Size :     4.21 GB <br>Quantized Model Size :   {model_size/MiB:.2f} MiB"116iface = gr.Interface(fn=generate_text_pip, inputs="text", outputs="text", description=description)117 118#iface_2 = gr.Interface(fn=generate_text_pip_q, inputs="text", outputs="text")119 120iface.launch()121#app = gr.TabbedInterface([iface, iface_2],["Normal", "Quantized"])122 123# Launch the Gradio app124#app.launch()125