lm_python_script.py

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import AutoConfig
import lm_eval
import argparse
import json
import torch
import torch.nn as nn
from torch import Tensor
from types import MethodType
import torch.nn.functional as F
from utils import build_model_and_enc
from omniuquant_utils import build_omniquant_model, apply_omni
from peft import AutoPeftModelForCausalLM
from duquant_utils import build_duquant
from tqdm import tqdm

RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
RESET = "\033[0m"

parser = argparse.ArgumentParser()
parser.add_argument("--model_id", type=str, default="checkpoints/models--meta-llama--Llama-2-7b-hf")
parser.add_argument("--quant_mode", type=str, default=None, choices=["bnb", "gptq", 'awq', 'aqlm', 'rtn', 'omniquant', 'duquant', 'svdquant'])
parser.add_argument("--bit", type=int, default=4)
parser.add_argument("--max_length", type=int, default=4096)
parser.add_argument("--tasks", type=str, nargs="+", default=None, choices=[
            "wikitext", "wikitext2",
            "winogrande","piqa","hellaswag","arc_easy","arc_challenge", 'mmlu',
            "leaderboard_bbh", "leaderboard_gpqa", "leaderboard_ifeval", "leaderboard_math_hard", "leaderboard_mmlu_pro", "leaderboard_musr"
        ])
# leaderboard_bbh leaderboard_gpqa leaderboard_ifeval leaderboard_math_hard leaderboard_mmlu_pro leaderboard_musr
parser.add_argument("--dataset", type=str, default="c4", choices=["c4", "wikitext2"])
parser.add_argument("--group_size", type=int, default=128)
parser.add_argument("--load_awq", type=str, default=None, help="awq qweights in model zoo")
parser.add_argument("--load_omniquant", type=str, default=None, help="omniquant qweights in model zoo")
parser.add_argument("--load_peft", type=str, default=None, help="peft qweights in model zoo")
parser.add_argument("--official_ominiquant", action="store_true", help="official omniquant")
parser.add_argument("--update_scale", action="store_true")
parser.add_argument("--boolq", action="store_true")
parser.add_argument("--svd_lora", action="store_true")
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--stop_quant_from", type=int, default=None)
parser.add_argument("--fbq_quant_mode", type=str, default=None)
parser.add_argument("--tanh_scale", type=float, default=None)

args = parser.parse_args()
if args.tasks is None:
    args.tasks = ["wikitext2","wikitext", 
        "winogrande","piqa","hellaswag","arc_easy","arc_challenge", 'mmlu']
else:
    args.tasks = (["wikitext2"] + args.tasks) if 'wikitext2' not in args.tasks else args.tasks
if args.quant_mode == "bnb":
    assert args.bit in [4,8], "BNB: bit must be 4 or 8"
print(args)

model_id = args.model_id
tokenizer = AutoTokenizer.from_pretrained(model_id)
if args.quant_mode == "gptq":
    from transformers import GPTQConfig
    gptq_config = GPTQConfig(
        bits=args.bit, dataset=args.dataset, tokenizer=tokenizer,
        group_size=args.group_size,
    )
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", quantization_config=gptq_config, torch_dtype=torch.float16)
elif args.quant_mode == "awq":
    print("AWQ from HanLab Model zoo!")
    model, enc = build_model_and_enc(model_id, args)
elif args.quant_mode == "omniquant":
    if args.official_ominiquant:
        print("Use official implementation of Omniquant")
        model = build_omniquant_model(args)
    else:
        print("Loading Omniquant weights and apply quantization with self-implementation")
        omni_results = torch.load(args.load_omniquant, map_location="cuda")
        model = apply_omni(model_id, omni_results, args)
elif args.quant_mode == "duquant":
    print("Duquant!!")
    model = build_duquant()

elif args.quant_mode == "aqlm":
    raise NotImplementedError
