Add files using upload-large-folder tool

lm-quant-toolkit/.deps/hqq/examples/hf/llama2_chat_hf_hub_example.py

@@ -0,0 +1,65 @@
+model_id  = 'mobiuslabsgmbh/Llama-2-7b-chat-hf-4bit_g64-HQQ'
+#model_id = 'mobiuslabsgmbh/Llama-2-13b-chat-hf-4bit_g64-HQQ'
+#model_id = 'mobiuslabsgmbh/Llama-2-70b-chat-hf-2bit_g16_s128-HQQ'
+
+from hqq.engine.hf import HQQModelForCausalLM, AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model     = HQQModelForCausalLM.from_quantized(model_id)
+
+##########################################################################################################
+import transformers
+from threading import Thread
+
+from sys import stdout
+def print_flush(data):
+    stdout.write("\r" + data)
+    stdout.flush()
+
+#Adapted from https://huggingface.co/spaces/huggingface-projects/llama-2-7b-chat/blob/main/app.py
+def process_conversation(chat):
+    system_prompt = chat['system_prompt']
+    chat_history  = chat['chat_history']
+    message       = chat['message']
+
+    conversation = []
+    if system_prompt:
+        conversation.append({"role": "system", "content": system_prompt})
+    for user, assistant in chat_history:
+        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+    conversation.append({"role": "user", "content": message})
+
+    return tokenizer.apply_chat_template(conversation, return_tensors="pt").to('cuda')
+
+def chat_processor(chat, max_new_tokens=100, do_sample=True):
+    tokenizer.use_default_system_prompt = False
+    streamer = transformers.TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+
+    generate_params = dict(
+        {"input_ids": process_conversation(chat)},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        top_p=0.90,
+        top_k=50,
+        temperature= 0.6,
+        num_beams=1,
+        repetition_penalty=1.2,
+    )
+
+    t = Thread(target=model.generate, kwargs=generate_params)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        print_flush("".join(outputs))
+
+    return outputs
+
+###################################################################################################
+
+outputs = chat_processor({'system_prompt':"You are a helpful assistant.",
+                        'chat_history':[],
+                        'message':"How can I build a car?"
+                        }, 
+                         max_new_tokens=1000, do_sample=False)

lm-quant-toolkit/.deps/hqq/examples/hf/whisper.py

@@ -0,0 +1,98 @@
+# Tested with torch nightly, 4090 
+# pip uninstall torch -y; pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
+##############################################################################################
+
+import torch
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+
+model_id      = "openai/whisper-medium"
+#model_id     = "distil-whisper/distil-large-v3"
+
+compute_dtype = torch.bfloat16
+device        = "cuda:0"
+
+model     = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=compute_dtype) 
+processor = AutoProcessor.from_pretrained(model_id)
+
+
+##############################################################################
+#No quantize
+#model = model.to(device)
+
+##############################################################################
+#Quantize
+from hqq.models.hf.base import AutoHQQHFModel
+from hqq.core.quantize import *
+
+quant_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_scale=False, quant_zero=False, axis=1) 
+HQQLinear.set_backend(HQQBackend.PYTORCH)
+
+AutoHQQHFModel.quantize_model(model.model.encoder, quant_config=quant_config, compute_dtype=compute_dtype, device=device)
+AutoHQQHFModel.quantize_model(model.model.decoder, quant_config=quant_config, compute_dtype=compute_dtype, device=device)
+
+from hqq.utils.patching import prepare_for_inference
+prepare_for_inference(model.model.encoder)
+prepare_for_inference(model.model.decoder, backend="torchao_int4")
+
+model.model.encoder.forward = torch.compile(model.model.encoder.forward, mode="reduce-overhead", fullgraph=True)
+model.model.decoder.forward = torch.compile(model.model.decoder.forward, mode="reduce-overhead", fullgraph=True)
+# ##############################################################################
+
+import time 
+import numpy as np 
+
+if(model_id=="openai/whisper-medium"):
+	encoder_input = torch.randn([1, 80, 3000], dtype=compute_dtype, device=device)
+if(model_id=="distil-whisper/distil-large-v3"):
+	encoder_input = torch.randn([1, 128, 3000], dtype=compute_dtype, device=device)
+
+def run_encoder():
+	with torch.no_grad():
+		model.model.encoder(encoder_input)
+	torch.cuda.synchronize()
+
+t = []
+for _ in range(200):
+	t1 = time.time()
+	run_encoder()
+	t2 = time.time()
+	t.append(t2-t1)
+print("Encoder", np.mean(t[-100:]), "sec / sample")
+
+
+decoder_input = torch.randint(0, 1000, [1, 1], dtype=torch.int64, device=device)
+def run_decoder():
+	with torch.no_grad():
+		out = model.model.decoder(decoder_input)
+	torch.cuda.synchronize()
+
+
+t = []
+for _ in range(200):
+	t1 = time.time()
+	run_decoder()
+	t2 = time.time()
+	t.append(t2-t1)
+print("Decoder", np.mean(t[-100:]), "sec / sample")
+
+
+#openai/whisper-medium | RTX 4090
+#Encoder: use default backend
+#fp16                         : 0.0234 sec / sample
+#hqq 4-bit (default,compiled) : 0.0124 sec / sample | 1.89x faster 
+
+
+#Decoder: use torchao backend to decode 1 token at a time
+#fp16                         : 0.01080 sec / sample
+#hqq 4-bit (ao_int4, compiled): 0.000928 sec / sample |  11.63x faster
+
+
+#distil-whisper/distil-large-v3 | RTX 4090
+#Encoder: use default backend
+#fp16                         : 0.03738  sec / sample
+#hqq 4-bit (default,compiled) : 0.01869 sec / sample | 2x faster 
+
+
+#Decoder: use torchao backend to decode 1 token at a time
+#fp16                         : 0.002592 sec / sample
+#hqq 4-bit (ao_int4, compiled): 0.000326 sec / sample |  7.95x faster

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/.gitignore

@@ -0,0 +1,3 @@
+snapshots/
+snapshots-*/
+results/

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/.pdbrc

@@ -0,0 +1,10 @@
+# b hqq/core/quantize.py:677
+# c
+# b eval_model.py:34
+# b bench.py:221
+# b hqq/models/base.py:132
+# b quant_llama2_hqq_demo.py:30
+# b hqq/models/base.py:265
+# b debug-quant-mem.py:18
+# b hqq/models/hf/llama.py:71
+# b hqq/models/base.py:260

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/adapter/autoawq.py

@@ -0,0 +1,32 @@
+import os
+import time
+import transformers
+
+from awq import AutoAWQForCausalLM
+
+
+def create_autoawq_model(model_id, quant_config, config_id, load_quantized, save_dir):
+    quantized = False
+    quant_path = f"{save_dir}/{model_id}-{config_id}-awq"
+    if load_quantized and os.path.exists(quant_path):
+        model = AutoAWQForCausalLM.from_quantized(quant_path, "", fuse_layers=False)
+        tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
+        quantized = True
+        model = model.cuda()
+    else:
+        tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
+        model = AutoAWQForCausalLM.from_pretrained(model_id)
+    return model, tokenizer, quantized
+
+
+def quantize_autoawq_model(model, tokenizer, quant_config, model_id, config_id, save_dir):
+    t1 = time.time()
+    model.quantize(tokenizer, quant_config=quant_config)
+    t2 = time.time()
+    print('Took ' + str(t2 - t1) + ' seconds to quantize the model with AutoAWQ')
+    quant_path = f"{save_dir}/{model_id}-{config_id}-awq"
+    model.save_quantized(quant_path)
+    tokenizer.save_pretrained(quant_path)
+    return model, t2 - t1
+
+

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/adapter/awq.py

@@ -0,0 +1,94 @@
+import os
+import time
+import torch
+
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import AutoConfig
+from accelerate import (
+    init_empty_weights,
+    infer_auto_device_map,
+    load_checkpoint_in_model,
+)
+
+from awq.quantize.pre_quant import run_awq
+from awq.quantize.pre_quant import apply_awq
+from awq.quantize.quantizer import real_quantize_model_weight
+from awq.utils.utils import simple_dispatch_model
+
+
+max_memory = {}
+
+
+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
+            )
+        # 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=True,
+        )
+        # 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
+
+
+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"
+    torch.save(model.cpu().state_dict(), f"{quant_path}/qmodel.pth")
+    tokenizer.save_pretrained(quant_path)
+    return model, t2 - t1

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/debug-bench.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python -m pdb bench.py

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/eval_model.py

@@ -0,0 +1,115 @@
+import gc
+import time
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+
+
+def cleanup():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def eval_ptb(model, tokenizer, max_length=1024, stride=512, verbose=True):
+    dataset = load_dataset("ptb_text_only", "penn_treebank", split="test")
+    return eval_ppl(
+        "ptb",
+        model,
+        tokenizer,
+        dataset,
+        text_column="sentence",
+        max_length=max_length,
+        stride=stride,
+        verbose=verbose,
+    )
+
+
+def eval_c4(model, tokenizer, max_length=1024, stride=512, verbose=True):
+    dataset = load_dataset(
+        "allenai/c4",
+        data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+        split="validation",
+        download_mode="reuse_dataset_if_exists",
+    )
+    # pick first 1100
+    dataset = dataset[:1100]
+    return eval_ppl(
+        "C4",
+        model,
+        tokenizer,
+        dataset,
+        text_column="text",
+        max_length=max_length,
+        stride=stride,
+        verbose=verbose,
+    )
+
+
+def eval_wikitext2(model, tokenizer, max_length=1024, stride=512, verbose=True):
+    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+    return eval_ppl(
+        "wikitext",
+        model,
+        tokenizer,
+        dataset,
+        text_column="text",
+        max_length=max_length,
+        stride=stride,
+        verbose=verbose,
+    )
+
+
+# Adapted from https://huggingface.co/transformers/v4.2.2/perplexity.html
+def eval_ppl(
+    ds_type,
+    model,
+    tokenizer,
+    dataset,
+    text_column="text",
+    max_length=1024,
+    stride=512,
+    verbose=True,
+):
+    model.eval()
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+    tokenizer.add_eos_token = False
+
+    encodings = tokenizer("\n\n".join(dataset[text_column]), return_tensors="pt")
+
+    encodings["input_ids"] = encodings["input_ids"].to("cuda")
+
+    lls, t = [], []
+    for i in tqdm(
+        range(0, encodings["input_ids"].size(1), stride),
+        desc=ds_type,
+        disable=not verbose,
+    ):
+        begin_loc = max(i + stride - max_length, 0)
+        end_loc = min(i + stride, encodings["input_ids"].size(1))
+        trg_len = end_loc - i
+        input_ids = encodings["input_ids"][:, begin_loc:end_loc]
+        target_ids = input_ids.clone()
+        target_ids[:, :-trg_len] = -100  # ignore context
+
+        t1 = time.time()
+        with torch.no_grad():
+            log_likelihood = model(input_ids, labels=target_ids).loss * trg_len
+        torch.cuda.synchronize()
+        t2 = time.time()
+        t.append((t2 - t1))
+        lls.append(log_likelihood)
+
+        del input_ids, target_ids
+
+    ppl = np.round(float(torch.exp(torch.stack(lls).sum() / end_loc)), 4)
+    pred_time = np.round(np.mean(t), 3)
+    if verbose:
+        print(f"{ds_type} perplexity: {ppl}, time: {pred_time} sec")
+
+    del encodings
+    cleanup()
+
+    return ppl, pred_time

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/fp16_llama2_demo.py

@@ -0,0 +1,26 @@
+import torch
+
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from eval_model import eval_wikitext2
+
+#Settings
+######################################################################################
+hf_auth    = None #HuggingFace token
+cache_path = ''   #cache directory to store data
+
+#Chose a model
+# model_id  = "meta-llama/Llama-2-7b-hf"
+# model_id  = "meta-llama/Meta-Llama-3-8B"
+model_id  = "meta-llama/Llama-2-13b-hf"
+#model_id  = "meta-llama/Llama-2-70b-hf"
+
+#Load model on the CPU
+######################################################################################
+model     = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+#Evaluate the quantized model
+######################################################################################
+eval_wikitext2(model, tokenizer, verbose=True)
+

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/quant_llama2_awq_demo.py

@@ -0,0 +1,40 @@
+import torch, transformers
+
+#Settings
+######################################################################################
+hf_auth    = None #HuggingFace token
+cache_path = ''   #cache directory to store data
+
+#Chose a model
+# model_id  = "meta-llama/Meta-Llama-3-8B"
+# model_id  = "meta-llama/Llama-2-7b-hf"
+model_id  = "meta-llama/Llama-2-13b-hf"
+#model_id  = "meta-llama/Llama-2-70b-hf"
+
+#AWQ settings
+######################################################################################
+from awq import AutoAWQForCausalLM
+import gc, time
+
+# Load model
+tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_auth)
+model     = AutoAWQForCausalLM.from_pretrained(model_id,         use_auth_token=hf_auth, resume_download=True)
+
+#quant_config = {"w_bit": 4, "q_group_size": 128, "zero_point": True, 'version':'GEMM'}
+#quant_config = {"w_bit": 4, "q_group_size": 64, "zero_point": True, 'version':'GEMM'}
+quant_config = {"w_bit": 3, "q_group_size": 64, "zero_point": True, 'version':'gemv_fast'}
+
+t1 = time.time()
+model.quantize(tokenizer, quant_config=quant_config)
+t2 = time.time()
+print('Took ' + str(t2-t1) + ' seconds to quantize the model with AWQ')
+
+model = model.cuda()
+torch.cuda.empty_cache()
+gc.collect()
+
+#Evaluate the quantized model
+######################################################################################
+from eval_model import eval_wikitext2
+eval_wikitext2(model, tokenizer, verbose=True)
+

lm-quant-toolkit/.deps/hqq/examples/llama2_benchmark/run-bench.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+# export HF_HOME=/data/hugginface
+
+if [ ! -d logs ]; then
+    mkdir logs
+fi
+python bench.py 2>&1 | tee logs/bench-$(date +%Y%m%d%H%M%S).log

lm-quant-toolkit/.deps/hqq/examples/lora/train_hqq_lora_example.py