elif args.quant_mode == "bnb":
    from transformers import BitsAndBytesConfig
    nf4_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16
    )
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", quantization_config=nf4_config, torch_dtype=torch.float16)
elif args.quant_mode == "rtn":
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
    def rtn_forward(self, input: Tensor) -> Tensor:
        if not self.rtn_inited:
            self.rtn_inited = True
            w = self.weight.data.clone()
            dtype = self.weight.dtype
            bit, group_size = args.bit, args.group_size
            if args.group_size<=0:
                group_size = w.shape[1]
            n = 2 ** bit -1
            assert len(w.shape)==2
            assert w.shape[1] % group_size==0

            qshape = (-1, group_size)
            wmin = torch.amin(w.view(qshape),dim=1,keepdim=True)
            wmax = torch.amax(w.view(qshape),dim=1,keepdim=True)
            scale = torch.clamp((wmax - wmin) / n, min=1e-6)
            zero = torch.round(wmin / scale)
            w = w.view(qshape).div(scale).round().sub(zero).clamp(0, n).add(zero).mul(scale).view(*self.weight.shape)
            w = w.to(dtype).contiguous()

            self.weight.copy_(w)
        return F.linear(input, self.weight, self.bias)
        # bit, group_size = args.bit, args.group_size
        # n = 2 ** bit -1
        # qshape = (-1, group_size)
        # wmin = torch.amin(w.view(qshape),dim=1,keepdim=True)
        # wmax = torch.amax(w.view(qshape),dim=1,keepdim=True)
        # scale = torch.clamp((wmax - wmin) / n, min=1e-6)
        # zero = torch.round(wmin / scale)
        # w = self.weight.view(qshape).div(scale).round().sub(zero).clamp(0, n).add(zero).mul(scale).view(*self.weight.shape)
        # print(w.shape, self.weight.shape)
        # print(w.dtype, self.weight.dtype)
        # return F.linear(input, w, self.bias)

    for layer in model.model.layers:
        for name, module in layer.named_modules():
            if isinstance(module, nn.Linear):
                module.rtn_inited = False
                # module.forward = MethodType(rtn_forward, module)
                with torch.no_grad():
                    w = module.weight.data.clone()
                    dtype = module.weight.dtype
                    bit, group_size = args.bit, args.group_size
                    if args.group_size<=0:
                        group_size = w.shape[1]
                    n = 2 ** bit -1
                    assert len(w.shape)==2
                    assert w.shape[1] % group_size==0

                    qshape = (-1, group_size)
                    wmin = torch.amin(w.view(qshape),dim=1,keepdim=True)
                    wmax = torch.amax(w.view(qshape),dim=1,keepdim=True)
                    scale = torch.clamp((wmax - wmin) / n, min=1e-6)
                    zero = torch.round(wmin / scale)
                    w = w.view(qshape).div(scale).round().sub(zero).clamp(0, n).add(zero).mul(scale).view(*module.weight.shape)
                    w = w.to(dtype).contiguous()
                    module.weight.copy_(w)
elif args.quant_mode == "svdquant":
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
    for layer in tqdm(model.model.layers):
        for name, module in layer.named_modules():
            if isinstance(module, nn.Linear):
                module.rtn_inited = False
                # module.forward = MethodType(rtn_forward, module)
                with torch.no_grad():
                    w = module.weight.data.clone()
                    dtype = module.weight.dtype
                    bit, group_size = args.bit, args.group_size
                    if args.group_size<=0:
                        group_size = w.shape[1]
                    n = 2 ** bit -1
                    assert len(w.shape)==2
                    assert w.shape[1] % group_size==0


                    w_sub = 0
                    u,s,v = torch.svd(w.float())
                    r = 128
                    s[r:] = 0
                    w_sub = u @ torch.diag(s) @ v.t()
                    w = w - w_sub.half()


                    qshape = (-1, group_size)
                    wmin = torch.amin(w.view(qshape),dim=1,keepdim=True)
                    wmax = torch.amax(w.view(qshape),dim=1,keepdim=True)
                    scale = torch.clamp((wmax - wmin) / n, min=1e-6)
                    zero = torch.round(wmin / scale)
                    w = w.view(qshape).div(scale).round().sub(zero).clamp(0, n).add(zero).mul(scale).view(*module.weight.shape)
                    w = w.to(dtype).contiguous()
                    print((w + w_sub - module.weight).pow(2).mean())
                    module.weight.copy_(w + w_sub)
else:
    model = model_id
    if args.load_peft is None:
        args.bit = 0
ppl = None
lm = None
if "wikitext2" in args.tasks:
    from utils import llama_eval
    from datasets import load_dataset
    from peft.tuners.lora import LoraLayer
    testdata = load_dataset('Salesforce/wikitext', 'wikitext-2-raw-v1', split='test')
    testenc = tokenizer("\n\n".join(testdata['text']), return_tensors='pt')

    if type(model) is str:
        lm = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
        # for module in lm.modules():
        #     if isinstance(module, LoraLayer):
        #         module.pseudo_quantize(group_size=args.group_size, bit=3)
        # lm = AutoPeftModelForCausalLM.from_pretrained(pretrained_model_name_or_path, model_id, device_map="auto", torch_dtype=torch.float16)
    else:
        # try:
        #     lm = model.to('cuda', torch.float16)
        # except Exception as e:
        #     print(e)
        #     lm = model
        lm = model

    if args.load_peft is not None:
        print("Loading PEFT weights from", args.load_peft)
        from peft import PeftModel
        lm = PeftModel.from_pretrained(lm, args.load_peft, torch_dtype=torch.float16)
        # lm = AutoPeftModelForCausalLM.from_pretrained(lm, args.load_peft, torch_dtype=torch.float16)
        for _name, module in lm.named_modules():
            if isinstance(module, LoraLayer):
                if args.svd_lora:
                    with torch.no_grad():
                        nn.init.zeros_(module.lora_A.default.weight)
                        nn.init.zeros_(module.lora_B.default.weight)
                if args.stop_quant_from is not None:
                    block_id = int(_name.split("layers.")[-1].split(".")[0])
                    if block_id >= args.stop_quant_from:
                        nn.init.zeros_(module.lora_A.default.weight)
                        nn.init.zeros_(module.lora_B.default.weight)
                        print(RED + f"Stop quantize {_name}!" + RESET)
                        continue
                setattr(module, "update_scale", args.update_scale)
                setattr(module, "quant_mode", args.fbq_quant_mode)
                setattr(module, "tanh_scale", args.tanh_scale)
                setattr(module, "group_size", args.group_size)
                setattr(module, "bit", args.bit)
                module.pseudo_quantize(group_size=args.group_size, bit=args.bit, quant_mode=args.fbq_quant_mode)
                if args.svd_lora:
                    with torch.no_grad():
                        delta = module.org_base_weight - module.base_layer.weight
                        u,s,v = torch.svd(delta.to(torch.float32))
                        top_singular_value = s[0]
                        u_top = u[:, 0].unsqueeze(1)  # Make u_top a column vector (Nx1)
                        v_top = v[:, 0].unsqueeze(0)  # Make v_top a row vector (1xM)
                        # delta_approx = top_singular_value * (u_top @ v_top)
                        module.lora_A.default.weight.copy_(v_top * top_singular_value.sqrt())
                        module.lora_B.default.weight.copy_(u_top * top_singular_value.sqrt())




    seqlen = lm.seqlen if hasattr(lm, "seqlen") else None
    lm.seqlen = args.max_length
    ppl = llama_eval(lm, testenc, "cuda")
    lm.seqlen = seqlen

if any(t in args.tasks for t in [
        "wikitext", 
        "winogrande","piqa","hellaswag","arc_easy","arc_challenge", 'mmlu',
        "leaderboard_bbh", "leaderboard_gpqa", "leaderboard_ifeval", "leaderboard_math_hard", "leaderboard_mmlu_pro", "leaderboard_musr"

    ]):
    if lm is not None:
        model = lm
    else:
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)

    model = lm_eval.models.huggingface.HFLM(
        pretrained=model,
        max_length=args.max_length,
        dtype = torch.float16,
        batch_size=args.batch_size,
    )

    if lm is None:
        for module in model.pretrained.modules():
            if isinstance(module, LoraLayer):
                setattr(module, "update_scale", args.update_scale)
                setattr(module, "quant_mode", args.fbq_quant_mode)
                setattr(module, "tanh_scale", args.tanh_scale)
                setattr(module, "group_size", args.group_size)
                setattr(module, "bit", args.bit)
                module.pseudo_quantize(group_size=args.group_size, bit=args.bit, quant_mode=args.fbq_quant_mode)

    task_manager = lm_eval.tasks.TaskManager()

    # Setting `task_manager` to the one above is optional and should generally be done
    # if you want to include tasks from paths other than ones in `lm_eval/tasks`.
    # `simple_evaluate` will instantiate its own task_manager if it is set to None here.
    results = lm_eval.simple_evaluate( # call simple_evaluate
        model=model,
        tasks= ['boolq'] if args.boolq else [t for t in args.tasks if t!="wikitext2"],
        # [
        #     "wikitext", 
        #     "winogrande",
        #     "piqa",
        #     "hellaswag",
        #     "arc_easy",
        #     "arc_challenge", 
        #     'mmlu',
        #     # 'boolq'
        # ],
        num_fewshot=0,
        task_manager=task_manager,
        batch_size=args.batch_size,
        # ...
    )
    results["results"]["wikitext2"] = {"word_perplexity,none": ppl.item(),}
    # print(results['results'])

    with open(f"output/results_m<{args.model_id.split('/')[-1]}>_q<{args.quant_mode}>_b{args.bit}g{args.group_size}_len<{args.max_length}>.json", "w") as f:
        json.dump(results["results"], f, indent=2)

    for task, res in results["results"].items():
        if res.get("acc,none", None):
            acc = res['acc,none']*100
            print(f"{task}, acc, {res['acc,none']*100:.2f}")
        if res.get("acc_norm,none", None):
            acc = res['acc_norm,none']*100
            print(f"{task}, acc_norm, {res['acc_norm,none']*100:.2f}")
        if res.get("word_perplexity,none", None):
            print(f"{task}, word_perplexity, {res['word_perplexity,none']:.2f}")
    print("="*40)

    tasks=[
            "wikitext2", "wikitext",
        ]
    # task = "wikitext2"
    h_results = []
    for task in tasks:
        if task not in results["results"].keys():
            continue
        res = results["results"][task]
        print(f"{task}, {res['word_perplexity,none']:.2f}")
        h_results += [f"{res['word_perplexity,none']:.2f}"]

    tasks=[
            "arc_challenge", 
            "arc_easy",
            "hellaswag",
            'mmlu',
            "piqa",
            "winogrande",
            'boolq',
        ]
    for task in tasks:
        if task not in results["results"].keys():
            continue
        res = results["results"][task]
        acc1 = res.get('acc,none', 0)
        acc2 = res.get('acc_norm,none', 0)
        acc = max(acc1, acc2)
        print(f"{task}, {acc*100:.2f}")
        h_results += [f"{acc*100:.2f}"]
    
    for r in h_results:
        print(r, end=" ")
    print()
    # import time
    # saved_model = f"{time.time()}_model.pt"
    # torch.save(model, saved_model)
    # print(f"Save to {saved_model}")