@@ -0,0 +1,220 @@
+#Settings
+######################################################################################
+hf_auth    = None #HuggingFace token
+cache_path = ''   #cache directory to store data
+
+#Chose a model
+model_id  = "meta-llama/Llama-2-7b-hf" 
+#model_id  = "meta-llama/Llama-2-13b-hf" 
+#model_id  = "meta-llama/Llama-2-70b-hf" 
+
+#HQQ Quantize
+######################################################################################
+from hqq.engine.hf import HQQModelForCausalLM, AutoTokenizer
+model     = HQQModelForCausalLM.from_pretrained(model_id, use_auth_token=hf_auth, cache_dir=cache_path)
+tokenizer = AutoTokenizer.from_pretrained(model_id,       use_auth_token=hf_auth, cache_dir=cache_path) 
+
+#Quantize the model
+from hqq.core.quantize import *
+quant_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_scale=False, quant_zero=False)
+model.quantize_model(quant_config=quant_config)
+
+#Add Peft
+######################################################################################
+from hqq.core.peft import PeftUtils
+from hqq.core.quantize import *
+
+train_dtype      = torch.float32 
+base_lora_params = {'lora_type':'default', 'r':32, 'lora_alpha':64, 'dropout':0.05, 'train_dtype':train_dtype}
+lora_params      = {'self_attn.q_proj': base_lora_params,
+		    'self_attn.k_proj': base_lora_params,
+		    'self_attn.v_proj': base_lora_params,
+		    'self_attn.o_proj': base_lora_params,
+		    'mlp.gate_proj'   : None,
+		    'mlp.up_proj'     : None,
+		    'mlp.down_proj'   : None}
+
+#Apply LoRA
+PeftUtils.add_lora(model, lora_params)
+
+#Dataset 
+######################################################################################
+from datasets import load_dataset, Dataset
+from tqdm import tqdm
+import transformers
+import numpy as np 
+import random
+
+tokenizer.pad_token     = tokenizer.eos_token 
+tokenizer.padding_side  = "right" 
+tokenizer.add_bos_token = False
+tokenizer.add_eos_token = False
+
+batch_size  = 1 
+num_epochs  = 1
+grad_acc    = 1
+max_tokens  = 256 
+max_samples = 5000
+
+#OpenAssistant
+##########################################################################
+dataset     = load_dataset("timdettmers/openassistant-guanaco", split="train") 
+dataset_val = load_dataset("timdettmers/openassistant-guanaco", split="test")  
+
+def pre_process_chat(chat):
+    #add proper chat preprocessing (bos/eos tokens, etc.)
+	return chat
+
+def assitant_prompt(prompt):
+	return '### Human:' + prompt + '\n### Assistant:'
+
+
+random.seed(100)
+idx         = random.sample(range(len(dataset)), min(max_samples, len(dataset)))
+dataset     = Dataset.from_dict({'text':[pre_process_chat(dataset[i]['text']) for i in tqdm(idx)]})
+dataset_val = Dataset.from_dict({'text':[pre_process_chat(dataset_val[i]['text']) for i in range(len(dataset_val))]})
+
+#####################################################################################
+#Train
+from trl import SFTTrainer
+
+grad_acc   = 2
+logging_st = 1
+max_steps  = -1
+lr         = 1e-4 
+batch_size = 1
+n_epochs   = 1
+
+training_args = transformers.TrainingArguments(
+    output_dir='.',	
+    per_device_train_batch_size=batch_size,
+    #per_device_eval_batch_size=batch_size,
+    gradient_accumulation_steps=grad_acc,
+    learning_rate=lr,
+    logging_steps=logging_st,
+    num_train_epochs=n_epochs,
+    max_steps=max_steps,
+    #evaluation_strategy = "epoch",
+    remove_unused_columns=False,
+    #logging_strategy="epoch",
+    fp16=train_dtype==torch.float32,
+    max_grad_norm=1.0,
+    save_steps=10000000,
+    lr_scheduler_type= "linear", 
+)
+
+#Wrap model to avoid accelerate issues 
+class WrappedModel(torch.nn.Module):
+	def __init__(self, model):
+		super().__init__()
+		self.model = model
+
+	def forward(self, *args, **kwargs):
+		return self.model.forward(*args, **kwargs)
+
+	def train(self):
+		self.model.train()
+
+	def eval(self):
+		self.model.eval()
+
+	def parameters(self):
+		return self.model.parameters()
+
+trainer = SFTTrainer(
+    model=WrappedModel(model),
+    tokenizer=tokenizer,
+    max_seq_length=max_tokens,
+    train_dataset=dataset,
+    eval_dataset=None,
+    peft_config=None,
+    args=training_args,
+    dataset_text_field="text",
+)
+
+model.is_parallelizable       = False
+trainer.is_model_parallel     = False
+trainer.place_model_on_device = False
+
+model.train()
+trainer.train()
+
+#Prediction/Eval
+######################################################################################
+#from #https://huggingface.co/spaces/evaluate-metric/perplexity/blob/main/perplexity.py
+def compute_perplexity_batched(model, tokenizer, predictions, encodings=None, batch_size=1, add_start_token=True, device='cuda', max_length=None):
+    if tokenizer.pad_token is None and batch_size > 1:
+        existing_special_tokens = list(tokenizer.special_tokens_map_extended.values())
+        # check that the model already has at least one special token defined
+        assert (len(existing_special_tokens) > 0), "If batch_size > 1, model must have at least one special token to use for padding. Please use a different model or set batch_size=1."
+        # assign one of the special tokens to also be the pad token
+        tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})
+
+    if add_start_token and max_length:
+        # leave room for <BOS> token to be added:
+        assert (tokenizer.bos_token is not None), "Input model must already have a BOS token if using add_start_token=True. Please use a different model, or set add_start_token=False"
+        max_tokenized_len = max_length - 1
+    else:
+        max_tokenized_len = max_length
+
+
+    if(encodings is None):
+        encodings = tokenizer(
+            predictions,
+            add_special_tokens=False,
+            padding=True,
+            truncation=True if max_tokenized_len else False,
+            max_length=max_tokenized_len,
+            return_tensors="pt",
+            return_attention_mask=True).to(device)
+
+    encoded_texts = encodings["input_ids"]
+    attn_masks    = encodings["attention_mask"]
+
+    # check that each input is long enough:
+    if add_start_token:
+        assert torch.all(torch.ge(attn_masks.sum(1), 1)), "Each input text must be at least one token long."
+    else:
+        assert torch.all(
+            torch.ge(attn_masks.sum(1), 2)
+        ), "When add_start_token=False, each input text must be at least two tokens long. Run with add_start_token=True if inputting strings of only one token, and remove all empty input strings."
+
+    ppls = []
+    loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+
+    for start_index in tqdm(range(0, len(encoded_texts), batch_size)):
+        end_index     = min(start_index + batch_size, len(encoded_texts))
+        encoded_batch = encoded_texts[start_index:end_index]
+        attn_mask     = attn_masks[start_index:end_index]
+
+        if add_start_token:
+            bos_tokens_tensor = torch.tensor([[tokenizer.bos_token_id]] * encoded_batch.size(dim=0)).to(device)
+            encoded_batch     = torch.cat([bos_tokens_tensor, encoded_batch], dim=1)
+            attn_mask         = torch.cat([torch.ones(bos_tokens_tensor.size(), dtype=torch.int64).to(device), attn_mask], dim=1)
+
+        labels = encoded_batch
+
+        with torch.no_grad():
+            out_logits = model(encoded_batch, attention_mask=attn_mask).logits
+
+        shift_logits = out_logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        shift_attention_mask_batch = attn_mask[..., 1:].contiguous()
+
+        perplexity_batch = torch.exp(
+            (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).sum(1)
+            / shift_attention_mask_batch.sum(1))
+
+        ppls += perplexity_batch.tolist()
+
+    return np.mean(ppls)
+
+
+tokenizer.add_bos_token = True
+tokenizer.add_eos_token = False
+model.eval()
+
+#Convert lora weights to the same model dtype for faster inference
+PeftUtils.cast_lora_weights(model, dtype=torch.half)
+
+print('perplexity', compute_perplexity_batched(model=model, tokenizer=tokenizer, predictions=[s['text'] for s in dataset_val], batch_size=1, max_length=max_tokens))

lm-quant-toolkit/.deps/hqq/examples/open_clip/.pdbrc

@@ -0,0 +1,4 @@
+# b /home/justin/work/hqq/hqq/engine/open_clip.py:57
+# b /home/justin/work/hqq/hqq/models/open_clip/base.py:77
+# b /home/justin/work/hqq/examples/open_clip/vit_clip_example.py:27
+b /home/justin/work/hqq/examples/open_clip/vit_clip_example.py:62

lm-quant-toolkit/.deps/hqq/examples/open_clip/vit_clip_example.py

@@ -0,0 +1,158 @@
+import os
+from pathlib import Path
+
+import open_clip
+import torch
+from open_clip import tokenizer
+from PIL import Image
+from skimage import data_dir
+
+from hqq.core.quantize import BaseQuantizeConfig
+from hqq.engine.open_clip import HQQOpenCLIP
+
+model_ids = [
+    "laion/CLIP-ViT-B-16-laion2B-s34B-b88K",
+    "laion/CLIP-ViT-B-32-laion2B-s34B-b79K",
+    "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+    "laion/CLIP-ViT-L-14-laion2B-s32B-b82K",
+]
+
+
+def quant_models(model_ids):
+    for model_id in model_ids:
+        model = HQQOpenCLIP.create_model(model_id, device="cpu")
+
+        # Quantize settings
+        # quant_config = BaseQuantizeConfig(nbits=8, group_size=128)
+        quant_config = BaseQuantizeConfig(nbits=4, group_size=64)
+        # quant_config = BaseQuantizeConfig(nbits=3, group_size=64)
+        # quant_config = BaseQuantizeConfig(nbits=2, group_size=16, quant_scale=True)
+
+        # Quantize
+        model.quantize_model(quant_config=quant_config)
+
+        # Save model
+        save_dir = "snapshots/" + model_id
+        Path(save_dir).mkdir(parents=True, exist_ok=True)
+        model.save_quantized(save_dir=save_dir)
+
+
+def load_quantized(model_id):
+    # Load model
+    save_dir = "snapshots/" + model_id
+    model = HQQOpenCLIP.from_quantized(save_dir)
+    return model
+
+
+def compare_weights_raw(model_id):
+    qnt_pt_file = "snapshots/" + model_id + "/qmodel.pt"
+    ref_pt_file = "/home/justin/.cache/huggingface/hub/models--laion--CLIP-ViT-H-14-laion2B-s32B-b79K/snapshots/de081ac0a0ca8dc9d1533eed1ae884bb8ae1404b/open_clip_pytorch_model.bin"
+    qnt = torch.load(qnt_pt_file)
+    ref = torch.load(ref_pt_file)
+    print("*" * 72)
+    print("from quantized file")
+    for k in qnt:
+        print(k)
+    print("*" * 72)
+    print("from original file")
+    for k in ref:
+        print(k)
+
+
+def compare_weights(model_id):
+    model, model_ref, _ = create_model_dual(model_id)
+    dict1 = dict(model.named_parameters())
+    dict2 = dict(model_ref.named_parameters())
+    for name, param in dict1.items():
+        param_ref = dict2.get(name, None)
+        if param_ref is not None:
+            if not torch.allclose(param_ref.data, param.data, rtol=0.0, equal_nan=True):
+                print(f"weight differs: {name}")
+                if name == "text_projection":
+                    print(param)
+                    print(param_ref)
+            else:
+                print(f"weight same: {name}")
+        else:
+            print(f"not in ref model: {name}")
+
+
+def create_model_dual(model_id, mask=3):
+    model = None
+    model_ref = None
+    preprocess = None
+
+    if mask & 1:
+        model = load_quantized(model_id)
+        model = model.half().cuda()
+
+    # Load reference model to compare with
+    if mask & 2:
+        comps = model_id.split("/")
+        elems = comps[1].split("-")
+        model_name = "-".join(elems[1:4])
+        pretrained = "-".join(elems[4:])
+        model_ref, _, preprocess = open_clip.create_model_and_transforms(
+            model_name, pretrained=pretrained
+        )
+        model_ref = model_ref.half().cuda()
+    return model, model_ref, preprocess
+
+
+def create_and_quant_model(model_id):
+    model = HQQOpenCLIP.create_model(model_id, device="cpu")
+    # Quantize settings
+    # quant_config = BaseQuantizeConfig(nbits=8, group_size=128)
+    quant_config = BaseQuantizeConfig(nbits=4, group_size=64)
+    # quant_config = BaseQuantizeConfig(nbits=3, group_size=64)
+    # quant_config = BaseQuantizeConfig(nbits=2, group_size=16, quant_scale=True)
+    # Quantize
+    model.quantize_model(quant_config=quant_config)
+    model = model.half().cuda()
+    return model
+
+
+def compare(model_id):
+    model = create_and_quant_model(model_id)
+    _, model_ref, preprocess = create_model_dual(model_id, mask=2)
+    model.eval()
+    model_ref.eval()
+
+    descriptions = {
+        "page": "a page of text about segmentation",
+        "chelsea": "a facial photo of a tabby cat",
+        "astronaut": "a portrait of an astronaut with the American flag",
+        "rocket": "a rocket standing on a launchpad",
+        "motorcycle_right": "a red motorcycle standing in a garage",
+        "camera": "a person looking at a camera on a tripod",
+        "horse": "a black-and-white silhouette of a horse",
+        "coffee": "a cup of coffee on a saucer",
+    }
+    texts = descriptions.values()
+    text_processed = tokenizer.tokenize(texts).cuda()
+
+    # preprocess image and text
+    img = Image.open(os.path.join(data_dir, "astronaut.png")).convert("RGB")
+    img_preprocessed = preprocess(img).cuda().unsqueeze(0)
+
+    with torch.amp.autocast("cuda"):
+        img_embedding, text_embedding, _ = model_ref(img_preprocessed, text_processed)
+    probs = (100 * img_embedding @ text_embedding.T).softmax(dim=-1)
+    print(probs)
+
+    with torch.amp.autocast("cuda"):
+        img_embedding, text_embedding, _ = model(img_preprocessed, text_processed)
+    probs = (100 * img_embedding @ text_embedding.T).softmax(dim=-1)
+    print(probs)
+
+
+def main():
+    # quant_models(model_ids)
+    # load_quantized(model_ids[0])
+    compare(model_ids[2])
+    # compare_weights(model_ids[2])
+    # compare_weights_raw(model_ids[2])
+
+
+if __name__ == "__main__":
+    main()

lm-quant-toolkit/.deps/hqq/examples/timm/vit_clip_example.py

@@ -0,0 +1,73 @@
+import numpy as np
+import timm
+import torch
+
+from hqq.core.quantize import BaseQuantizeConfig
+from hqq.engine.timm import HQQtimm
+
+# Model
+# model_id = 'vit_base_patch32_clip_224.laion2b' #ViT-B-32
+# model_id = 'vit_large_patch14_clip_224.laion2b'#ViT-L-14
+model_id = "vit_huge_patch14_clip_224.laion2b"  # ViT-H-14
+
+# Load model (on CPU)
+model = HQQtimm.create_model(model_id, pretrained=True)
+
+# Quantize settings
+# quant_config = BaseQuantizeConfig(nbits=8, group_size=128)
+quant_config = BaseQuantizeConfig(nbits=4, group_size=64)
+# quant_config = BaseQuantizeConfig(nbits=3, group_size=64)
+# quant_config = BaseQuantizeConfig(nbits=2, group_size=16, quant_scale=True)
+
+# Quantize
+model.quantize_model(quant_config=quant_config)
+
+###############################################################
+# #Save model
+# save_dir = "repo/" + model_id
+# model.save_quantized(model, save_dir=save_dir)
+
+# #Load model
+# model = HQQtimm.from_quantized(save_dir)
+###############################################################
+
+# Load reference model to compare with
+model_ref = timm.create_model(model_id, pretrained=True)
+model_ref = model_ref.half().cuda()
+model_ref.eval()
+
+# Pre-processing
+mean_clip = np.array([0.4815, 0.4578, 0.4082], "float32")
+std_clip = np.array([0.2686, 0.2613, 0.2758], "float32")
+
+
+def normalize_images_clip(data_np_in, BCHW=True):
+    data_t = (
+        torch.from_numpy(data_np_in).float()
+        if (type(data_np_in) is np.ndarray)
+        else data_np_in.float()
+    )
+    data_t = (data_t / 255.0 - mean_clip) / std_clip
+    data_t = data_t.swapaxes(2, 3).swapaxes(1, 2) if (BCHW) else data_t
+    return data_t
+
+
+###############################################################
+# Compare the compressed model with the original
+x = np.random.rand(16, 224, 224, 3)
+x = normalize_images_clip(x).half().cuda()
+
+# Quantize
+with torch.no_grad():
+    y_q = model(x)
+    y_q /= torch.norm(y_q, p=2, dim=-1, keepdim=True)
+
+# Full-precision
+with torch.no_grad():
+    y_r = model_ref(x)
+    y_r /= torch.norm(y_r, p=2, dim=-1, keepdim=True)
+
+# We want the dot product to be as close as possible to 1
+print(
+    "Average dot-product score", float(torch.diag(torch.matmul(y_q, y_r.t())).mean())
+)  # ~0.998 (ViT-H-14 @4bit)

lm-quant-toolkit/.deps/hqq/examples/vllm/llama2_example.py

@@ -0,0 +1,25 @@
+model_id = 'meta-llama/Llama-2-7b-chat-hf'
+
+#Load VLLM un-quantized model 
+from hqq.engine.vllm import HQQLLM
+model = HQQLLM(model=model_id)
+
+#Quantize the model
+from hqq.core.quantize import *
+model.quantize_model(BaseQuantizeConfig(nbits=4, group_size=64))
+
+#Optional: Save the model
+#model.save_quantized(model_id.split('/')[-1] + '_quantized')
+
+#Optional: Set backend
+HQQLinear.set_backend(HQQBackend.PYTORCH_COMPILE) #set backend
+
+#Generation
+from vllm.entrypoints.llm import SamplingParams
+sampling_params = SamplingParams(temperature=0.6, top_p=0.90, max_tokens=1000, repetition_penalty=1.2)
+
+prompt = "How can I build a car?"
+
+output = model.generate([prompt], sampling_params)[0]
+print(output.prompt)
+print(output.outputs[0].text)

lm-quant-toolkit/.deps/hqq/hqq.egg-info/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: hqq
+Version: 0.1.7.post2
+Summary: Half-Quadratic Quantization (HQQ)
+Home-page: https://github.com/mobiusml/hqq/
+Author: Dr. Hicham Badri
+Author-email: hicham@mobiuslabs.com
+License: Apache 2
+License-File: LICENSE
+Requires-Dist: numpy>=1.24.4
+Requires-Dist: tqdm>=4.64.1
+Requires-Dist: einops
+Requires-Dist: accelerate
+Requires-Dist: transformers>=4.36.1
+Requires-Dist: huggingface_hub
+Requires-Dist: termcolor
+Dynamic: author
+Dynamic: author-email
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary

lm-quant-toolkit/.deps/hqq/hqq.egg-info/SOURCES.txt

@@ -0,0 +1,53 @@
+LICENSE
+setup.py
+hqq/__init__.py
+hqq.egg-info/PKG-INFO
+hqq.egg-info/SOURCES.txt
+hqq.egg-info/dependency_links.txt
+hqq.egg-info/requires.txt
+hqq.egg-info/top_level.txt
+hqq/backends/__init__.py
+hqq/backends/marlin.py
+hqq/backends/torchao.py
+hqq/core/__init__.py
+hqq/core/bitpack.py
+hqq/core/optimize.py
+hqq/core/peft.py
+hqq/core/quantize.py
+hqq/core/utils.py
+hqq/engine/__init__.py
+hqq/engine/base.py
+hqq/engine/hf.py
+hqq/engine/open_clip.py
+hqq/engine/timm.py
+hqq/engine/vllm.py
+hqq/kernels/__init__.py
+hqq/kernels/hqq_aten_cuda.cpp
+hqq/kernels/hqq_aten_cuda_kernel.cu
+hqq/kernels/hqq_aten_torch.cpp
+hqq/kernels/setup_cuda.py
+hqq/kernels/setup_torch.py
+hqq/models/__init__.py
+hqq/models/base.py
+hqq/models/hf/__init__.py
+hqq/models/hf/base.py
+hqq/models/hf/llama.py
+hqq/models/hf/mistral.py
+hqq/models/hf/mixtral.py
+hqq/models/hf/phi.py
+hqq/models/hf/phi_opt.py
+hqq/models/open_clip/__init__.py
+hqq/models/open_clip/base.py
+hqq/models/open_clip/vit_clip.py
+hqq/models/timm/__init__.py
+hqq/models/timm/base.py
+hqq/models/timm/vit_clip.py
+hqq/models/vllm/__init__.py
+hqq/models/vllm/base.py
+hqq/models/vllm/llama.py
+hqq/utils/__init__.py
+hqq/utils/generation_hf.py
+hqq/utils/optimizer.py
+hqq/utils/patching.py
+tests/test_bitpack.py
+tests/test_quantize.py

lm-quant-toolkit/.deps/hqq/hqq.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

lm-quant-toolkit/.deps/hqq/hqq.egg-info/requires.txt

@@ -0,0 +1,7 @@
+numpy>=1.24.4
+tqdm>=4.64.1
+einops
+accelerate
+transformers>=4.36.1
+huggingface_hub
+termcolor

lm-quant-toolkit/.deps/hqq/hqq.egg-info/top_level.txt

@@ -0,0 +1 @@
+hqq

lm-quant-toolkit/.deps/hqq/hqq/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.7.post2"
+__author__  = 'Dr. Hicham Badri'
+__credits__ = 'Mobius Labs GmbH'

lm-quant-toolkit/.deps/hqq/hqq/backends/marlin.py

@@ -0,0 +1,118 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2024
+#####################################################
+import torch
+
+try:
+    import marlin
+except Exception:
+    marlin = None
+from ..core.quantize import Quantizer
+
+
+class MarlinLinear(torch.nn.Module):
+    def __init__(
+        self, W: torch.Tensor, scales: torch.Tensor, u=None, bias=None, groupsize=-1
+    ):
+        super().__init__()
+
+        m, n = W.shape
+        device = W.device
+        _linear = torch.nn.Linear(m, n)
+        _linear.weight.data = W.half().t()
+
+        effective_groupsize = m if (groupsize == -1) else groupsize
+
+        _layer = marlin.Layer(m, n, groupsize=groupsize)
+        _layer.k = m
+        _layer.n = n
+        _layer.groupsize = effective_groupsize
+        _layer.B = torch.empty((m // 16, n * 16 // 8), dtype=torch.int, device=device)
+        _layer.s = torch.empty(
+            (m // effective_groupsize, n), dtype=torch.half, device=device
+        )
+        _layer.pack(_linear, scales.t())
+
+        self.bias = bias.half() if (bias is not None) else None
+        self.Wq_packed = _layer.B.clone()
+        self.scales = _layer.s.clone()
+        self.workspace_fp = torch.zeros(n // 128 * 16, device=device)
+        self.in_features = m
+        self.out_features = n
+        self.group_size = effective_groupsize
+        self.axis = 1
+        self.device = device
+        self.compute_dtype = torch.float16
+        self.u = torch.nn.Parameter(u, requires_grad=False) if (u is not None) else None
+
+        del _linear, _layer
+        torch.cuda.empty_cache()
+
+    @torch.no_grad()
+    def matmul(self, x):
+        out = torch.empty(
+            x.shape[:-1] + (self.scales.shape[1],), dtype=x.dtype, device=x.device
+        )
+        marlin.mul(
+            x.view((-1, x.shape[-1])),
+            self.Wq_packed,
+            out.view((-1, out.shape[-1])),
+            self.scales,
+            self.workspace_fp,
+        )
+        return out
+
+    @torch.jit.ignore
+    def forward(self, x):
+        out = self.matmul(x)
+
+        if self.u is not None:
+            out += torch.matmul(x.sum(axis=-1, keepdim=True), self.u)
+
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+
+# ONLY WORKS WITH AXIS=1, group_size= - 1
+def patch_hqq_to_marlin(layer, patch_params):
+    if marlin is None:
+        return layer
+
+    z_shift = 8.0
+    hqq_layer = layer.linear_layer if hasattr(layer, "linear_layer") else layer
+
+    # Check config suppport
+    if (
+        (hqq_layer.meta["axis"] == 0)
+        or (hqq_layer.meta["group_size"] is not None)
+        or (hqq_layer.meta["nbits"] != 4)
+    ):
+        print("Skipping marlin conversion for", hqq_layer.name)
+        return layer
+
+    W_r = Quantizer.unpack[hqq_layer.meta["packing"]](
+        hqq_layer.W_q, dtype=hqq_layer.compute_dtype
+    ).t()
+    z = hqq_layer.meta["zero"]
+    s = hqq_layer.meta["scale"].t()
+    W_r = (W_r - z_shift) * s
+
+    if type(z) in [torch.Tensor, torch.nn.Parameter]:
+        z = z.t()
+        u = (s * (-z + z_shift)).view([1, -1])
+    else:
+        u = None
+
+    marlin_layer = MarlinLinear(W_r, s, u=u, bias=hqq_layer.bias)
+
+    if hasattr(layer, "linear_layer"):
+        del layer.linear_layer
+        layer.linear_layer = marlin_layer
+    else:
+        del hqq_layer
+        layer = marlin_layer
+
+    torch.cuda.empty_cache()
+
+    return layer

lm-quant-toolkit/.deps/hqq/hqq/backends/torchao.py

@@ -0,0 +1,388 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2024
+#####################################################
+
+# Makes HQQ 4-bit (axis=1) compatbile with torch.ops.aten._weight_int4pack_mm
+
+# The code is partially based on: https://github.com/pytorch-labs/gpt-fast/blob/main/quantize.py (BSD-3-Clause license)
+# Only the packing logic is copied, the rest is rewritten to support HQQ's logic
+
+# Only works with: bfloat16, axis=1.
+# Only tested on Ada gpus.
+
+import torch
+import copy
+from torch import bfloat16, nn, Tensor
+import torch.nn.functional as F
+from typing import Union
+
+from ..core.quantize import HQQLinear, Quantizer
+from ..core.peft import HQQLinearLoRA
+from ..core.utils import cleanup
+
+
+class HQQLinearTorchWeightOnlynt4(torch.nn.Module):
+    def __init__(
+        self,
+        linear_layer: Union[nn.Module, None],
+        quant_config: dict,
+        del_orig: bool = True,
+        compute_dtype: torch.dtype = bfloat16,
+        device: str = "cuda",
+        initialize: bool = True,
+        inner_k_tiles=8,
+        padding=True,
+    ):
+        super().__init__()
+
+        self.ready = False
+        self.in_gpu = False
+        self.bias = None
+        self.device = device
+        self.compute_dtype = compute_dtype
+        self.quant_config = (
+            copy.deepcopy(quant_config) if (quant_config is not None) else None
+        )
+        self.del_orig = del_orig
+
+        if (quant_config is None) and (linear_layer is None):
+            raise Exception(
+                "Invalid parameters: Both quant_config and linear_layer are None."
+            )
+
+        if (quant_config is None) and (initialize is True):
+            raise Exception(
+                "Invalid parameters: setting initialize to True requires a quant_config."
+            )
+
+        if self.quant_config is not None:
+            weight_quant_params = self.quant_config["weight_quant_params"]
+            self.groupsize = weight_quant_params["group_size"]
+            self.nbits = weight_quant_params["nbits"]
+            self.axis = weight_quant_params["axis"]
+
+        if linear_layer is not None:
+            self.groupsize = linear_layer.meta["group_size"]
+            self.nbits = linear_layer.meta["nbits"]
+            self.axis = linear_layer.meta["axis"]
+
+        self.inner_k_tiles = inner_k_tiles
+        self.padding = padding
+
+        assert self.axis==1, "Only axis==1 is supported"
+        assert self.nbits in [4], "Unsupported nbits."
+        assert (
+            self.compute_dtype is bfloat16
+        ), "Only bfloat16 compute_dtype is supported."
+        assert self.groupsize in [None, 32, 64, 128, 256], "Unsupported groupsize."
+        assert self.inner_k_tiles in [2, 4, 8], "Unsupported tile."
+
+        self.linear_layer = linear_layer
+
+        if initialize:
+            self.initialize()
+
+    ###################### Initializers ######################
+    def initialize_with_hqq_quants(self, W_q, meta, bias=None):
+        self.padding = (
+            False  # Force padding off, a bit tricky to post-pad with grouping
+        )
+
+        self.set_shape(meta["shape"])
+        self.process_hqq_quants(W_q, meta)
+        self.bias = bias
+        self.ready = True
+        self.in_gpu = True
+        torch.cuda.empty_cache()
+
+        return self
+
+    def initialize(self):
+        if self.linear_layer is not None:
+            W = self.linear_layer.weight.data
+            self.set_shape(W.shape)
+
+            if self.in_features_diff > 0:
+                W = F.pad(W, pad=(0, self.in_features_diff), value=0)
+
+            W_q, meta = self.quantize(W, **self.quant_config)
+            self.process_hqq_quants(W_q, meta)
+            del W_q, meta
+
+            self.bias = (
+                None
+                if (self.linear_layer.bias is None)
+                else self.linear_layer.bias.to(
+                    dtype=self.compute_dtype, device=self.device
+                )
+            )
+
+        if self.del_orig:
+            del self.linear_layer
+
+        self.ready = True
+        self.in_gpu = True
+        torch.cuda.empty_cache()
+
+        return self
+
+    ###################### Quantize/packing ######################
+
+    def quantize(
+        self,
+        W: Tensor,
+        weight_quant_params: dict,
+        scale_quant_params=Union[dict,None],
+        zero_quant_params=Union[dict,None],
+        offload_meta=False,
+    ):
+        W_q, meta = Quantizer.quantize(
+            W,
+            **weight_quant_params,
+            device=self.device,
+            compute_dtype=self.compute_dtype,
+            bitpack=False,
+        )
+
+        # ToDO: meta quantization
+
+        return W_q, meta
+
+    # TODO: move these to utils
+    @torch.no_grad()
+    def reshape_meta_axis1(self, meta_tensor, new_group_size, shape):
+        meta_tensor = meta_tensor.repeat([1, shape[1]]).reshape(shape)
+        meta_tensor = torch.mean(
+            meta_tensor.reshape([-1, new_group_size]), axis=1, keepdim=True
+        )
+        return meta_tensor
+
+    def find_multiple(self, n: int, k: int) -> int:
+        if n % k == 0:
+            return n
+        return n + k - (n % k)
+
+    def set_shape(self, shape):
+        self.shape = shape
+        self.in_features = shape[1]
+        self.out_features = shape[0]
+
+        self.origin_in_features = self.in_features
+        if self.padding:
+            self.in_features = self.find_multiple(self.in_features, 1024)
+
+        self.in_features_diff = self.in_features - self.origin_in_features
+
+    @torch.no_grad()
+    def process_hqq_quants(self, W_q, meta):
+        scales = meta["scale"]
+        zeros = meta["zero"]
+        shape = meta["shape"]
+
+        if meta["packing"] is not None:
+            W_q = Quantizer.unpack[meta["packing"]](W_q)
+
+        if self.groupsize is None:
+            self.groupsize = 128
+            W_q = W_q.reshape([-1, self.groupsize])
+            scales = self.reshape_meta_axis1(scales, self.groupsize, shape)
+            zeros = self.reshape_meta_axis1(zeros, self.groupsize, shape)
+
+        W_q_torch, scales_torch, zeros_torch = self.hqq_quants_to_torch_quants(
+            W_q=W_q, scales=scales, zeros=zeros, shape=shape, nbits=self.nbits
+        )
+        self.weight_int4pack = torch.ops.aten._convert_weight_to_int4pack(
+            W_q_torch, self.inner_k_tiles
+        )
+        self.scales_and_zeros = self.pack_scales_and_zeros(scales_torch, zeros_torch)
+
+        del W_q_torch, scales_torch, zeros_torch
+        torch.cuda.empty_cache()
+
+    @torch.no_grad()
+    def hqq_quants_to_torch_quants(
+        self, W_q: Tensor, scales: Tensor, zeros: Tensor, shape, nbits=4
+    ):
+        W_q = W_q.to(dtype=self.compute_dtype, device=self.device)
+        scales = scales.to(dtype=self.compute_dtype, device=self.device)
+        zeros = zeros.to(dtype=self.compute_dtype, device=self.device)
+
+        max_int = 2**nbits - 1
+        min_int = 0
+        dump = 2 ** (nbits - 1)
+
+        # HQQ -> torch logic
+        new_zeros = (scales * dump) - zeros * scales
+
+        min_val = new_zeros - scales * dump
+
+        # group_quantize_tensor_from_qparams
+        W_r = (W_q - zeros) * scales
+
+        W_q = (
+            W_r.sub(min_val)
+            .div(scales)
+            .round()
+            .clamp_(min_int, max_int)
+            .to(torch.int32)
+            .reshape(shape)
+            .contiguous()
+        )
+
+        # group_dequantize_tensor_from_qparams
+        # W_r = W_q*scales + min_val
+
+        scales = scales.contiguous().reshape(shape[0], -1)
+        new_zeros = new_zeros.contiguous().reshape(shape[0], -1)
+
+        return W_q, scales, new_zeros
+
+    def pack_scales_and_zeros(self, scales, zeros):
+        assert scales.shape == zeros.shape
+        assert scales.dtype == bfloat16
+        assert zeros.dtype == bfloat16
+        return (
+            torch.cat(
+                [
+                    scales.reshape(scales.size(0), scales.size(1), 1),
+                    zeros.reshape(zeros.size(0), zeros.size(1), 1),
+                ],
+                2,
+            )
+            .transpose(0, 1)
+            .contiguous()
+        )
+
+    ###################### Forward/matmul ######################
+
+    # @torch.jit.ignore()
+    def matmul(self, x: Tensor) -> Tensor:
+        origin_x_size = x.size()
+        x = x.reshape(-1, origin_x_size[-1])
+        c = torch.ops.aten._weight_int4pack_mm(
+            x, self.weight_int4pack, self.groupsize, self.scales_and_zeros
+        )
+        new_shape = origin_x_size[:-1] + (self.out_features,)
+        c = c.reshape(new_shape)
+        return c
+
+    # TODO without matmul
+    def dequantize(self) -> Tensor:
+        return self.matmul(
+            torch.eye(self.in_features, dtype=self.compute_dtype, device=self.device)
+        )[: self.origin_in_features].t()
+
+    # TODO: backward
+    def forward(self, x: Tensor) -> Tensor:
+        if self.in_features_diff > 0:
+            x = F.pad(x, pad=(0, self.in_features_diff))
+
+        out = self.matmul(x)
+
+        if self.bias is not None:
+            out += self.bias
+        return out
+
+
+###################### Patching ######################
+def patch_linearlayers(model, fct, patch_param=None):
+    model.base_class.patch_linearlayers(
+        model,
+        fct,
+        {lin_tag: patch_param for lin_tag in model.base_class.get_linear_tags()},
+    )
+
+
+def patch_hqq_to_aoint4(layer, patch_params):
+    hqq_layer = None
+    if type(layer) is HQQLinear:
+        hqq_layer = layer
+    if type(layer) is HQQLinearLoRA:
+        hqq_layer = layer.linear_layer
+
+    if hqq_layer is None:
+        return layer
+
+    if hqq_layer.meta["nbits"] != 4 or hqq_layer.meta["axis"] != 1:
+        print("Skipping aoint4 conversion for ", hqq_layer.name)
+        return layer
+
+    quant_config = getattr(hqq_layer, "quant_config", None)
+
+    hqq_aoint4_layer = HQQLinearTorchWeightOnlynt4(
+        None,
+        quant_config=quant_config,
+        compute_dtype=hqq_layer.compute_dtype,
+        device=hqq_layer.device,
+        del_orig=False,
+        initialize=False,
+        padding=False,
+    )
+    hqq_aoint4_layer.initialize_with_hqq_quants(
+        hqq_layer.W_q, hqq_layer.meta, hqq_layer.bias
+    )
+
+    del hqq_layer
+    torch.cuda.empty_cache()
+
+    if type(layer) is HQQLinear:
+        return hqq_aoint4_layer
+
+    if type(layer) is HQQLinearLoRA:
+        layer.linear_layer = hqq_aoint4_layer
+
+    return layer
+
+
+def replace_with_torchInt4(model):
+    patch_linearlayers(model, patch_hqq_to_aoint4)
+    cleanup()
+
+
+# Force requantize, mainly to check if the padding with int4mm is faster
+def patch_hqq_to_aoint4_force_requantize(layer, patch_params):
+    hqq_layer = None
+    if type(layer) is HQQLinear:
+        hqq_layer = layer
+    if type(layer) is HQQLinearLoRA:
+        hqq_layer = layer.linear_layer
+
+    if hqq_layer is None:
+        return layer
+
+    if hqq_layer.meta["nbits"] != 4 or hqq_layer.meta["axis"] != 1:
+        print("Skipping aoint4 conversion for ", hqq_layer.name)
+        return layer
+
+    # Create dummy linear layer to store dequantize weights
+    dummy_linear = torch.nn.Linear(1, 1, bias=False)
+    dummy_linear.weight.data = hqq_layer.dequantize()
+
+    # Disable optimizer on already dequantized weights
+    quant_config = hqq_layer.quant_config
+    quant_config["weight_quant_params"]["optimize"] = False
+
+    hqq_aoint4_layer = HQQLinearTorchWeightOnlynt4(
+        dummy_linear,
+        quant_config=quant_config,
+        compute_dtype=hqq_layer.compute_dtype,
+        device=hqq_layer.device,
+        del_orig=True,
+        initialize=True,
+        padding=True,
+    )
+
+    del hqq_layer
+    torch.cuda.empty_cache()
+
+    if type(layer) is HQQLinear:
+        return hqq_aoint4_layer
+
+    if type(layer) is HQQLinearLoRA:
+        layer.linear_layer = hqq_aoint4_layer
+
+    return layer
+
+
+def replace_with_torchInt4_force_requantize(model):
+    patch_linearlayers(model, patch_hqq_to_aoint4_force_requantize)
+    cleanup()

lm-quant-toolkit/.deps/hqq/hqq/core/bitpack.py

@@ -0,0 +1,144 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+
+import torch
+from torch import uint8, int32, Tensor
+import numpy as np
+
+
+# Bit packing logic. format: pack/unpack_nBits_target-<uint8 or int32>
+class BitPack:
+    # 8-bit
+    ################################################
+    @staticmethod
+    def pack_8bit_u8(W_q: Tensor) -> Tensor:
+        return W_q.to(uint8)
+
+    @staticmethod
+    def unpack_8bit_u8(W_q: Tensor, dtype=uint8) -> Tensor:
+        return W_q.to(dtype)
+
+    # 4-bit
+    ################################################
+    @staticmethod
+    def pack_4bit_u8(W_q: Tensor) -> Tensor:  # uint8 > uint8/2
+        W_q = W_q.to(uint8)
+        _step = int(len(W_q) / 2)
+
+        return (W_q[:_step] << 4) | W_q[_step:]
+
+    @staticmethod
+    def unpack_4bit_u8(W_q: Tensor, dtype=uint8) -> Tensor:  # uint8/2 > uint8
+        _step = W_q.shape[0]
+        tmp = torch.empty([2 * _step, W_q.shape[1]], dtype=dtype, device=W_q.device)
+
+        tmp[:_step] = (W_q & 0b11110000) >> 4
+        tmp[_step:] = W_q & 0b00001111
+
+        return tmp
+
+    # 2-bit
+    ################################################
+    @staticmethod
+    def pack_2bit_u8(W_q: Tensor) -> Tensor:  # uint8 > uint8/4
+        W_q = W_q.to(uint8)
+        _step = int(len(W_q) / 4)
+
+        return (
+            W_q[:_step] << 6
+            | W_q[_step : 2 * _step] << 4
+            | W_q[2 * _step : 3 * _step] << 2
+            | W_q[3 * _step :]
+        )
+
+    @staticmethod
+    def unpack_2bit_u8(W_q: Tensor, dtype=uint8) -> Tensor:
+        _step = W_q.shape[0]
+        tmp = torch.empty([4 * _step, W_q.shape[1]], dtype=dtype, device=W_q.device)
+
+        tmp[0 * _step : 1 * _step] = (W_q & 0b11000000) >> 6
+        tmp[1 * _step : 2 * _step] = (W_q & 0b00110000) >> 4
+        tmp[2 * _step : 3 * _step] = (W_q & 0b00001100) >> 2
+        tmp[3 * _step : 4 * _step] = W_q & 0b00000011
+
+        return tmp
+
+    # 3-bit
+    ################################################
+    @staticmethod
+    def pack_3bit_32(W_q_in: Tensor) -> Tensor:
+        W_q = torch.zeros(
+            [int(10 * np.ceil(W_q_in.shape[0] / 10.0)), W_q_in.shape[1]],
+            device=W_q_in.device,
+            dtype=int32,
+        )
+        W_q[: len(W_q_in)] = W_q_in
+        _step = int(len(W_q) / 10)
+
+        W_q = (
+            (W_q[:_step] << 27)
+            | (W_q[1 * _step : 2 * _step] << 24)
+            | (W_q[2 * _step : 3 * _step] << 21)
+            | (W_q[3 * _step : 4 * _step] << 18)
+            | (W_q[4 * _step : 5 * _step] << 15)
+            | (W_q[5 * _step : 6 * _step] << 12)
+            | (W_q[6 * _step : 7 * _step] << 9)
+            | (W_q[7 * _step : 8 * _step] << 6)
+            | (W_q[8 * _step : 9 * _step] << 3)
+            | (W_q[9 * _step : 10 * _step])
+        )
+
+        return W_q
+
+    # A bit faster than _cat version
+    @staticmethod
+    def unpack_3bit_32(W_q: Tensor, dtype=uint8) -> Tensor:
+        _step = W_q.shape[0]
+        tmp = torch.empty([10 * _step, W_q.shape[1]], dtype=dtype, device=W_q.device)
+
+        tmp[0 * _step : 1 * _step] = (W_q & 0b00111000000000000000000000000000) >> 27
+        tmp[1 * _step : 2 * _step] = (W_q & 0b00000111000000000000000000000000) >> 24
+        tmp[2 * _step : 3 * _step] = (W_q & 0b00000000111000000000000000000000) >> 21
+        tmp[3 * _step : 4 * _step] = (W_q & 0b00000000000111000000000000000000) >> 18
+        tmp[4 * _step : 5 * _step] = (W_q & 0b00000000000000111000000000000000) >> 15
+        tmp[5 * _step : 6 * _step] = (W_q & 0b00000000000000000111000000000000) >> 12
+        tmp[6 * _step : 7 * _step] = (W_q & 0b00000000000000000000111000000000) >> 9
+        tmp[7 * _step : 8 * _step] = (W_q & 0b00000000000000000000000111000000) >> 6
+        tmp[8 * _step : 9 * _step] = (W_q & 0b00000000000000000000000000111000) >> 3
+        tmp[9 * _step : 10 * _step] = W_q & 0b00000000000000000000000000000111
+
+        return tmp
+
+    # 1-bit
+    ################################################
+    @staticmethod
+    def pack_1bit_u8(W_q: Tensor) -> Tensor:
+        W_q = W_q.to(uint8)
+        _step = int(len(W_q) / 8)
+
+        return (
+            W_q[:_step] << 7
+            | W_q[1 * _step : 2 * _step] << 6
+            | W_q[2 * _step : 3 * _step] << 5
+            | W_q[3 * _step : 4 * _step] << 4
+            | W_q[4 * _step : 5 * _step] << 3
+            | W_q[5 * _step : 6 * _step] << 2
+            | W_q[6 * _step : 7 * _step] << 1
+            | W_q[7 * _step : 8 * _step]
+        )
+
+    @staticmethod
+    def unpack_1bit_u8(W_q: Tensor, dtype=uint8) -> Tensor:
+        _step = W_q.shape[0]
+        tmp = torch.empty([8 * _step, W_q.shape[1]], dtype=dtype, device=W_q.device)
+
+        tmp[0 * _step : 1 * _step] = (W_q & 0b10000000) >> 7
+        tmp[1 * _step : 2 * _step] = (W_q & 0b01000000) >> 6
+        tmp[2 * _step : 3 * _step] = (W_q & 0b00100000) >> 5
+        tmp[3 * _step : 4 * _step] = (W_q & 0b00010000) >> 4
+        tmp[4 * _step : 5 * _step] = (W_q & 0b00001000) >> 3
+        tmp[5 * _step : 6 * _step] = (W_q & 0b00000100) >> 2
+        tmp[6 * _step : 7 * _step] = (W_q & 0b00000010) >> 1
+        tmp[7 * _step : 8 * _step] = W_q & 0b00000001
+
+        return tmp

lm-quant-toolkit/.deps/hqq/hqq/core/optimize.py

@@ -0,0 +1,476 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+import torch
+import numpy as np
+from torch import float32, float16, Tensor
+from functools import partial
+from typing import Union
+
+
+# re-estimate the scale based on the inverse median: Only tested with axis==0
+def update_scale_inverse_median(
+    W_f: Tensor, scale: Tensor, zero: Tensor, axis: int, min_max: list
+) -> tuple:
+    scale_rng = 2e4
+    z_val = 1e-4
+    delta = 1e-2
+
+    W_q = torch.round(W_f * scale + zero).clamp(min_max[0], min_max[1])
+
+    # Correct zero to avoid W_q==zero
+    zero_c = zero.clone()
+    zero_c_indx = torch.sum(1.0 * ((W_q - zero) == 0), axis=axis, keepdim=True) > 0
+    zero_c[zero_c_indx] = zero_c[zero_c_indx] + delta
+
+    # Build scale tensor
+    W_f_c = W_f.clone()
+    W_f_c_mask = torch.abs(W_f_c) < z_val
+    W_f_c[W_f_c_mask] = z_val
+
+    scale_tensor = (W_q - zero_c).float() / W_f_c.float()
+    # W_r               = (W_q - zero_c)/scale_tensor
+
+    # Normalize scale_tensor
+    scale_b = torch.median(scale_tensor, axis=axis, keepdim=True)[0]
+    scale_b = scale_b.clamp(min=-scale_rng, max=scale_rng).half()
+
+    # Mix with older scale
+    W_r = (W_q - zero_c) / scale_b
+    err_b = torch.abs(W_f - W_r).mean(axis=axis, keepdim=True)
+
+    W_r = (W_q - zero_c) / scale
+    err_a = torch.abs(W_f - W_r).mean(axis=axis, keepdim=True)
+
+    mask = (err_b < err_a).half()
+    scale_b = mask * scale_b + (1 - mask) * scale
+
+    # W_r   = (W_q - zero_c)/scale_b
+    return scale_b, zero_c
+
+
+# Greedy local search: Only tested with axis==0
+def update_scale_grid_search(
+    W_f: Tensor, scale: Tensor, zero: Tensor, axis: int, min_max: list, N: int = 128 + 1
+) -> Tensor:
+    # Make sure it's an odd number so that the original scale is included
+    assert N % 2 == 1, "Please check whether N: odd number"
+    rng_dump = 0.05  # 0.05 / 1.
+    z_val = 2e-4
+
+    device = scale.device
+    dtype = scale.dtype
+    ###############################
+    W_q = torch.round(W_f * scale + zero).clamp(min_max[0], min_max[1])
+    n_clusters = max(W_q.shape[0], W_q.shape[1])
+    rng = torch.abs(scale).mean() * rng_dump if (rng_dump < 1.0) else rng_dump
+
+    scale_shifted = (
+        torch.linspace(-rng, rng, N)[:, None]
+        .to(dtype=dtype, device=device)
+        .repeat(1, n_clusters)
+        + scale
+    )
+
+    # Safe inverse
+    scale_shifted[
+        torch.logical_and(scale_shifted >= 0, torch.abs(scale_shifted) <= z_val)
+    ] = z_val
+    scale_shifted[
+        torch.logical_and(scale_shifted < 0, torch.abs(scale_shifted) <= z_val)
+    ] = -z_val
+
+    err = torch.empty([N, n_clusters], dtype=dtype, device=device)
+    for i in range(N):
+        W_r = (W_q - zero) / scale_shifted[i][None, :]
+        err[i] = torch.abs(W_f - W_r).mean(axis=axis, keepdim=True)
+
+    ind_r = torch.argmin(err, axis=axis).to(torch.int32)
+    ind_c = torch.arange(len(ind_r), dtype=torch.int32, device=device)
+    scale_b = scale_shifted[ind_r, ind_c]
+
+    return scale_b
+
+
+# Shrinking operator
+def shrink_lp_op(x: Tensor, beta: float, lp_norm: float) -> Tensor:
+    if lp_norm == 1:
+        return torch.sign(x) * torch.nn.functional.relu(torch.abs(x) - 1.0 / beta)
+    else:
+        return torch.sign(x) * torch.nn.functional.relu(
+            torch.abs(x) - (1.0 / beta) * torch.pow(torch.abs(x), lp_norm - 1)
+        )
+
+
+# Proximal solver || W - dequantize(quantize(W))||_p^p - Experimental
+@torch.inference_mode()
+def optimize_weights_proximal_v2(
+    tensor: Tensor,
+    scale: Tensor,
+    zero: Tensor,
+    min_max: list,
+    axis: int = 0,
+    device: Union[str, None] = None,
+    dtype: Union[torch.dtype, None] = None,
+    opt_params: dict = {
+        "lp_norm": 0.7,
+        "beta": 1e1,
+        "kappa": 1.01,
+        "iters": 20,
+        "tol": 0.0,
+        "early_stop": True,
+        "scale_gridsearch": False,
+    },
+    verbose: bool = False,
+) -> tuple:
+    # Params
+    lp_norm = max(opt_params["lp_norm"], 0.1)
+    beta = opt_params["beta"]
+    kappa = opt_params["kappa"]
+    iters = opt_params["iters"]
+    early_stop = opt_params["early_stop"]
+    tol = opt_params["tol"]
+
+    # Check
+    assert lp_norm <= 1.0, "lp_norm should be <=1"
+    assert beta > 0.0, "beta should be > 0"
+    assert kappa > 1.0, "kappa should be > 1"
+    assert iters > 1, "iters should be > 1"
+
+    # Cast/device
+    if device is None:
+        device = tensor.device
+    else:
+        device = torch.device(device)
+
+    if dtype is None:
+        dtype = float16 if (device.type == "cuda") else float32
+
+    W_f = tensor.to(device=device, dtype=dtype)
+    scale = scale.to(device=device, dtype=dtype)
+    zero = zero.to(device=device, dtype=dtype)
+
+    # Update scale: works slightly better. Tested on Llama2 only
+    if opt_params["scale_gridsearch"]:
+        scale = update_scale_grid_search(W_f, scale, zero, axis, min_max)
+
+    # Optimize for zero-point
+    best_error = 1e4
+    scale_prev, zero_prev = scale.clone(), zero.clone()
+    for i in range(iters):
+        W_q = torch.round(W_f * scale + zero).clamp(min_max[0], min_max[1])
+        W_r = (W_q - zero) / scale
+
+        # current_error = float(torch.pow(torch.abs(W_f - W_r), max(0.80, lp_norm)).mean())
+        current_error = float(torch.abs(W_f - W_r).mean())
+
+        if verbose:
+            print(i, np.round(current_error, 6))
+
+        if early_stop:
+            if best_error - current_error > tol:
+                best_error = current_error
+                scale_prev, zero_prev = scale.clone(), zero.clone()
+            else:
+                scale, zero = scale_prev.clone(), zero_prev.clone()
+                break
+
+        W_e = shrink_lp_op(W_f - W_r, beta, lp_norm)
+        zero = torch.mean(W_q - (W_f - W_e) * scale, axis=axis, keepdim=True)
+        beta *= kappa
+
+    # Clean-up
+    scale = scale.to(tensor.device)
+    zero = zero.to(tensor.device)
+    del W_f, W_q, W_r, W_e, scale_prev, zero_prev
+    torch.cuda.empty_cache()
+
+    W_q = torch.round(tensor * scale + zero).clamp(min_max[0], min_max[1])
+
+    return W_q, scale, zero
+
+
+# Proximal solver || W - dequantize(quantize(W))||_p^p
+@torch.inference_mode()
+def optimize_weights_proximal_legacy(
+    tensor: Tensor,
+    scale: Tensor,
+    zero: Tensor,
+    min_max: list,
+    axis: int = 0,
+    device: Union[str, None] = None,
+    opt_params: dict = {"lp_norm": 0.7, "beta": 1e1, "kappa": 1.01, "iters": 20},
+    verbose: bool = False,
+) -> tuple:
+    lp_norm, beta, kappa, iters = (
+        opt_params["lp_norm"],
+        opt_params["beta"],
+        opt_params["kappa"],
+        opt_params["iters"],
+    )
+
+    if device is None:
+        device = tensor.device
+    else:
+        device = torch.device(device)
+
+    dtype = float16 if (device.type == "cuda") else float32
+    W_f = tensor.to(dtype=dtype, device=device)
+    scale = scale.to(dtype=dtype, device=device)
+    zero = zero.to(dtype=dtype, device=device)
+
+    best_error = 1e4
+    for i in range(iters):
+        W_q = torch.round(W_f * scale + zero).clamp(min_max[0], min_max[1])
+        W_r = (W_q - zero) / scale
+        W_e = shrink_lp_op(W_f - W_r, beta, lp_norm)
+        zero = torch.mean(W_q - (W_f - W_e) * scale, axis=axis, keepdim=True)
+        beta *= kappa
+
+        current_error = float(torch.abs(W_f - W_r).mean())
+        if verbose:
+            print(i, np.round(current_error, 6))
+        if current_error < best_error:
+            best_error = current_error
+        else:
+            break
+
+    scale = scale.to(tensor.device)
+    zero = zero.to(tensor.device)
+    del W_f, W_q, W_r, W_e
+    torch.cuda.empty_cache()
+
+    W_q = torch.round(tensor * scale + zero).clamp(min_max[0], min_max[1])
+    return W_q, scale, zero
+
+
+# Default: fast with early stopping
+optimize_weights_proximal = optimize_weights_proximal_legacy
+
+# Slower, better quality: no early stoppping, more iterations
+optimize_weights_proximal_slow = partial(
+    optimize_weights_proximal_v2,
+    dtype=torch.float32,
+    opt_params={
+        "lp_norm": 0.7,
+        "beta": 1e1,
+        "kappa": 1.01,
+        "iters": 100,
+        "tol": 0.0,
+        "early_stop": False,
+        "scale_gridsearch": False,
+    },
+)
+
+##############################################################################################################
+# L1 with SGD optimizer: supports scale and W_q updates. L{p<1} fails with SGD
+
+
+class LinearSchedulerWithWarmStart(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, lr_start, lr_end, iters, warm_start=0, last_epoch=-1):
+        iters_wrm = max(0, int(iters * warm_start))
+        self.lr_wrm = (
+            np.linspace(lr_end, lr_start, iters_wrm)
+            if (iters_wrm > 0)
+            else np.array([])
+        )
+        self.lr_mid = np.linspace(lr_start, lr_end, iters - iters_wrm)
+        self.lr_sch = np.concatenate([self.lr_wrm, self.lr_mid])
+        self.idx = 0
+        super(LinearSchedulerWithWarmStart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        self.idx = min(self.idx, len(self.lr_sch) - 1)
+        out = [self.lr_sch[self.idx] for base_lr in self.base_lrs]
+        self.idx += 1
+        return out
+
+
+# SGD solver  || W - dequantize(quantize(W))||_1 (p=1 only, with additional fake inputs x)
+def optimize_weights_autograd(
+    tensor: Tensor,
+    scale: Tensor,
+    zero: Tensor,
+    min_max: list,
+    axis: int = 0,
+    device: Union[str, None] = None,
+    dtype: Union[torch.dtype, None] = float32,
+    opt_params: dict = {
+        "lr": 2e-3,
+        "iters": 2500,
+        "lr_schedule": False,
+        "update_Wq": False,
+        "use_fake_data": False,
+    },
+    data_params: dict = {"normalize": False, "data_rng": 10.0, "data_ctx": 32},
+    compile: bool = True,
+    verbose: bool = False,
+) -> tuple:
+    ref_device = scale.device
+    ref_dtype = scale.dtype
+
+    if device is None:
+        device = tensor.device
+    else:
+        device = torch.device(device)
+
+    W_f = tensor.to(dtype=dtype, device=device)
+
+    params = {}
+    params["scale"] = torch.nn.Parameter(
+        scale.to(dtype=dtype, device=device), requires_grad=True
+    )
+    params["zero"] = torch.nn.Parameter(
+        zero.to(dtype=dtype, device=device), requires_grad=True
+    )
+
+    if opt_params["update_Wq"]:
+        with torch.no_grad():
+            params["W_q"] = torch.round(W_f * params["scale"] + params["zero"]).clamp(
+                min_max[0], min_max[1]
+            )
+        params["W_q"] = torch.nn.Parameter(params["W_q"], requires_grad=True)
+
+    optimizer = torch.optim.AdamW(
+        [params[k] for k in params],
+        lr=opt_params["lr"],
+        betas=(0.9, 0.99),
+        eps=1e-06,
+        weight_decay=0.0,
+    )
+
+    if opt_params["lr_schedule"]:
+        scheduler = LinearSchedulerWithWarmStart(
+            optimizer,
+            lr_start=opt_params["lr"],
+            lr_end=1e-6,
+            iters=opt_params["iters"],
+            warm_start=0,
+        )
+    else:
+        scheduler = None
+
+    with torch.no_grad():
+        if data_params["normalize"]:
+            scale_loss = 1.0 / (tensor.abs().mean() + 1e-4)
+        else:
+            scale_loss = 1.0
+
+    def _loss_fct(output, target):
+        return torch.mean(torch.abs(scale_loss * (target - output)))  # L1
+
+    def _fake_quant_fixed_Wq(W_f):
+        # Quantize
+        W_q = torch.round(W_f * params["scale"] + params["zero"]).clamp(
+            min_max[0], min_max[1]
+        )
+        # Dequantize
+        W_r = (W_q - params["zero"]) / params["scale"]
+        return W_r
+
+    def _fake_quant_update_Wq(W_f):
+        # Quantize
+        W_q = torch.round(params["W_q"]).clamp(min_max[0], min_max[1])
+        # Dequantize
+        W_r = (W_q - params["zero"]) / params["scale"]
+        return W_r
+
+    if opt_params["update_Wq"]:
+        _fake_quant = _fake_quant_update_Wq
+    else:
+        _fake_quant = _fake_quant_fixed_Wq
+
+    def _fake_quant_loss(W_f):
+        return _loss_fct(_fake_quant(W_f), W_f)
+
+    def _fake_quant_loss_with_fake_data(W_f):
+        x = (
+            torch.rand(
+                [data_params["data_ctx"], W_f.shape[1]], device=device, dtype=dtype
+            )
+            - 0.5
+        ) * 2 ** data_params["data_rng"]
+        y_ref = torch.matmul(x, W_f.T)
+        y_pred = torch.matmul(x, _fake_quant(W_f).T)
+        return _loss_fct(y_pred, y_ref)
+
+    if opt_params["use_fake_data"]:
+        _fake_quant_loss = _fake_quant_loss_with_fake_data
+    else:
+        _fake_quant_loss = _fake_quant_loss
+
+    if compile:
+        _fake_quant_loss = torch.compile(_fake_quant_loss)
+
+    def _step(W_f):
+        optimizer.zero_grad()
+        loss = _fake_quant_loss(W_f)
+        loss.backward()
+        optimizer.step()
+        if scheduler is not None:
+            scheduler.step()
+        return np.round(loss.item(), 10)
+
+    with torch.no_grad():
+        _init_loss = _fake_quant_loss(W_f).item()
+
+    for i in range(opt_params["iters"]):
+        loss_out = _step(W_f)
+        if verbose and (i % 100) == 0:
+            print(i, loss_out)
+
+    with torch.no_grad():
+        _final_loss = _fake_quant_loss(W_f).item()
+
+    if _final_loss < _init_loss:
+        for k in params:
+            params[k] = params[k].data.detach()
+    else:
+        if verbose:
+            print("optimization failed...")
+        params = {"scale": scale, "zero": zero}
+
+    scale = params["scale"].to(device=ref_device, dtype=ref_dtype)
+    zero = params["zero"].to(device=ref_device, dtype=ref_dtype)
+
+    if "W_q" in params:
+        W_q = params["W_q"].to(device=ref_device, dtype=ref_dtype)
+    else:
+        W_q = (
+            torch.round(tensor * scale + zero)
+            .clamp(min_max[0], min_max[1])
+            .to(device=ref_device, dtype=ref_dtype)
+        )
+
+    del W_f
+    torch.cuda.empty_cache()
+    return W_q, scale, zero
+
+
+optimize_weights_autograd_main = partial(
+    optimize_weights_autograd,
+    dtype=torch.float32,
+    opt_params={
+        "lr": 2e-3,
+        "iters": 1000,
+        "lr_schedule": True,
+        "update_Wq": True,
+        "use_fake_data": False,
+    },
+    verbose=False,
+    compile=True,
+)
+
+optimize_weights_autograd_fakedata = partial(
+    optimize_weights_autograd,
+    dtype=torch.float32,
+    opt_params={
+        "lr": 2e-3,
+        "iters": 1000,
+        "lr_schedule": True,
+        "update_Wq": True,
+        "use_fake_data": True,
+    },
+    data_params={"normalize": False, "data_rng": 10.0, "data_ctx": 32},
+    verbose=False,
+    compile=True,
+)

lm-quant-toolkit/.deps/hqq/hqq/core/peft.py

@@ -0,0 +1,545 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+import torch
+from torch import float16, bfloat16, float32
+from torch import Tensor, nn
+import numpy as np
+from .quantize import HQQLinear, Quantizer
+from .utils import cleanup
+from typing import Union
+
+
+# Return trainable weight matrix
+def _get_dense_param(
+    in_features: int,
+    out_features: int,
+    device="cuda",
+    trainable: bool = True,
+    dtype: torch.dtype = float32,
+) -> Tensor:
+    W = (
+        nn.Linear(in_features, out_features, bias=None)
+        .weight.data.t()
+        .to(dtype)
+        .to(device)
+        .contiguous()
+    )
+    return nn.Parameter(W, requires_grad=trainable)
+
+
+# Applies LoRA to a linear module (include HQQLinear)
+class HQQLinearLoRA(nn.Module):
+    def __init__(self, linear_layer: nn.Module, peft_config: dict):
+        super().__init__()
+
+        # Device
+        self.device = (
+            linear_layer.device
+            if hasattr(linear_layer, "device")
+            else next(linear_layer.parameters()).device
+        )
+        self.train_dtype = (
+            peft_config["train_dtype"]
+            if ("train_dtype" in peft_config)
+            else torch.float
+        )
+
+        # Linear layer
+        self.linear_layer = linear_layer
+        self.in_features = linear_layer.in_features
+        self.out_features = linear_layer.out_features
+
+        # Bias
+        self.bias = None if (linear_layer.bias is None) else linear_layer.bias.clone()
+        self.linear_layer.bias = None
+        self.train_bias = (
+            peft_config["train_bias"] if ("train_bias" in peft_config) else False
+        )
+
+        if self.bias is not None:
+            self.bias = nn.Parameter(self.bias, requires_grad=self.train_bias)
+            if self.train_bias:
+                self.bias.data = self.bias.data.to(self.train_dtype)
+
+        if (self.bias is None) and self.train_bias:
+            self.bias = nn.Parameter(
+                torch.zeros(
+                    (self.out_features,), device=self.device, dtype=self.train_dtype
+                ),
+                requires_grad=True,
+            )
+
+        # Dropout
+        if "dropout" in peft_config:
+            self.peft_drop = (
+                nn.Dropout(p=peft_config["dropout"])
+                if (peft_config["dropout"] > 0.0)
+                else nn.Identity()
+            )
+        else:
+            self.peft_drop = nn.Identity()
+
+        # LoRA A/B
+        self.peft_config = peft_config
+        self.lora_alpha = peft_config["lora_alpha"]
+        self.r = peft_config["r"]
+        self.scaling = self.lora_alpha / self.r
+
+        self.lora_A = _get_dense_param(
+            self.in_features,
+            self.r,
+            device=self.device,
+            trainable=True,
+            dtype=self.train_dtype,
+        )
+
+        self.lora_B = _get_dense_param(
+            self.r,
+            self.out_features,
+            device=self.device,
+            trainable=True,
+            dtype=self.train_dtype,
+        )
+
+        # LoRA weights init
+        if "lora_init" in peft_config:
+            # Set lora init
+            assert (
+                peft_config["lora_init"]["lora_A"].shape[0],
+                peft_config["lora_init"]["lora_B"].shape[1],
+            ) == (self.in_features, self.out_features), (
+                "Invalid init LoRA weight shapes. Expected: lora_A: "
+                + str(self.in_features)
+                + " x r , lora_B: r x "
+                + str(self.out_features)
+                + ")"
+            )
+            self.lora_A.data = peft_config["lora_init"]["lora_A"].to(
+                device=self.device, dtype=self.train_dtype
+            )
+            self.lora_B.data = peft_config["lora_init"]["lora_B"].to(
+                device=self.device, dtype=self.train_dtype
+            )
+        else:
+            # Init weights, as as the original LoRA implementation
+            nn.init.kaiming_uniform_(self.lora_A, a=np.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def forward_lora(self, x: Tensor) -> Tensor:  # output is self.train_dtype
+        return (
+            torch.matmul(
+                torch.matmul(self.peft_drop(x.to(self.lora_A.dtype)), self.lora_A),
+                self.lora_B,
+            )
+            * self.scaling
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+
+        # Forward with the base linear layer
+        out = self.linear_layer(x)
+
+        # LoRA
+        if self.train_bias:
+            out += (self.forward_lora(x) + self.bias).to(x_dtype)
+
+        else:
+            out += self.forward_lora(x).to(x_dtype)
+            if self.bias is not None:
+                out += self.bias
+
+        return out
+
+    def merge_and_quantize(self, quant_config: dict) -> nn.Module:
+        # not easy to figure out the type for "any" potential linear layer (nn.Linear, HQQLinear, etc.)
+        for _dtype in [float16, bfloat16, float32]:
+            try:
+                # Get initial weights
+                W = self.linear_layer(
+                    torch.eye(self.in_features, device=self.device, dtype=_dtype)
+                ).t()  # == self.linear_layer.dequantize()
+            except Exception:
+                pass
+
+        # Merge weights
+        W += (
+            (torch.matmul(self.lora_A.data, self.lora_B.data) * self.scaling)
+            .t()
+            .to(W.dtype)
+        )
+
+        # New HQQ layer
+        new_hqq_layer = HQQLinear(None, quant_config)
+        new_hqq_layer.bias = None if (self.bias is None) else self.bias.clone()
+        new_hqq_layer.quantize(W, **quant_config)
+
+        return new_hqq_layer
+
+    def cast(self, dtype: torch.dtype = float16):
+        self.lora_A.data = self.lora_A.data.to(dtype)
+        self.lora_B.data = self.lora_B.data.to(dtype)
+        if self.bias is not None:
+            self.bias.data = self.bias.data.to(dtype)
+        if isinstance(self.scaling, nn.Parameter):
+            self.scaling.data = self.scaling.data.to(dtype)
+        if isinstance(self.scaling, Tensor):
+            self.scaling = self.scaling.to(dtype)
+        return self
+
+    def state_dict(self):
+        return {
+            "lora_A": self.lora_A.data,
+            "lora_B": self.lora_B.data,
+            "scaling": self.scaling,
+            "bias": self.bias,
+        }
+
+    def load_state_dict(self, state_dict):
+        self.lora_A.data = state_dict["lora_A"].data.to(self.device)
+        self.lora_B.data = state_dict["lora_B"].data.to(self.device)
+
+        if state_dict["bias"] is not None:
+            self.bias.data = state_dict["bias"].data.to(self.device)
+
+        # Handle different use-cases of scaling
+        if isinstance(state_dict["scaling"], (int, float)):
+            self.scaling = state_dict["scaling"]
+        if isinstance(state_dict["scaling"], nn.Parameter):
+            self.scaling.data = state_dict["scaling"].data.to(self.device)
+        if isinstance(state_dict["scaling"], Tensor):
+            self.scaling = state_dict["scaling"].to(self.device)
+
+
+# LoRA with fake quantization
+class HQQLinearLoRAWithFakeQuant(HQQLinearLoRA):
+    def __init__(self, linear_layer: nn.Module, peft_config: dict):
+        super(HQQLinearLoRAWithFakeQuant, self).__init__(linear_layer, peft_config)
+        self.quant_param = peft_config["quant_param"]
+
+    # @torch.no_grad()
+    # @torch.compile()
+    def fake_quant(self, weight):
+        if self.quant_param:
+            W_q, meta = Quantizer.quantize(
+                weight, **self.quant_param, bitpack=False
+            )  # todo: clone() tensor
+            weight_est = Quantizer.dequantize(W_q, meta)
+        else:
+            weight_est = weight
+        return weight_est
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+
+        # Get initial weights
+        W = self.linear_layer(
+            torch.eye(self.in_features, device=self.device, dtype=x_dtype)
+        ).t()  # == self.linear_layer.dequantize()
+
+        # Merge weights
+        W += (torch.matmul(self.lora_A, self.lora_B) * self.scaling).t().to(W.dtype)
+
+        # Fake quant
+        W = self.fake_quant(W).to(x_dtype)
+
+        # Matmul
+        out = torch.matmul(x, W.t())
+
+        # Bias
+        if self.bias is not None:
+            out += self.bias
+
+        out = out.to(x_dtype)
+
+        return out
+
+
+# Experimental
+class HQQLinearGroupedProj(nn.Module):
+    def __init__(self, linear_layer: nn.Module, peft_config: dict):
+        super().__init__()
+
+        # Device
+        self.device = (
+            linear_layer.device
+            if hasattr(linear_layer, "device")
+            else next(linear_layer.parameters()).device
+        )
+        self.train_dtype = (
+            peft_config["train_dtype"]
+            if ("train_dtype" in peft_config)
+            else torch.float
+        )
+
+        # Linear layer
+        self.linear_layer = linear_layer
+        self.in_features = linear_layer.in_features
+        self.out_features = linear_layer.out_features
+        self.bias = None if (linear_layer.bias is None) else linear_layer.bias.clone()
+
+        # Turn-off bias in the linear layer
+        self.linear_layer.bias = None
+
+        # Group proj
+        self.peft_config = peft_config
+        self.proj_size = peft_config["proj_size"]
+        self.proj_num = peft_config["proj_num"]
+        self.proj = nn.Parameter(
+            torch.stack(
+                [torch.eye(self.proj_size, dtype=self.train_dtype, device=self.device)]
+                * self.proj_num
+            )
+        )
+        if peft_config["zero_trainable"]:
+            self.linear_layer.meta["zero"] = nn.Parameter(
+                self.linear_layer.meta["zero"].to(self.train_dtype), requires_grad=True
+            )
+
+    @torch.compile()
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+
+        # Forward with base linear
+        with torch.no_grad():
+            W = self.linear_layer.dequantize().clone()
+        # W = self.linear_layer(torch.eye(self.in_features, device=self.device, dtype=x_dtype)).t()
+        shape = W.shape
+
+        # Grouped proj
+        proj_b, gs = self.proj.shape[0], self.proj.shape[1]
+        W = (
+            torch.matmul(self.proj, W.reshape((proj_b, gs, -1)).to(self.proj.dtype))
+            .to(x_dtype)
+            .reshape(shape)
+        )
+
+        # Matmul
+        out = torch.matmul(x, W.t())
+
+        # Bias
+        if self.bias is not None:
+            out += self.bias
+
+        out = out.to(x_dtype)
+
+        return out
+
+    def cast(self, dtype: torch.dtype = float16):
+        self.proj.data = self.proj.data.to(dtype)
+        self.linear_layer.meta["zero"] = self.linear_layer.meta["zero"].to(dtype)
+        if self.bias is not None:
+            if self.bias.requires_grad:
+                self.bias.data = self.bias.data.to(dtype)
+            else:
+                self.bias = self.bias.to(dtype)
+        return self
+
+    def state_dict(self):
+        return {
+            "proj": self.proj.data,
+            "bias": self.bias,
+            "peft_config": self.peft_config,
+        }
+
+    def load_state_dict(self, state_dict):
+        self.proj.data = state_dict["proj"].data.to(self.device)
+        self.bias = state_dict["bias"] if ("bias" in state_dict) else None
+        self.bias = self.bias.to(self.device) if (self.bias is not None) else None
+        self.peft_config = state_dict["peft_config"]
+
+
+_HQQ_LORA_CLASSES = [HQQLinearLoRA, HQQLinearLoRAWithFakeQuant, HQQLinearGroupedProj]
+_HQQ_LORA_MAPPING = {
+    "default": HQQLinearLoRA,
+    "lora_with_fakequant": HQQLinearLoRAWithFakeQuant,
+    "grouped_proj": HQQLinearGroupedProj,
+}
+
+
+def is_hqq_lora_layer(layer):
+    return type(layer) in _HQQ_LORA_CLASSES
+
+
+##################################################################################################################
+def autoname_modules(model):
+    for name, module in model.named_modules():
+        module.name = name
+
+
+# Patching functions
+def patch_linear_add_peft(layer: nn.Module, patch_params: Union[dict,None]) -> nn.Module:
+    _peft_config = patch_params
+    if _peft_config:
+        lora_type = (
+            _peft_config["lora_type"] if ("lora_type" in _peft_config) else "default"
+        )
+        new_layer = _HQQ_LORA_MAPPING[lora_type](layer, _peft_config)
+    else:
+        new_layer = layer
+    return new_layer
+
+
+def patch_linear_merge_peft(layer: nn.Module, patch_params: Union[dict,None]) -> nn.Module:
+    _quant_config = patch_params
+    if _quant_config:
+        new_layer = layer.merge_and_quantize(_quant_config)
+        del layer
+        cleanup()
+    else:
+        new_layer = layer
+    return new_layer
+
+
+def patch_linear_cast_peft(layer: nn.Module, patch_params: Union[dict,None]) -> nn.Module:
+    if is_hqq_lora_layer(layer):
+        layer.cast(patch_params)
+    return layer
+
+
+# Putting it all together
+class PeftUtils:
+    @classmethod
+    def get_base_class(cls, model, base_class):
+        # Get base class
+        if base_class is None:
+            base_class = getattr(model, "base_class", None)
+
+        assert (
+            base_class is not None
+        ), "You need to provide the base HQQ class (LlamaHQQ, MixtralHQQ, etc.) as model.base_class or as an argument base_class=LlamaHQQ. Use AutoHQQHFMode if the model architecture is not supported."
+
+        return base_class
+
+    @classmethod
+    def add_lora(
+        cls, model, peft_config, base_class: bool = None, verbose: bool = True
+    ) -> None:
+        # Base classs
+        base_class = cls.get_base_class(model, base_class)
+
+        base_class.setup_model(model)
+
+        # Freeze
+        for param in model.parameters():
+            param.requires_grad = False
+
+        # Patch
+        base_class.patch_linearlayers(
+            model, patch_linear_add_peft, peft_config, verbose=verbose
+        )
+
+        # Rename modules
+        autoname_modules(model)
+
+        # add config to model
+        model.peft_config = peft_config
+
+    @classmethod
+    def merge_lora(
+        cls, model, merge_lora_params, base_class: bool = None, verbose: bool = True
+    ) -> None:
+        # Base classs
+        base_class = cls.get_base_class(model, base_class)
+
+        # Patch
+        base_class.patch_linearlayers(
+            model, patch_linear_merge_peft, merge_lora_params, verbose=verbose
+        )
+
+    @classmethod
+    def cast_lora_weights(
+        cls, model, dtype: torch.dtype, base_class: bool = None, verbose: bool = True
+    ) -> None:
+        # Base classs
+        base_class = cls.get_base_class(model, base_class)
+
+        base_class.setup_model(model)
+
+        # Linear tags
+        linear_tags = model.linear_tags
+
+        # Patch
+        base_class.patch_linearlayers(
+            model,
+            patch_linear_cast_peft,
+            dict([(linear_tag, dtype) for linear_tag in linear_tags]),
+            verbose=verbose,
+        )
+
+    @classmethod
+    def save_lora_weights(
+        cls, model, filename: str, base_class: bool = None, verbose: bool = True
+    ) -> None:
+        # Base classs
+        base_class = cls.get_base_class(model, base_class)
+
+        base_class.setup_model(model)
+
+        lora_global_params = {}
+
+        def _patch_linear_save_weights(layer, patch_params, return_layer=True):
+            if is_hqq_lora_layer(layer):
+                lora_global_params[layer.name] = layer.state_dict()
+            if return_layer:
+                return layer
+
+        # Linear tags
+        linear_tags = model.linear_tags
+
+        # Patch
+        base_class.patch_linearlayers(
+            model,
+            _patch_linear_save_weights,
+            dict([(linear_tag, None) for linear_tag in linear_tags]),
+            verbose=verbose,
+        )
+
+        # save
+        torch.save(
+            {"peft_config": model.peft_config, "parameters": lora_global_params},
+            filename,
+        )
+
+    @classmethod
+    def load_lora_weights(
+        cls, model, filename: str, base_class: bool = None, verbose: bool = True
+    ) -> None:
+        # Base classs
+        base_class = cls.get_base_class(model, base_class)
+
+        base_class.setup_model(model)
+
+        lora_data = torch.load(filename, map_location="cpu")
+
+        # V0.2 format: automatically creates lora modules if the model doesn't contain it
+        if ("peft_config" in lora_data) and ("parameters" in lora_data):
+            peft_config = lora_data["peft_config"]
+            lora_global_params = lora_data["parameters"]
+            if not hasattr(model, "peft_config"):
+                cls.add_lora(model=model, peft_config=peft_config)
+
+        # v0.1 format
+        else:
+            if not hasattr(model, "peft_config"):
+                raise Exception(
+                    "Using older version of lora weights. LoRa modules should be manually added in this case."
+                )
+            lora_global_params = lora_data
+
+        def _patch_linear_load_weights(layer, patch_params, return_layer=True):
+            if is_hqq_lora_layer(layer):
+                layer.load_state_dict(lora_global_params[layer.name])
+            if return_layer:
+                return layer
+
+        # Linear tags
+        linear_tags = model.linear_tags
+
+        # Patch
+        base_class.patch_linearlayers(
+            model,
+            _patch_linear_load_weights,
+            dict([(linear_tag, None) for linear_tag in linear_tags]),
+            verbose=verbose,
+        )

lm-quant-toolkit/.deps/hqq/hqq/core/quantize.py

@@ -0,0 +1,948 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+import copy
+from enum import Enum
+from typing import Union
+
+import torch
+from torch import Tensor, float16, int32, nn, uint8
+
+from .bitpack import BitPack
+from .optimize import optimize_weights_proximal
+from .utils import is_divisible
+
+
+# Main HQQ Quantizer
+class Quantizer:
+    SUPPORTED_BITS = [8, 6, 5, 4, 3, 2, 1]
+    optimize_weights = optimize_weights_proximal
+
+    bit_to_packing = {
+        8: "8bit_u8",
+        6: "8bit_u8",  # todo: bitpacking
+        5: "8bit_u8",  # todo: bitpacking
+        4: "4bit_u8",
+        3: "3bit_32",
+        2: "2bit_u8",
+        1: "1bit_u8",
+    }
+
+    pack = {
+        "8bit_u8": BitPack.pack_8bit_u8,
+        "4bit_u8": BitPack.pack_4bit_u8,
+        "3bit_32": BitPack.pack_3bit_32,
+        "2bit_u8": BitPack.pack_2bit_u8,
+        "1bit_u8": BitPack.pack_1bit_u8,
+    }
+
+    unpack = {
+        "8bit_u8": BitPack.unpack_8bit_u8,
+        "4bit_u8": BitPack.unpack_4bit_u8,
+        "3bit_32": BitPack.unpack_3bit_32,
+        "2bit_u8": BitPack.unpack_2bit_u8,
+        "1bit_u8": BitPack.unpack_1bit_u8,
+    }
+
+    unpack_view_dtype = {
+        "8bit_u8": uint8,
+        "4bit_u8": uint8,
+        "3bit_32": int32,
+        "2bit_u8": uint8,
+        "1bit_u8": uint8,
+    }
+
+    @classmethod
+    def quantize(
+        cls,
+        tensor: Tensor,
+        nbits: int = 4,
+        channel_wise: bool = True,
+        group_size: int = 64,
+        optimize: bool = False,
+        round_zero: bool = False,
+        axis: int = 0,
+        bitpack: bool = True,
+        compute_dtype: Union[torch.dtype, None] = None,
+        view_as_float: bool = False,
+        device: str = "cuda",
+    ) -> tuple:
+        assert nbits in Quantizer.SUPPORTED_BITS, (
+            "nbits=" + str(nbits) + " not supported."
+        )
+        assert axis in [0, 1], "axis should be either 0 or 1"
+        if group_size is not None:
+            assert is_divisible(tensor.numel(), group_size), (
+                "group_size should be divisble by the total tensor dimensions. shape: "
+                + str(tensor.shape)
+                + ", group_size: "
+                + str(group_size)
+            )
+
+        W = tensor.float()
+        shape = W.shape
+
+        # Reshape for grouping
+        if (group_size is not None) and channel_wise:
+            W = (
+                W.reshape([-1, group_size])
+                if (axis == 1)
+                else W.reshape([group_size, -1])
+            )
+
+        # Get min/max values
+        if not channel_wise:
+            _min, _max = W.min(), W.max()
+            optimize = False
+        else:
+            _min = W.min(axis=axis, keepdim=True)[0]
+            _max = W.max(axis=axis, keepdim=True)[0]
+
+        max_v = 2**nbits - 1
+        min_v = 0
+        min_max = [min_v, max_v]
+
+        # Note: here we work with the inverse of the scale to avoid division and quantize instead via W*scale + zero, the scale is inverted later on.
+        scale = (max_v / (_max - _min)).clamp(
+            max=2e4
+        )  # clamp to avoid half-precision problems
+        zero = -_min * scale
+
+        # Round zero as in: https://github.com/casper-hansen/AutoAWQ/blob/main/awq/quantize/quantizer.py#L42C9-L42C14
+        if round_zero:
+            zero = torch.round(zero)
+
+        # Fine-tune weights
+        # if False:
+        if optimize:
+            W_q, scale, zero = Quantizer.optimize_weights(
+                tensor=W,
+                scale=scale,
+                zero=zero,
+                min_max=min_max,
+                axis=axis,
+                device=device,
+            )
+        else:
+            W_q = torch.round(W * scale + zero).clamp(min_max[0], min_max[1])
+
+        # Store meta-data (we invert the scale for dequantization)
+        meta = {
+            "nbits": nbits,
+            "group_size": group_size,
+            "shape": shape,
+            "scale": 1.0 / scale,
+            "zero": zero,
+            "axis": axis,
+            "packing": Quantizer.bit_to_packing[nbits],
+        }
+        meta["unpack_view_dtype"] = Quantizer.unpack_view_dtype[meta["packing"]]
+
+        # Pack bits
+        meta["view_as_float"] = view_as_float
+        if bitpack:
+            W_q = Quantizer.pack[meta["packing"]](W_q)
+            if view_as_float:
+                W_q = W_q.view(
+                    torch.float32 if compute_dtype is None else compute_dtype
+                )  # store quantized weights as compute_dtype
+        else:
+            W_q = W_q.to(tensor.dtype)
+            meta["packing"] = None
+
+        # cleanup
+        del W, _min, _max
+        torch.cuda.empty_cache()
+
+        return W_q, meta
+
+    # Main dequantization: bit_unpacking > (W_q - z)*s > reshape
+    @classmethod
+    def dequantize(cls, W_q: Tensor, meta: dict) -> Tensor:
+        compute_dtype = meta["compute_dtype"] if ("compute_dtype" in meta) else float16
+        if meta["packing"]:
+            if meta["view_as_float"]:
+                W_q = W_q.view(meta["unpack_view_dtype"])
+            W_r = Quantizer.unpack[meta["packing"]](W_q, dtype=compute_dtype)
+            if meta["nbits"] == 3:
+                W_r = W_r[
+                    : meta["group_size"]
+                    if meta["axis"] == 0
+                    else meta["shape"][0] * meta["shape"][1] // meta["group_size"]
+                ]
+        else:
+            W_r = W_q.to(compute_dtype)
+        zero = meta["zero"].to(compute_dtype)
+        scale = meta["scale"].to(compute_dtype)
+        W_r = ((W_r - zero) * scale).reshape(meta["shape"])
+        return W_r
+
+    @classmethod
+    def to_inplace(cls, W_q: Tensor, meta: dict, device) -> tuple:
+        compute_dtype = meta["compute_dtype"] if ("compute_dtype" in meta) else float16
+        if W_q is not None:
+            W_q = W_q.to(device).contiguous()
+        for key in meta:
+            if type(meta[key]) == torch.Tensor:
+                meta[key] = (
+                    (
+                        meta[key].to(compute_dtype)
+                        if torch.is_floating_point(meta[key])
+                        else meta[key]
+                    )
+                    .to(device)
+                    .contiguous()
+                )
+        return W_q, meta
+
+    @classmethod
+    def to_ooplace(cls, W_q: Tensor, meta: dict, device) -> tuple:
+        compute_dtype = meta["compute_dtype"] if ("compute_dtype" in meta) else float16
+        if W_q is not None:
+            W_q_c = W_q.to(device).contiguous()
+        else:
+            W_q_c = None
+        meta_c = {}
+        for key in meta:
+            if type(meta[key]) == torch.Tensor:
+                meta_c[key] = (
+                    (
+                        meta[key].to(compute_dtype)
+                        if torch.is_floating_point(meta[key])
+                        else meta[key]
+                    )
+                    .to(device)
+                    .contiguous()
+                )
+            else:
+                meta_c[key] = meta[key]
+        return W_q_c, meta_c
+
+    @classmethod
+    def cuda(cls, W_q: Tensor, meta: dict, device) -> tuple:
+        return Quantizer.to_inplace(W_q, meta, device=device)
+
+    @classmethod
+    def cpu(cls, W_q: Tensor, meta: dict) -> tuple:
+        return Quantizer.to_ooplace(W_q, meta, device="cpu")
+
+
+# Check Aten/CUDA backend if available
+hqq_aten_is_available = False
+try:
+    import hqq_aten
+
+    hqq_aten_is_available = True
+except Exception:
+    hqq_aten = None
+    hqq_aten_is_available = False
+
+
+class HQQBackend(Enum):
+    # Name of the forward functions
+    PYTORCH = "forward_pytorch_backprop"
+    PYTORCH_COMPILE = "forward_pytorch_backprop_compile"
+    ATEN = "forward_aten_backprop"
+
+    # Alias for backward compatibility
+    PYTORCH_BACKPROP = "forward_pytorch_backprop"
+    PYTORCH_BACKPROP_COMPILE = "forward_pytorch_backprop_compile"
+    ATEN_BACKPROP = "forward_aten_backprop"
+
+    PYTORCH_FORWARD = "forward_pytorch"
+    PYTORCH_FORWARD_COMPILE = "forward_pytorch_compile"
+    ATEN_FORWARD = "forward_aten"
+
+    # Experimental
+    ATEN_FORWARD_INT8 = "forward_aten_int8"
+
+
+# No cache: less memory, slower
+class HQQMatmulNoCacheDeq(torch.autograd.Function):
+    @staticmethod
+    def forward(x: Tensor, dequantize, bias: Tensor):
+        out = torch.matmul(x, dequantize().t())
+        if bias is not None:
+            out += bias
+        return out
+
+    @staticmethod
+    def setup_context(ctx, inputs, outputs):
+        x, dequantize, bias = inputs
+        ctx.save_for_backward(x, bias)
+        ctx.dequantize = dequantize
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, bias = ctx.saved_tensors
+
+        grad_input = grad_weight = grad_bias = None
+
+        if ctx.needs_input_grad[0]:
+            grad_input = torch.matmul(grad_output, ctx.dequantize())
+
+        # weight grad for frozen quantized weights not defined
+        # if ctx.needs_input_grad[1]:
+        #   grad_weight = torch.matmul(grad_output.t(), x)
+
+        if bias is not None and ctx.needs_input_grad[2]:
+            grad_bias = grad_output.sum(0)
+
+        return grad_input, grad_weight, grad_bias
+
+
+class HQQMatmulNoCacheMul(torch.autograd.Function):
+    @staticmethod
+    def forward(x, matmul, bias):
+        out = matmul(x, transpose=True)
+        if bias is not None:
+            out += bias
+        return out
+
+    @staticmethod
+    def setup_context(ctx, inputs, outputs):
+        x, matmul, bias = inputs
+        ctx.save_for_backward(x, bias)
+        ctx.matmul = matmul
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, bias = ctx.saved_tensors
+
+        grad_input = grad_weight = grad_bias = None
+
+        if ctx.needs_input_grad[0]:
+            grad_input = ctx.matmul(grad_output, transpose=False)
+
+        # weight grad for frozen quantized weights not defined
+        # if ctx.needs_input_grad[1]:
+        #   grad_weight = torch.matmul(grad_output.t(), x)
+
+        if bias is not None and ctx.needs_input_grad[2]:
+            grad_bias = grad_output.sum(0)
+
+        return grad_input, grad_weight, grad_bias
+
+
+# Cache dequantized tensor: Faster but needs more memory
+class HQQMatmulCachedDeq(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, hqq_layer, bias):
+        weight_tmp = hqq_layer.dequantize()
+        out = torch.matmul(x, weight_tmp.t())
+        if bias is not None:
+            out += bias
+
+        ctx.save_for_backward(x, bias, weight_tmp)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, bias, weight_tmp = ctx.saved_tensors
+
+        grad_input = grad_weight = grad_bias = None
+
+        if ctx.needs_input_grad[0]:
+            grad_input = torch.matmul(grad_output, weight_tmp)
+
+        del weight_tmp
+
+        if bias is not None and ctx.needs_input_grad[2]:
+            grad_bias = grad_output.sum(0)
+
+        return grad_input, grad_weight, grad_bias
+
+
+# Main linear layer
+class HQQLinear(nn.Module):
+    # Default backend
+    backend = HQQBackend.PYTORCH
+
+    def __init__(
+        self,
+        linear_layer: Union[nn.Module, None],
+        quant_config: dict,
+        del_orig: bool = True,
+        compute_dtype: torch.dtype = float16,
+        device: str = "cuda",
+        initialize: bool = True,
+    ):
+        super().__init__()
+        self.ready = False
+        self.in_gpu = False
+        self.bias = None
+        self.device = device
+        self.compute_dtype = compute_dtype
+        self.quant_config = copy.deepcopy(quant_config)
+        self.del_orig = del_orig
+        self.offload_meta = (
+            self.quant_config.pop("offload_meta")
+            if (self.quant_config is not None)
+            else None
+        )
+
+        self.set_backend(HQQLinear.backend)
+
+        self.linear_layer = linear_layer
+        self.W_q = None
+        self.meta = None
+
+        if initialize:
+            self.initialize()
+
+    def initialize(self):
+        if self.linear_layer is not None:
+            self.quantize(self.linear_layer.weight.data, **self.quant_config)
+            self.bias = (
+                None
+                if (self.linear_layer.bias is None)
+                else self.linear_layer.bias.to(
+                    device=self.device, dtype=self.compute_dtype
+                )
+            )
+
+        if self.del_orig:
+            del self.linear_layer
+        torch.cuda.empty_cache()
+
+    def extra_repr(self) -> str:
+        out = ""
+        if hasattr(self, "meta"):
+            if self.meta is not None:
+                in_features, out_features = self.meta["shape"][::-1]
+                out = f"in_features={in_features}, out_features={out_features}, bias={self.bias is not None}"
+        return out
+
+    # Set backends
+    @classmethod
+    def set_backend(cls, backend: HQQBackend):
+        if "aten" in backend.value:
+            if hqq_aten_is_available is False:
+                print(
+                    "ATEN/CUDA backend not availabe. Make sure you install the hqq_aten library."
+                )
+                return
+            print(
+                "Warning: the ATEN/CUDA backend only supports axis=0 and GPU runtime."
+            )
+        HQQLinear.backend = backend
+        cls.forward = getattr(cls, backend.value)
+
+    # TODO: rewrite this mess
+    def cuda(self, device):
+        self.meta["compute_dtype"] = self.compute_dtype
+
+        if type(self.W_q) == nn.parameter.Parameter:
+            self.W_q.data, self.meta = Quantizer.cuda(self.W_q.data, self.meta, device)
+        else:
+            self.W_q, self.meta = Quantizer.cuda(self.W_q, self.meta, device)
+
+        if self.meta["quant_zero"]:
+            if "zero_q" in self.meta:
+                self.meta["zero_q"], self.meta["meta_zero"] = Quantizer.cuda(
+                    self.meta["zero_q"], self.meta["meta_zero"], device
+                )
+            else:
+                _, self.meta["meta_zero"] = Quantizer.cuda(
+                    None, self.meta["meta_zero"], device
+                )
+        else:
+            self.meta["zero"] = self.meta["zero"].to(device)
+
+        if self.meta["quant_scale"]:
+            if "scale_q" in self.meta:
+                self.meta["scale_q"], self.meta["meta_scale"] = Quantizer.cuda(
+                    self.meta["scale_q"], self.meta["meta_scale"], device
+                )
+            else:
+                _, self.meta["meta_scale"] = Quantizer.cuda(
+                    None, self.meta["meta_scale"], device
+                )
+        else:
+            self.meta["scale"] = self.meta["scale"].to(device)
+
+        # #Use zero/scale with streams for dequantization is faster than packing in "zero_scale"
+        # for key in ["zero", "zero_q", "scale", "scale_q"]:
+        #     if((key in self.meta) and self.offload_meta):
+        #         self.meta[key] = self.meta[key].contiguous().cpu().pin_memory()
+
+        if self.offload_meta:
+            if "zero_scale" not in self.meta:
+                if self.meta["quant_scale"] and self.meta["quant_zero"]:
+                    self.meta["zero_scale"] = torch.stack(
+                        (self.meta["zero_q"], self.meta["scale_q"])
+                    )
+                    del self.meta["scale_q"], self.meta["zero_q"]
+                else:
+                    self.meta["zero_scale"] = torch.stack(
+                        (self.meta["zero"], self.meta["scale"])
+                    ).to(self.compute_dtype)
+                    del self.meta["scale"], self.meta["zero"]
+
+            self.meta["zero_scale"] = (
+                self.meta["zero_scale"].contiguous().cpu().pin_memory()
+            )
+
+        if self.bias is not None:
+            self.bias = self.bias.to(device=device, dtype=self.compute_dtype)
+
+        self.W_q = nn.Parameter(self.W_q, requires_grad=False)
+        self.device = device
+        self.in_gpu = True
+
+        torch.cuda.empty_cache()
+
+        return self
+
+    def to(self, *args, **kwargs):
+        # TODO: later
+        return self
+
+    # TODO: later
+    # def to_empty(self, device, recurse=True):
+    #     return self.cuda(device)
+
+    def type(self, dst_type):
+        # TODO: later
+        return self
+
+    def half(self, *args, **kwargs):
+        return self
+
+    def bfloat16(self, *args, **kwargs):
+        # TODO: later
+        return self
+
+    def float(self, *args, **kwargs):
+        # TODO: later
+        return self
+
+    def double(self, *args, **kwargs):
+        return self
+
+    def cpu(self):
+        # TODO: later
+        return self
+
+    def state_dict(self, *args, **kwargs):  # nn.Module override compatible
+        state = {"W_q": self.W_q, "meta": self.meta, "bias": self.bias}
+        if "destination" in kwargs and "prefix" in kwargs:
+            for key, value in state.items():
+                kwargs["destination"][kwargs["prefix"] + key] = value
+        return state
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        W_q_key = prefix + "W_q"
+        meta_key = prefix + "meta"
+        bias_key = prefix + "bias"
+
+        if W_q_key not in state_dict:
+            missing_keys.append(W_q_key)
+        if meta_key not in state_dict:
+            missing_keys.append(meta_key)
+        if missing_keys:
+            return  # Can't load weights if either weight or meta is missing
+
+        W_q = nn.Parameter(state_dict.pop(W_q_key), requires_grad=False)
+        meta = state_dict.pop(meta_key)
+        bias = state_dict.pop(bias_key, None)
+
+        unexpected_keys += state_dict.keys()
+
+        self.load_state_dict({"W_q": W_q, "meta": meta, "bias": bias}, strict)
+
+    def load_state_dict(self, state_dict, strict=True, assign=False):
+        self.W_q = state_dict["W_q"]
+        self.meta = state_dict["meta"]
+        self.bias = state_dict["bias"] if ("bias" in state_dict) else None
+
+        # Meta-data offloading
+        self.offload_meta = False
+        for key in ["zero", "zero_q", "scale", "scale_q", "zero_scale"]:
+            if key in self.meta:
+                if self.meta[key].device.type == "cpu":
+                    self.offload_meta = True
+                    self.meta[key] = self.meta[key].contiguous().pin_memory()
+
+        # Float view settings
+        if "unpack_view_dtype" not in self.meta:
+            self.meta["unpack_view_dtype"] = Quantizer.unpack_view_dtype[
+                self.meta["packing"]
+            ]
+
+        if "view_as_float" not in self.meta:
+            self.meta["view_as_float"] = False
+
+        if "meta_scale" in self.meta:
+            if "view_as_float" not in self.meta["meta_scale"]:
+                self.meta["meta_scale"]["view_as_float"] = False
+
+        if "meta_zero" in self.meta:
+            if "view_as_float" not in self.meta["meta_zero"]:
+                self.meta["meta_zero"]["view_as_float"] = False
+
+        # Check GPU
+        self.cuda(self.device)
+        self.ready = True
+
+        # Set in_features/out_features
+        self.in_features, self.out_features = self.meta["shape"][::-1]
+
+    def quantize(
+        self,
+        W: Tensor,
+        weight_quant_params: dict,
+        scale_quant_params: dict,
+        zero_quant_params: dict,
+    ) -> None:
+        quant_scale = scale_quant_params is not None
+        quant_zero = zero_quant_params is not None
+
+        self.in_features, self.out_features = W.t().shape
+
+        # Quantize
+        W_q, meta = Quantizer.quantize(
+            W,
+            device=self.device,
+            compute_dtype=self.compute_dtype,
+            **weight_quant_params,
+        )
+        meta.update({"quant_scale": quant_scale, "quant_zero": quant_zero})
+
+        if meta["quant_zero"]:
+            meta["zero_q"], meta["meta_zero"] = Quantizer.quantize(
+                meta["zero"],
+                device=self.device,
+                view_as_float=False,
+                **zero_quant_params,
+            )
+            del meta["zero"]
+            meta["meta_zero"]["compute_dtype"] = self.compute_dtype
+
+        if meta["quant_scale"]:
+            meta["scale_q"], meta["meta_scale"] = Quantizer.quantize(
+                meta["scale"],
+                device=self.device,
+                view_as_float=False,
+                **scale_quant_params,
+            )
+            del meta["scale"]
+            meta["meta_scale"]["compute_dtype"] = self.compute_dtype
+
+        self.W_q = W_q
+        self.meta = meta
+        self.cuda(self.device)
+        self.ready = True
+
+    def dequantize(self):
+        assert self.ready, "model was not quantized"
+        W_q, meta = self.W_q, self.meta
+        device = W_q.device
+        del_keys = set()
+
+        # Zero/Scale packed together
+        if "zero_scale" in meta:
+            zero_scale = meta["zero_scale"].to(device=device)
+
+            if zero_scale.dtype == uint8:
+                meta["zero_q"], meta["scale_q"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero_q", "scale_q"})
+            else:
+                meta["zero"], meta["scale"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero", "scale"})
+
+        if meta["quant_zero"]:
+            meta["zero"] = Quantizer.dequantize(
+                meta["zero_q"].to(device=device), meta["meta_zero"]
+            )
+            del_keys.add("zero")
+
+        if meta["quant_scale"]:
+            meta["scale"] = Quantizer.dequantize(
+                meta["scale_q"].to(device=device), meta["meta_scale"]
+            )
+            del_keys.add("scale")
+
+        W_est = Quantizer.dequantize(W_q, meta)
+
+        # Cleanup
+        for key in del_keys:
+            del meta[key]
+        return W_est
+
+    def matmul(self, x: Tensor, transpose: bool = True) -> Tensor:
+        weight = self.dequantize()
+        return torch.matmul(x, weight.t() if (transpose) else weight)
+
+    @torch.compile()
+    def matmul_compile(self, *args, **kwargs):
+        return self.matmul(*args, **kwargs)
+
+    def forward_pytorch_backprop(self, x: Tensor) -> Tensor:
+        return HQQMatmulNoCacheMul.apply(x, self.matmul, self.bias)
+
+    def forward_pytorch_backprop_compile(self, x: Tensor) -> Tensor:
+        return HQQMatmulNoCacheMul.apply(x, self.matmul_compile, self.bias)
+
+    def forward_pytorch(self, x: Tensor) -> Tensor:
+        out = torch.matmul(x, self.dequantize().t())
+        if self.bias is not None:
+            out += self.bias
+        return out
+
+    @torch.compile()
+    def forward_pytorch_compile(self, x: Tensor) -> Tensor:
+        return self.forward_pytorch(x)
+
+    ############################################################################################
+    # ATen C++ / CUDA Bacekdn
+    ##########################################################################################
+    # Requires building the aten backend
+    @torch.jit.ignore
+    def dequantize_Wq_aten(self, W_q: Tensor, meta: dict):
+        if meta["view_as_float"]:
+            W_q = W_q.view(meta["unpack_view_dtype"])
+
+        return hqq_aten.dequantize(
+            W_q,
+            meta["scale"],
+            meta["zero"],
+            meta["shape"],
+            meta["group_size"] if (meta["group_size"]) else -1,
+            meta["nbits"],
+            meta["axis"],
+            meta["packing"],
+        )
+
+    def dequantize_aten(self):
+        # Dequantize
+        assert self.ready, "model was not quantized"
+        assert (
+            self.meta["axis"] == 0
+        ), "only axis=0 is supported. Use HQQLinear.set_backend(HQQBackend.PYTORCH) instead."
+
+        W_q, meta = self.W_q, self.meta
+        device = W_q.device
+        del_keys = set()
+
+        # Zero/Scale packed together
+        if "zero_scale" in meta:
+            zero_scale = meta["zero_scale"].to(device=device, non_blocking=True)
+            if zero_scale.dtype == uint8:
+                meta["zero_q"], meta["scale_q"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero_q", "scale_q"})
+            else:
+                meta["zero"], meta["scale"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero", "scale"})
+
+        # Dequantize zero_q / scale_q with device loading
+        if meta["quant_zero"]:
+            if meta["meta_zero"]["group_size"]:
+                meta["zero"] = self.dequantize_Wq_aten(
+                    meta["zero_q"].to(device=device), meta["meta_zero"]
+                )
+            else:
+                meta["zero"] = Quantizer.dequantize(
+                    meta["zero_q"].to(device=device), meta["meta_zero"]
+                )
+            del_keys.add("zero")
+
+        if meta["quant_scale"]:
+            if meta["meta_scale"]["group_size"]:
+                meta["scale"] = self.dequantize_Wq_aten(
+                    meta["scale_q"].to(device=device), meta["meta_scale"]
+                )
+            else:
+                meta["scale"] = Quantizer.dequantize(
+                    meta["scale_q"].to(device=device), meta["meta_scale"]
+                )
+            del_keys.add("scale")
+
+        # Reconstruct the weights
+        W_est = self.dequantize_Wq_aten(W_q, meta)
+
+        # Cleanup
+        for key in del_keys:
+            del meta[key]
+
+        return W_est
+
+    # Much faster with data-offloading zero_q/scale_q but takes more VRAM
+    def dequantize_aten_with_streams(self):
+        # Create streams
+        if hasattr(self, "stream_zero") is False:
+            self.stream_zero = torch.cuda.Stream()
+            self.stream_scale = torch.cuda.Stream()
+
+        # Dequantize
+        assert self.ready, "model was not quantized"
+        W_q, meta = self.W_q, self.meta
+        device = W_q.device
+        del_keys = set()
+
+        # Zero/Scale packed together
+        if "zero_scale" in meta:
+            zero_scale = meta["zero_scale"].to(device=device, non_blocking=True)
+            if zero_scale.dtype == uint8:
+                meta["zero_q"], meta["scale_q"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero_q", "scale_q"})
+            else:
+                meta["zero"], meta["scale"] = zero_scale[0], zero_scale[1]
+                del_keys.update({"zero", "scale"})
+
+        # Using non_blocking=False for the moment, otherwise it can result in strange behaviors
+        non_blocking = False
+        with torch.cuda.stream(self.stream_zero):
+            if meta["quant_zero"]:
+                if meta["meta_zero"]["group_size"]:
+                    meta["zero"] = self.dequantize_Wq_aten(
+                        meta["zero_q"].to(device=device, non_blocking=non_blocking),
+                        meta["meta_zero"],
+                    )
+                else:
+                    meta["zero"] = Quantizer.dequantize(
+                        meta["zero_q"].to(device=device, non_blocking=non_blocking),
+                        meta["meta_zero"],
+                    )
+                del_keys.add("zero")
+
+        with torch.cuda.stream(self.stream_scale):
+            if meta["quant_scale"]:
+                if meta["meta_scale"]["group_size"]:
+                    meta["scale"] = self.dequantize_Wq_aten(
+                        meta["scale_q"].to(device=device, non_blocking=non_blocking),
+                        meta["meta_scale"],
+                    )
+                else:
+                    meta["scale"] = Quantizer.dequantize(
+                        meta["scale_q"].to(device=device, non_blocking=non_blocking),
+                        meta["meta_scale"],
+                    )
+                del_keys.add("scale")
+
+        # Wait for streams to finish
+        torch.cuda.synchronize()
+
+        # Reconstruct the weights
+        W_est = self.dequantize_Wq_aten(W_q, meta)
+
+        # Cleanup
+        for key in del_keys:
+            del meta[key]
+
+        return W_est
+
+    def forward_aten(self, x: Tensor) -> Tensor:
+        W_est = self.dequantize_aten()
+        out = torch.matmul(x, W_est.t())
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+    def forward_aten_backprop(self, x: Tensor) -> Tensor:
+        return HQQMatmulNoCacheDeq.apply(x, self.dequantize_aten, self.bias)
+
+    # TODO: as fused kernel in CUDA
+    def _get_int8_matrix(self, M):
+        scale = torch.abs(M).amax() / 127.0
+        return torch.round(M / scale).to(torch.int8), scale.float()
+
+    # TODO: in ATEN
+    @torch.compile()
+    def _matmul_int8(self, A, B):
+        dtype = A.dtype
+        A, scale_A = self._get_int8_matrix(A)
+        B, scale_B = self._get_int8_matrix(B)
+        return (torch._int_mm(A, B) * (scale_A * scale_B)).to(dtype)
+
+    def forward_aten_int8(self, x: Tensor) -> Tensor:
+        W_est = self.dequantize_aten()
+        out = self._matmul_int8(x[0], W_est.t())[None, ...]
+        if self.bias is not None:
+            out += self.bias
+
+        return out
+
+    # TODO: fix this hack later for open_clip
+    @property
+    def weight(self):
+        return self.dequantize()
+
+
+def hqq_base_quant_config(
+    nbits: int = 4,
+    group_size: int = 64,
+    quant_zero: bool = True,
+    quant_scale: bool = False,
+    offload_meta: bool = False,  # meta-data should be quantized with the same settings to use offload_meta
+    view_as_float: bool = False,
+    axis: int = 0,
+    mixed: bool = False,  # Auto tune nbits and group_size according to weight distribution
+    budget: float = 0.0,  # overall quantization budget as bits per parameter
+):
+    assert (
+        nbits in Quantizer.SUPPORTED_BITS
+    ), "nbits value not supported. Check Quantizer.SUPPORTED_BITS."
+    if group_size is not None:
+        assert is_divisible(
+            group_size, 8
+        ), "Invalid group_size param: the value should be a multiple of 8."
+    weight_quant_params = {
+        "nbits": nbits,
+        "channel_wise": True,
+        "group_size": group_size,
+        "optimize": True,
+        "round_zero": True if nbits == 4 else False,
+        "axis": axis,
+        "view_as_float": view_as_float,
+    }
+
+    if offload_meta:
+        if quant_scale != quant_zero:
+            # print(colored("quant_zero and quant_scale must be the same when offload_meta is set to True. Setting quant_scale=quant_zero." , 'yellow'))
+            quant_scale = quant_zero
+
+        scale_quant_params = (
+            {"nbits": 8, "channel_wise": True, "group_size": 128, "optimize": False}
+            if (quant_scale)
+            else None
+        )
+        zero_quant_params = (
+            {"nbits": 8, "channel_wise": True, "group_size": 128, "optimize": False}
+            if (quant_zero)
+            else None
+        )
+
+    else:
+        scale_quant_params = (
+            {"nbits": 8, "channel_wise": True, "group_size": 128, "optimize": False}
+            if (quant_scale)
+            else None
+        )
+        zero_quant_params = (
+            {"nbits": 8, "channel_wise": False, "group_size": None, "optimize": False}
+            if (quant_zero)
+            else None
+        )
+
+    return {
+        "weight_quant_params": weight_quant_params,
+        "scale_quant_params": scale_quant_params,
+        "zero_quant_params": zero_quant_params,
+        "offload_meta": offload_meta,
+        "mixed": mixed,
+        "budget": budget,
+    }
+
+
+# Alias: follow similar Auto-GPTQ naming
+BaseQuantizeConfig = hqq_base_quant_config

lm-quant-toolkit/.deps/hqq/hqq/core/utils.py

@@ -0,0 +1,31 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+import torch
+import gc
+import numpy as np
+from typing import Union
+
+
+def cleanup() -> None:
+    try:
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
+    gc.collect()
+
+
+def is_divisible(val1: int, val2: int) -> bool:
+    return int(val2 * np.ceil(val1 / val2)) == val1
+
+
+def zero_pad_row(
+    tensor: torch.Tensor, num_rows: int, dtype: Union[torch.dtype, None] = None
+) -> torch.Tensor:
+    out = torch.zeros(
+        [num_rows, tensor.shape[1]],
+        device=tensor.device,
+        dtype=tensor.dtype if (dtype is None) else dtype,
+    )
+    out[: len(tensor)] = tensor
+
+    return out

lm-quant-toolkit/.deps/hqq/hqq/engine/base.py

@@ -0,0 +1,105 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+from abc import abstractmethod
+from typing import Union
+
+import torch
+from torch import float16
+
+from ..models.base import BaseHQQModel
+
+# Wrapper that makes it easier to add quantization support to different engines (HF, VLLM, etc.)
+
+
+class HQQWrapper:
+    @abstractmethod
+    def _get_arch_key_from_save_dir(cls, save_dir: str):
+        pass
+
+    @classmethod
+    def _get_hqq_class(cls, arg):
+        arch = arg if (isinstance(arg, str)) else arg.arch_key
+        return cls._HQQ_REGISTRY[arch]
+
+    @classmethod
+    def _validate_params(cls, params: dict):
+        pass
+
+    @classmethod
+    def _is_quantizable(cls, model):
+        return hasattr(model, "hqq_quantized")
+
+    @classmethod
+    def _make_quantizable(cls, model, quantized: bool):
+        model.hqq_quantized = quantized
+        model.base_class = cls._get_hqq_class(model)
+
+    @classmethod
+    def _check_arch_support(cls, arg):
+        arch = arg if (isinstance(arg, str)) else getattr(arg, "arch_key", None)
+        assert arch in cls._HQQ_REGISTRY, (
+            "Model architecture " + arch + " not supported yet."
+        )
+
+    @classmethod
+    def _check_if_already_quantized(cls, model):
+        assert not model.hqq_quantized, "Model already quantized"
+
+    @classmethod
+    def _check_if_not_quantized(cls, model):
+        assert model.hqq_quantized, "Model not quantized."
+
+    @classmethod
+    def _set_quantized(cls, model, quantized: bool):
+        model.hqq_quantized = quantized
+
+    #####################################################
+    @classmethod
+    def quantize_model_(
+        cls, model, quant_config, compute_dtype: torch.dtype = float16, device="cuda"
+    ):
+        if cls._is_quantizable(model) is False:
+            cls._make_quantizable(model, quantized=False)
+        cls._check_arch_support(model)
+        cls._check_if_already_quantized(model)
+        cls._get_hqq_class(model).quantize_model(
+            model, quant_config=quant_config, compute_dtype=compute_dtype, device=device
+        )
+        cls._set_quantized(model, True)
+
+    @classmethod
+    def save_quantized_(cls, model, save_dir: str):
+        cls._check_if_not_quantized(model)
+        cls._get_hqq_class(model).save_quantized(model, save_dir=save_dir)
+
+    @classmethod
+    def from_quantized(
+        cls,
+        save_dir_or_hub,
+        compute_dtype: torch.dtype = float16,
+        device="cuda",
+        cache_dir: Union[str, None] = "",
+        adapter: str = None,
+        **kwargs,
+    ):
+        # Both local and hub-support
+        save_dir = BaseHQQModel.try_snapshot_download(
+            save_dir_or_hub, cache_dir=cache_dir
+        )
+        arch_key = cls._get_arch_key_from_save_dir(save_dir)
+        cls._check_arch_support(arch_key)
+
+        model = cls._get_hqq_class(arch_key).from_quantized(
+            save_dir,
+            compute_dtype=compute_dtype,
+            device=device,
+            cache_dir=cache_dir,
+            adapter=adapter,
+        )
+
+        cls._make_quantizable(model, quantized=True)
+        return model
+
+    @classmethod
+    def get_linear_tags(cls, model):
+        return cls._get_hqq_class(model).get_linear_tags()

lm-quant-toolkit/.deps/hqq/hqq/engine/hf.py

@@ -0,0 +1,73 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+
+import transformers
+import torch
+from .base import HQQWrapper
+
+from ..models.hf.llama import LlamaHQQ
+from hqq.models.hf.mixtral import MixtralHQQ
+from hqq.models.hf.phi import PhiHQQ
+from hqq.models.hf.mistral import MistralHQQ
+
+_HQQ_REGISTRY = {}
+_HQQ_REGISTRY["LlamaForCausalLM"] = LlamaHQQ
+_HQQ_REGISTRY["MixtralForCausalLM"] = MixtralHQQ
+_HQQ_REGISTRY["PhiForCausalLM"] = PhiHQQ
+_HQQ_REGISTRY["MistralForCausalLM"] = MistralHQQ
+
+# Alias
+AutoTokenizer = transformers.AutoTokenizer
+
+# Used to call super() on classmethods
+_Parent = transformers.AutoModelForCausalLM
+
+
+class HQQModelForCausalLM(_Parent, HQQWrapper):
+    _HQQ_REGISTRY = _HQQ_REGISTRY
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def _make_quantizable(cls, model, quantized: bool) -> None:
+        model.hqq_quantized = quantized
+        model.arch_key = model.config.architectures[0]
+        model.quantize_model = (
+            lambda quant_config,
+            compute_dtype=torch.float16,
+            device="cuda": cls.quantize_model_(
+                model=model,
+                quant_config=quant_config,
+                compute_dtype=compute_dtype,
+                device=device,
+            )
+        )
+        model.save_quantized = lambda save_dir: cls.save_quantized_(
+            model=model, save_dir=save_dir
+        )
+        model.cuda = lambda *args, **kwargs: model if (quantized) else model.cuda
+        model.to = lambda *args, **kwargs: model if (quantized) else model.to
+        model.float = lambda *args, **kwargs: model if (quantized) else model.float
+        model.half = lambda *args, **kwargs: model if (quantized) else model.half
+        model.base_class = cls._get_hqq_class(model)
+
+    # Force loading the model on CPU and unquantized
+    @classmethod
+    def _validate_params(cls, params: dict) -> None:
+        for p in ["load_in_4bit", "load_in_8bit"]:  # ignore these
+            if p in params:
+                params[p] = False
+        params["device_map"] = None
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        cls._validate_params(kwargs)
+        model = super(_Parent, cls).from_pretrained(*args, **kwargs)
+        cls._make_quantizable(model, quantized=False)
+        return model
+
+    @classmethod
+    def _get_arch_key_from_save_dir(cls, save_dir: str):
+        config = transformers.AutoConfig.from_pretrained(save_dir)
+        return config.architectures[0]

lm-quant-toolkit/.deps/hqq/hqq/engine/open_clip.py

@@ -0,0 +1,112 @@
+import json
+
+import torch
+
+from ..models.base import BaseHQQModel
+from ..models.open_clip.vit_clip import ViTCLIPHQQ
+from .base import HQQWrapper
+
+_HQQ_REGISTRY = {}
+_HQQ_REGISTRY["ViT-B-16"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-plus-240"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-plus"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-quickgelu"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-SigLIP-256"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-SigLIP-384"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-SigLIP-512"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-SigLIP-i18n-256"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-16-SigLIP"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-32-256"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-32"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-32-plus-256"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-B-32-quickgelu"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-bigG-14-CLIPA-336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-bigG-14-CLIPA"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-bigG-14"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-e-14"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-g-14"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-14-378-quickgelu"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-14-CLIPA-336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-14-CLIPA"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-14"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-14-quickgelu"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-H-16"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14-280"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14-336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14-CLIPA-336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14-CLIPA"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-14-quickgelu"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-16-320"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-16"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-16-SigLIP-256"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-L-16-SigLIP-384"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-M-16-alt"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-M-16"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-M-32-alt"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-M-32"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-S-16-alt"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-S-16"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-S-32-alt"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-S-32"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-SO400M-14-SigLIP-384"] = ViTCLIPHQQ
+_HQQ_REGISTRY["ViT-SO400M-14-SigLIP"] = ViTCLIPHQQ
+
+
+class HQQOpenCLIP(HQQWrapper):
+    _HQQ_REGISTRY = _HQQ_REGISTRY
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def _make_quantizable(cls, model, quantized: bool):
+        model.hqq_quantized = quantized
+        model.quantize_model = (
+            lambda quant_config,
+            compute_dtype=torch.float16,
+            device="cuda": cls.quantize_model_(
+                model=model,
+                quant_config=quant_config,
+                compute_dtype=compute_dtype,
+                device=device,
+            )
+        )
+        model.save_quantized = lambda save_dir: cls.save_quantized_(
+            model=model, save_dir=save_dir
+        )
+        # model.cuda = lambda *args, **kwargs: model if (quantized) else model.cuda
+        # model.to = lambda *args, **kwargs: model if (quantized) else model.to
+        # model.float = lambda *args, **kwargs: model if (quantized) else model.float
+        # model.half = lambda *args, **kwargs: model if (quantized) else model.half
+        model.base_class = ViTCLIPHQQ
+
+    @classmethod
+    def _validate_params(cls, params: dict):
+        pass
+
+    @classmethod
+    def create_model(cls, model_id, **kwargs):
+        cls._validate_params(kwargs)
+        # check if the first positional argument is an directory that exists
+        comps = model_id.split("/")
+        elems = comps[1].split("-")
+        model_name = "-".join(elems[1:4])
+        pretrained = "-".join(elems[4:])
+        model = cls._get_hqq_class(model_name).create_model(
+            model_name, model_name=model_name, pretrained=pretrained, **kwargs
+        )
+        cls._make_quantizable(model, quantized=False)
+        return model
+
+    @classmethod
+    def wrap_model(cls, model, model_name):
+        model.arch_key = model_name
+        cls._make_quantizable(model, quantized=False)
+        return model
+
+    @classmethod
+    def _get_arch_key_from_save_dir(cls, save_dir: str):
+        with open(BaseHQQModel.get_config_file(save_dir), "r") as file:
+            config = json.load(file)
+        return config["architecture"]

lm-quant-toolkit/.deps/hqq/hqq/engine/timm.py

@@ -0,0 +1,67 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+import timm
+import json
+import torch
+from .base import HQQWrapper
+
+from ..models.base import BaseHQQModel
+from ..models.timm.vit_clip import ViTCLIPHQQ
+
+_HQQ_REGISTRY = {}
+_HQQ_REGISTRY["vit_huge_patch14_clip_336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_huge_patch14_clip_224"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_large_patch14_clip_224"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_large_patch14_clip_336"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_base_patch16_clip_384"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_base_patch32_clip_448"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_base_patch32_clip_384"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_base_patch16_clip_224"] = ViTCLIPHQQ
+_HQQ_REGISTRY["vit_base_patch32_clip_224"] = ViTCLIPHQQ
+
+
+class HQQtimm(HQQWrapper):
+    _HQQ_REGISTRY = _HQQ_REGISTRY
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def _make_quantizable(cls, model, quantized: bool):
+        model.hqq_quantized = quantized
+        model.arch_key = model.default_cfg["architecture"]
+        model.quantize_model = (
+            lambda quant_config,
+            compute_dtype=torch.float16,
+            device="cuda": cls.quantize_model_(
+                model=model,
+                quant_config=quant_config,
+                compute_dtype=compute_dtype,
+                device=device,
+            )
+        )
+        model.save_quantized = lambda save_dir: cls.save_quantized_(
+            model=model, save_dir=save_dir
+        )
+        model.cuda = lambda *args, **kwargs: model if (quantized) else model.cuda
+        model.to = lambda *args, **kwargs: model if (quantized) else model.to
+        model.float = lambda *args, **kwargs: model if (quantized) else model.float
+        model.half = lambda *args, **kwargs: model if (quantized) else model.half
+        model.base_class = cls._get_hqq_class(model)
+
+    @classmethod
+    def _validate_params(cls, params: dict):
+        pass
+
+    @classmethod
+    def create_model(cls, *args, **kwargs):
+        cls._validate_params(kwargs)
+        model = timm.create_model(*args, **kwargs)
+        cls._make_quantizable(model, quantized=False)
+        return model
+
+    @classmethod
+    def _get_arch_key_from_save_dir(cls, save_dir: str):
+        with open(BaseHQQModel.get_config_file(save_dir), "r") as file:
+            config = json.load(file)
+        return config["architecture"]

lm-quant-toolkit/.deps/hqq/hqq/engine/vllm.py

@@ -0,0 +1,149 @@
+# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2023
+#####################################################
+
+# Import VLLM architectures here with the dummy_load trick
+import torch
+from torch import float16
+import transformers
+from ..models.vllm.llama import LlamaForCausalLM, LlamaHQQ
+from ..models.base import BaseHQQModel
+from .base import HQQWrapper
+from termcolor import colored
+
+# Set them in the model registry
+import vllm
+
+vllm.model_executor.model_loader._MODEL_REGISTRY["LlamaForCausalLM"] = LlamaForCausalLM
+
+_HQQ_REGISTRY = {}
+_HQQ_REGISTRY["LlamaForCausalLM"] = LlamaHQQ
+
+# VLLM requires model input as string, so this is the fallback for a dummy init used to load the model from quantized weights + load the tokenizer
+_ARCH_TO_DEFAULT = {}
+_ARCH_TO_DEFAULT["LlamaForCausalLM"] = "meta-llama/Llama-2-7b-chat-hf"
+
+_Parent = vllm.entrypoints.llm.LLM
+
+
+# Similar to the from vllm import LLM class, with some extra parameters and only loads a dummy model
+class HQQLLM(_Parent, HQQWrapper):
+    _HQQ_REGISTRY = _HQQ_REGISTRY
+    INIT_CACHE_MEM = 1.0  # Default cache mem in GB. This is extra-memory, but you need it otherwise runtime is slower
+
+    def __init__(self, *args, **kwargs):
+        self._validate_params(kwargs)
+
+        if "force_skip" in kwargs:
+            force_skip = kwargs["force_skip"]
+            del kwargs["force_skip"]
+        else:
+            force_skip = False
+
+        # This will init a dummy model
+        super().__init__(*args, **kwargs)
+
+        cache_dir = kwargs["cache_dir"] if ("cache_dir" in kwargs) else ""
+        model_id = kwargs["model"]
+
+        # Here we load the model on CPU so that we can quantize it while avoiding extra GPU cost
+        self.arch_key = self.llm_engine.model_config.hf_config.architectures[0]
+        model_config = self.llm_engine.model_config.hf_config
+        workers = self.llm_engine.workers
+
+        if force_skip is False:
+            for i in range(len(workers)):
+                workers[i].model.__init__(model_config, dummy_load=False)
+                workers[i].model.load_weights(
+                    model_name_or_path=model_id, cache_dir=cache_dir
+                )
+
+        self.hqq_quantized = False
+
+    def _validate_params(self, kwargs: dict):
+        if "gpu_memory_utilization" not in kwargs:
+            total_gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9  # GB
+            kwargs["gpu_memory_utilization"] = HQQLLM.INIT_CACHE_MEM / total_gpu_mem
+
+    # In case the user wants to use the fp16 model and skip quantization
+    def cuda(self):
+        self._check_if_already_quantized(self)
+        workers = self.llm_engine.workers
+        for i in range(len(workers)):
+            workers[i].model = workers[i].model.half().cuda(i)
+        return self
+
+    def quantize_model(
+        self, quant_config: dict, compute_dtype: torch.dtype = float16, device="cuda"
+    ):
+        return self.quantize_model_(
+            model=self,
+            quant_config=quant_config,
+            compute_dtype=compute_dtype,
+            device=device,
+        )
+
+    def save_quantized(self, save_dir: str):
+        return self.save_quantized_(model=self, save_dir=save_dir)
+
+    @classmethod
+    def _get_arch_key_from_save_dir(cls, save_dir: str):
+        config = transformers.AutoConfig.from_pretrained(save_dir)
+        return config.architectures[0]
+
+    # This requires custom loading because VLLM requires a str input model to initalize the workers
+    @classmethod
+    def from_quantized(
+        cls,
+        save_dir_or_hub: str,
+        compute_dtype: torch.dtype = float16,
+        cache_dir: str = "",
+        tensor_parallel_size: int = 1,
+    ):
+        assert tensor_parallel_size == 1, "Only single GPU is supported."
+        # Both local and hub-support
+        save_dir = BaseHQQModel.try_snapshot_download(save_dir_or_hub)
+        arch_key = cls._get_arch_key_from_save_dir(save_dir)
+        cls._check_arch_support(arch_key)
+
+        # Trick to initialize the tokenizer and a dummy model inside a VLLM instance
+        instance = cls(
+            model=_ARCH_TO_DEFAULT[arch_key], force_skip=True, tensor_parallel_size=1
+        )
+        workers = instance.llm_engine.workers
+        for i in range(len(workers)):
+            workers[i].model = cls._get_hqq_class(
+                arch_key
+            ).from_quantized_single_worker(
+                save_dir_or_hub=save_dir_or_hub,
+                compute_dtype=compute_dtype,
+                cache_dir=cache_dir,
+                device="cuda:" + str(i),
+            )
+
+        cls._make_quantizable(instance, quantized=True)
+        return instance
+
+
+# From https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/llms/vllm.py
+
+
+try:
+    from langchain.llms import VLLM as LangchainVLLMBase
+    from langchain_core.pydantic_v1 import root_validator
+
+    class LangchainVLLM(LangchainVLLMBase):
+        def set(self, model):
+            self.client = model
+            return self
+
+        @root_validator()
+        def validate_environment(cls, values: dict) -> dict:
+            return values
+except Exception:
+    LangchainVLLM = None
+    print(
+        colored(
+            'Langchain not installed. You can install it via "pip install langchain"',
+            "cyan",
+        )
+    )