__version__ = "0.0.3"
from .models import AutoQuantForCausalLM, AutoQuantConfig
from .quant.rtn.models import AutoRTNForCausalLM, RTNConfig, BaseRTNForCausalLM
from .quant.awq.models import AutoAWQForCausalLM, AWQConfig, BaseAWQForCausalLM
from .quant.gptq.models import AutoGPTQForCausalLM, GPTQConfig, BaseGPTQForCausalLM
from .modules.qlinear_exllama import exllama_set_max_input_length
from .utils.data import get_calib_dataset
from .utils.lm_eval_adaptor import LMEvalAdaptor
auto_quant/__init__.py
qwopqwop200/AutoQuant
from .base import BaseQuantConfig, BaseQuantForCausalLM
from .auto import AutoQuantConfig, AutoQuantForCausalLM
auto_quant/models/__init__.py
qwopqwop200/AutoQuant
import os
import json
from inspect import signature
from typing import Dict, Optional, Union

import torch
from transformers import AutoConfig
from transformers.utils.hub import cached_file
from .base import BaseQuantConfig, BaseQuantForCausalLM
from auto_quant.quant.rtn.models import AutoRTNForCausalLM, RTNConfig
from auto_quant.quant.awq.models import AutoAWQForCausalLM, AWQConfig
from auto_quant.quant.gptq.models import AutoGPTQForCausalLM, GPTQConfig

QUANT_CONFIG_MAP = {
    "RTN": RTNConfig,
    "AWQ": AWQConfig,
    "GPTQ": GPTQConfig,
}

QUANT_CAUSAL_LM_MODEL_MAP = {
    "RTN": AutoRTNForCausalLM,
    "AWQ": AutoAWQForCausalLM,
    "GPTQ": AutoGPTQForCausalLM,
}

def check_quant_type(quant_type):
    if quant_type not in QUANT_CAUSAL_LM_MODEL_MAP.keys():
        raise TypeError(f"{quant_type} isn't supported yet. Only support: {list(QUANT_CAUSAL_LM_MODEL_MAP.keys())}")
    return quant_type
    
class AutoQuantConfig:
    def __init__(self, quant_type: str, **kwargs):
        quant_type = check_quant_type(quant_type)
        self.quant_config =  QUANT_CONFIG_MAP[quant_type](**kwargs)

    @classmethod
    def from_pretrained(self, save_dir: str, **kwargs):
        # Parameters related to loading from Hugging Face Hub
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", None)
        commit_hash = kwargs.pop("_commit_hash", None)

        quant_config_filename = "quant_config.json"
        if os.path.isdir(save_dir):  # Local
            resolved_config_file = os.path.join(save_dir, quant_config_filename)
        else: # Remote
            resolved_config_file = cached_file(
                save_dir,
                quant_config_filename,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                use_auth_token=use_auth_token,
                revision=revision,
                local_files_only=local_files_only,
                subfolder=subfolder,
                _raise_exceptions_for_missing_entries=False,
                _raise_exceptions_for_connection_errors=False,
                _commit_hash=commit_hash,
            )
        with open(resolved_config_file, "r", encoding="utf-8") as f:
            quant_type = json.loads(f.read())['quant_type']
        quant_type = check_quant_type(quant_type)
        return QUANT_CONFIG_MAP[quant_type].from_pretrained(save_dir, **kwargs)
        
    def __repr__(self):
        return self.quant_config.__repr__()
            
    def __str__(self):
        return self.quant_config.__str__()
        
class AutoQuantForCausalLM:
    def __init__(self):
        raise EnvironmentError(
            "AutoQuantForCausalLM is designed to be instantiated\n"
            "using `AutoQuantForCausalLM.from_pretrained` if want to quantize a pretrained model.\n"
            "using `AutoQuantForCausalLM.from_quantized` if want to inference with quantized model."
        )

    @classmethod
    def from_pretrained(
        self, 
        model_path, 
        quant_config, 
        max_memory: Optional[dict] = None,
        torch_dtype: torch.dtype = torch.float16,
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseQuantForCausalLM:
        if isinstance(quant_config, AutoQuantConfig):
            quant_config = quant_config.quant_config
        quant_type = check_quant_type(quant_config.quant_type)
        return QUANT_CAUSAL_LM_MODEL_MAP[quant_type].from_pretrained(
            model_path, 
            quant_config=quant_config,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            **kwargs
        )

    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseQuantForCausalLM:
        quant_type = AutoQuantConfig.from_pretrained(quant_path).quant_type
        quant_func = QUANT_CAUSAL_LM_MODEL_MAP[quant_type].from_quantized
        # A static list of kwargs needed for huggingface_hub
        huggingface_kwargs = [
            "cache_dir",
            "force_download",
            "proxies",
            "resume_download",
            "local_files_only",
            "use_auth_token",
            "revision",
            "subfolder",
            "_raise_exceptions_for_missing_entries",
            "_commit_hash"
        ]
        keywords = {
            key: kwargs[key]
            for key in list(signature(quant_func).parameters.keys()) + huggingface_kwargs
            if key in kwargs
        }
        return quant_func(
            quant_path,
            device=device,
            device_map=device_map,
            max_memory=max_memory,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype, 
            trust_remote_code=trust_remote_code, 
            **kwargs
        )
auto_quant/models/auto.py
qwopqwop200/AutoQuant
import os
import json
import logging
from typing import Optional, Union
from dataclasses import field, fields

import torch
import torch.nn as nn
from safetensors.torch import save_file as safe_save

from transformers import AutoModelForCausalLM, AutoConfig, PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.utils.hub import PushToHubMixin, cached_file, create_repo, create_commit, CommitOperationAdd
from transformers.modeling_utils import shard_checkpoint, WEIGHTS_NAME, WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, SAFE_WEIGHTS_INDEX_NAME
from accelerate import init_empty_weights, infer_auto_device_map

from auto_quant import __version__
from auto_quant.modules.qlinear_exllama import check_exllama_can_save

class BaseQuantConfig(PushToHubMixin):
    quant_path: Optional[str] = None
    quant_type = None
    
    def save_pretrained(self, save_dir: str, **kwargs):
        with open(os.path.join(save_dir, "quant_config.json"), "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def from_pretrained(self, save_dir: str, **kwargs):
        # Parameters related to loading from Hugging Face Hub
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", None)
        commit_hash = kwargs.pop("_commit_hash", None)

        quant_config_filename = "quant_config.json"
        if os.path.isdir(save_dir):  # Local
            resolved_config_file = os.path.join(save_dir, quant_config_filename)
        else: # Remote
            resolved_config_file = cached_file(
                save_dir,
                quant_config_filename,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                use_auth_token=use_auth_token,
                revision=revision,
                local_files_only=local_files_only,
                subfolder=subfolder,
                _raise_exceptions_for_missing_entries=False,
                _raise_exceptions_for_connection_errors=False,
                _commit_hash=commit_hash,
            )
        
        field_names = [field.name for field in fields(self)]
        with open(resolved_config_file, "r", encoding="utf-8") as f:
            args_from_json = json.load(f)
            if args_from_json['quant_type'] != self.quant_type:
                raise ValueError(f"Expected quant type: {self.quant_type}, quant type in config: {args_from_json['quant_type']}")
            
            filtered_args = {}
            for key, val in args_from_json.items():
                if key in field_names:
                    filtered_args[key] = val
                else:
                    if key not in ['quant_type']:
                        logging.warning(f"ignoring unknown parameter in {quant_config_filename}: {key}.")
            return self(**filtered_args)
            
class BaseQuantForCausalLM(nn.Module, PushToHubMixin, GenerationMixin):
    def __init__(self, model, quant_config, is_quantized):
        super().__init__()
        self.model:PreTrainedModel = model
        self.config = self.model.config
        self.model_type:str = self.model.config.model_type
        self.is_quantized:bool = is_quantized
        if not isinstance(quant_config, BaseQuantConfig):
            raise Exception('Please use QuantConfig, which Inheritance BaseQuantConfig.')
        self.quant_config = quant_config

    @property
    def device(self):
        if not self.hf_device_map:
            return self.model.device
        else:
            device = [d for d in self.hf_device_map.values() if d not in {'cpu', 'disk'}][0]
            return torch.device(device)

    @property
    def hf_device_map(self):
        return getattr(self.model, "hf_device_map", None)
    
    def to(self, device: Union[str, torch.device]):
        self.model.to(device)
        return self

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def generate(self, *args, **kwargs):
        with torch.inference_mode(), torch.amp.autocast(device_type=self.device.type):
            return self.model.generate(*args, **kwargs)

    def prepare_inputs_for_generation(self, *args, **kwargs):
        return self.model.prepare_inputs_for_generation(*args, **kwargs)

    def get_input_embeddings(self, *args, **kwargs):
        return self.model.get_input_embeddings(*args, **kwargs)
            
    def save_pretrained(self, save_dir: str, max_shard_size: Union[int, str] = "10GB", use_safetensors: bool = False):
        """alias of save_quantized"""
        logging.warning("you are using save_pretrained, which will re-direct to save_quantized.")
        self.save_quantized(save_dir, max_shard_size)

    def save_quantized(self, save_dir: str, max_shard_size: Union[int, str] = "10GB", use_safetensors: bool = False):    
        if not self.is_quantized:
            raise EnvironmentError("can only save quantized model, please execute .quantize first.")
        if not check_exllama_can_save(self.model):
            raise EnvironmentError("desc_act is enabled and Exllama is post-initialized. The model cannot be saved.")
        
        os.makedirs(save_dir, exist_ok=True)        
        
        state_dict = self.model.state_dict()
        state_dict = {k: v.to('cpu').clone().contiguous() for k, v in state_dict.items()}
        shards, index = shard_checkpoint(state_dict, max_shard_size, weights_name=SAFE_WEIGHTS_NAME if use_safetensors else WEIGHTS_NAME)
        
        for shard_file, shard in shards.items():
            if use_safetensors:
                safe_save(shard, os.path.join(save_dir, shard_file), {'format': "pt", 'auto_quant_version': str(__version__)})
            else:
                torch.save(shard, os.path.join(save_dir, shard_file))
        if index is not None:
            save_index_file = os.path.join(save_dir, SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME)
            # Save the index as well
            with open(save_index_file, "w", encoding="utf-8") as f:
                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
                f.write(content)
            
        self.model.config.save_pretrained(save_dir)
        self.quant_config.save_pretrained(save_dir)
        self.quant_config.quant_path = save_dir
        
    @classmethod
    def from_pretrained(
        self, 
        model_path, 
        quant_config, 
        max_memory: Optional[dict] = None,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **model_init_kwargs,
    ):  
        from .auto import AutoQuantConfig
        if isinstance(quant_config, AutoQuantConfig):
            quant_config = quant_config.quant_config
        
        self.check_quant_config(quant_config)
        # Parameters related to loading from Hugging Face Hub
        cache_dir = model_init_kwargs.pop("cache_dir", None)
        force_download = model_init_kwargs.pop("force_download", False)
        resume_download = model_init_kwargs.pop("resume_download", False)
        proxies = model_init_kwargs.pop("proxies", None)
        local_files_only = model_init_kwargs.pop("local_files_only", False)
        use_auth_token = model_init_kwargs.pop("use_auth_token", None)
        revision = model_init_kwargs.pop("revision", None)
        subfolder = model_init_kwargs.pop("subfolder", "")
        commit_hash = model_init_kwargs.pop("_commit_hash", None)

        cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "use_auth_token": use_auth_token,
            "revision": revision,
            "subfolder": subfolder,
        }
    
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, **cached_file_kwargs)
        model_init_kwargs["torch_dtype"] = torch_dtype
        model_init_kwargs["trust_remote_code"] = trust_remote_code
        if max_memory:
            if "disk" in max_memory:
                raise NotImplementedError("disk offload not support yet.")
                
            with accelerate.init_empty_weights():
                model = AutoModelForCausalLM.from_config(config, torch_dtype=torch_dtype, trust_remote_code=True)
            model.tie_weights()
            max_memory = get_balanced_memory(
                model,
                max_memory=max_memory,
                no_split_module_classes=[self.layer_type],
                low_zero=False
            )
            model_init_kwargs["device_map"] = infer_auto_device_map(
                model,
                max_memory=max_memory,
                no_split_module_classes=[self.layer_type],
            )
            model_init_kwargs["low_cpu_mem_usage"] = True
            del model
        else:
            model_init_kwargs["device_map"] = None
            model_init_kwargs["low_cpu_mem_usage"] = False
        torch.cuda.empty_cache()
        merged_kwargs = {**model_init_kwargs, **cached_file_kwargs}
        model = AutoModelForCausalLM.from_pretrained(model_path, **merged_kwargs)
        model.eval()
        return self(model, quant_config, is_quantized=False)
        
    def push_to_hub(
        self,
        repo_id: str,
        save_dir: Optional[str] = None,
        max_shard_size: Union[int, str] = "10GB",
        use_safetensors: bool = False,
        commit_message: Optional[str] = "Upload of AutoQuant quantized model",
        use_auth_token: Optional[Union[bool, str]] = None,
        private: Optional[bool] = None,
        token: Optional[Union[bool, str]] = None,
        create_pr: Optional[bool] = False,
    ) -> str:
        """
        Upload the model to the Hugging Face Hub.

        Parameters:
            repo_id (`str`):
                The name of the repository you want to push your tool to. It should contain your organization name when
                pushing to a given organization.
            save_dir (`str`, *optional*):
                The name of the local folder to save the model to.
                If the model has already been saved, this parameter can be omitted.
            max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
                Only applicable for models. The maximum size for a checkpoint before being sharded. 
                Checkpoints shard will then be each of size lower than this size.
            commit_message (`str`, *optional*, defaults to `"Upload of AutoQuant quantized model"`):
                Message to commit while pushing.
            use_auth_token (`bool` or `str`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `huggingface-cli login` (stored in `~/.huggingface`). Will default to `True` if `repo_url`
                is not specified.
            private (`bool`, *optional*):
                Whether or not the repository created should be private.
            token (`bool` or `str`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If unset, will use the token generated
                when running `huggingface-cli login` (stored in `~/.huggingface`).
            create_pr (`bool`, *optional*, defaults to `False`):
                Whether or not to create a PR with the uploaded files or directly commit.
        """
        if (self.quant_config.quant_path is None or not os.path.isdir(self.quant_config.quant_path)) and save_dir is None:
            raise ValueError("Quantized model should be saved first, or you can provide save_dir to make sure model is saved to local disk before uploading.")
        
        if save_dir is not None:
            logging.info(f"Saving model to {save_dir}")
            self.save_quantized(save_dir, max_shard_size)
            
        repo_url = create_repo(
            repo_id=repo_id, token=token, private=private, exist_ok=True, repo_type="model"
        )
        repo_id = repo_url.repo_id

        if self.quant_config.quant_path is not None:
            work_dir = self.quant_config.quant_path
            operations = [
                CommitOperationAdd(path_or_fileobj=os.path.join(work_dir, f), path_in_repo=f)
                for f in os.listdir(work_dir)
            ]
            logging.info(f"Uploading the following files to {repo_id}: {','.join(os.listdir(work_dir))}")
            return create_commit(
                repo_id=repo_id,
                operations=operations,
                commit_message=commit_message,
                token=use_auth_token,
                create_pr=create_pr,
                repo_type="model",
            )
        
    def __getattr__(self, item):
        try:
            return super().__getattr__(item)
        except:
            return getattr(self.model, item)
auto_quant/models/base.py
qwopqwop200/AutoQuant
import gc
import math
import logging
from typing import Optional
import torch
import torch.nn as nn
import transformers
from exllama_kernels import make_q4, q4_matmul, prepare_buffers, set_tuning_params, cleanup_buffers_cuda

EXLLAMA_DEFAULT_MAX_INPUT_LENGTH = 2048

# Dummy tensor to pass instead of g_idx since there is no way to pass "None" to a C++ extension
none_tensor = torch.empty((1, 1), device="meta")

def ext_make_q4(qweight, qzeros, scales, g_idx, device):
    """Construct Q4Matrix, return handle"""
    return make_q4(qweight,
                   qzeros,
                   scales,
                   g_idx if g_idx is not None else none_tensor,
                   device)

def ext_q4_matmul(x, q4, q4_width):
    """Matrix multiplication, returns x @ q4"""
    outshape = x.shape[:-1] + (q4_width,)
    x = x.view(-1, x.shape[-1])
    output = torch.empty((x.shape[0], q4_width), dtype=torch.float16, device=x.device)
    q4_matmul(x, q4, output)
    return output.view(outshape)

class ExllamaLinear(nn.Module):
    def __init__(self, bits, group_size, in_features, out_features, bias, act_order=False, device='cpu'):
        super().__init__()
        if bits not in [4]:
            raise NotImplementedError("Only 4-bit are supported for now.")
        
        self.in_features = in_features
        self.out_features = out_features
        self.bits = bits
        self.group_size = group_size if group_size != -1 else in_features
        # quick sanity check (make sure aligment)
        assert self.in_features % self.group_size == 0
        assert out_features % (32 // self.bits) == 0
        
        self.register_buffer('qweight',torch.zeros((in_features // 32 * bits), out_features, dtype=torch.int32, device=device))
        self.register_buffer('qzeros',torch.zeros((math.ceil(in_features / self.group_size), out_features // 32 * bits), dtype=torch.int32, device=device))
        self.register_buffer('scales',torch.zeros((math.ceil(in_features / self.group_size), out_features), dtype=torch.float16, device=device))
        if bias:
            self.register_buffer('bias', torch.zeros((out_features), dtype=torch.float16, device=device))
        else:
            self.bias = None
        if act_order:
            self.register_buffer('g_idx',torch.tensor([i // self.group_size for i in range(in_features)], dtype=torch.int32))
        else:
            self.g_idx = None
        self.is_post_init = False
        
    def post_init(self):
        assert self.qweight.device.type == "cuda"
        assert self.qweight.device.index is not None
        self.q4 = ext_make_q4(
            self.qweight,
            self.qzeros,
            self.scales,
            self.g_idx.to("cpu") if self.g_idx is not None else None,
            self.qweight.device.index # device index
        )
        self.is_post_init = True
        
    @classmethod
    def from_linear(self, linear, bits, group_size, act_order=False, init_only=False, scales=None, zeros=None, g_idx=None):
        awq_linear = self(bits, group_size, linear.in_features, linear.out_features, linear.bias is not None, act_order, linear.weight.device)
        if init_only:  # just prepare for loading sd
            return awq_linear
        W = linear.weight.data.clone()
        if isinstance(linear, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(linear, transformers.pytorch_utils.Conv1D):
            W = W.t()
        # need scales and zeros info for real quantization
        assert scales is not None and zeros is not None
        scales = scales.t().contiguous()
        zeros = zeros.t().contiguous()
        
        awq_linear.scales = scales.clone().half()
        awq_linear.bias = linear.bias.clone().half()
        awq_linear.g_idx = g_idx.clone() if act_order else None
        
        scale_zeros = zeros * scales
        pack_num = 32 // awq_linear.bits
        
        intweight = []
        for idx in range(awq_linear.in_features):
            if act_order:
                intweight.append(torch.round((W[:, idx] + scale_zeros[g_idx[idx]]) / awq_linear.scales[g_idx[idx]]).to(torch.int)[:, None])
            else:
                intweight.append(torch.round((W[:, idx] + scale_zeros[idx // awq_linear.group_size]) / awq_linear.scales[idx // awq_linear.group_size]).to(torch.int)[:, None])
                
        intweight = torch.cat(intweight, dim=1)
        intweight = intweight.t().contiguous()
        intweight = intweight.to(dtype=torch.int32)
        
        qweight = torch.zeros((intweight.shape[0]  // 32 * awq_linear.bits, intweight.shape[1]), dtype=torch.int32, device=intweight.device)           
        for row in range(intweight.shape[0] // pack_num):
            for i in range(pack_num):
                qweight[row] |= intweight[row * pack_num + i] << (i * awq_linear.bits)
        awq_linear.qweight = qweight

        zeros = zeros.to(dtype=torch.int32)
        qzeros = torch.zeros((zeros.shape[0], zeros.shape[1] // 32 * awq_linear.bits), dtype=torch.int32, device=zeros.device)
        for col in range(zeros.shape[1] // pack_num):     
            for i in range(pack_num):
                qzeros[:, col] |= zeros[:, col * pack_num + i] << (i * awq_linear.bits)
        awq_linear.qzeros = qzeros
        return awq_linear
        
    @torch.no_grad()
    def forward(self, x):
        if not self.is_post_init:
            raise EnvironmentError('Exllama has not been post-initialized.')
        out = ext_q4_matmul(x.half(), self.q4, self.out_features)
        out = out + self.bias if self.bias is not None else out
        return out
    
    def extra_repr(self) -> str:
        return 'in_features={}, out_features={}, bias={}, bits={}, group_size={}'.format(
            self.in_features, self.out_features, self.bias is not None, self.bits, self.group_size
        )
        
def check_exllama_can_save(model):
    for name, submodule in model.named_modules():
        if isinstance(submodule, ExllamaLinear):
            if submodule.is_post_init and submodule.g_idx is not None:
                return False
    return True
        
def check_exllama_device(model):
    for name, submodule in model.named_modules():
        if isinstance(submodule, ExllamaLinear) and submodule.qweight.device.type != "cuda":
                return False
    return True
    
def post_init(model, use_act_order: bool, max_input_length: Optional[int] = None):
    """
    The max_input_length argument is specific to the exllama backend, that requires to initialize a buffer temp_state.
    """
    device_to_buffers_size = {}
    for name, submodule in model.named_modules():
        if isinstance(submodule, ExllamaLinear):
            device = submodule.qweight.device
            if device not in device_to_buffers_size:
                device_to_buffers_size[device] = {"max_dq_buffer_size": 1,"max_inner_outer_dim": 1}

            device_to_buffers_size[device]["max_dq_buffer_size"] = max(device_to_buffers_size[device]["max_dq_buffer_size"], submodule.qweight.numel() * 8)
            if use_act_order:
                device_to_buffers_size[device]["max_inner_outer_dim"] = max(device_to_buffers_size[device]["max_inner_outer_dim"], submodule.in_features, submodule.out_features)
    device_to_buffers = {}

    if use_act_order:
        if max_input_length is None:
            max_input_len = EXLLAMA_DEFAULT_MAX_INPUT_LENGTH
        else:
            max_input_len = max_input_length
    else:
        if max_input_length is not None:
            logging.info("Using exllama backend without act-order, the parameter max_input_length was set although not needed, it will be ignored.")
        max_input_len = 1

    for device, buffers_size in device_to_buffers_size.items():
        # The temp_state buffer is required to reorder X in the act-order case.
        # The temp_dq buffer is required to dequantize weights when using cuBLAS, typically for the prefill.
        device_to_buffers[device] = {
            "temp_state": torch.zeros((max_input_len, buffers_size["max_inner_outer_dim"]), dtype=torch.float16, device=device),
            "temp_dq": torch.zeros((1, buffers_size["max_dq_buffer_size"]), dtype=torch.float16, device=device),
            "max_dq_buffer_size": buffers_size["max_dq_buffer_size"],
            "max_inner_outer_dim": buffers_size["max_inner_outer_dim"],
        }

    # Buffers need to be persistent to avoid any bug.
    model.device_to_buffers = device_to_buffers
    
    for device, buffers in model.device_to_buffers.items():
        prepare_buffers(device, buffers["temp_state"], buffers["temp_dq"])

    # Using the default from exllama repo here.
    matmul_recons_thd = 8
    matmul_fused_remap = False
    matmul_no_half2 = False
    set_tuning_params(matmul_recons_thd, matmul_fused_remap, matmul_no_half2)

    # The buffers need to have been initialized first before calling make_q4.
    for name, submodule in model.named_modules():
        if isinstance(submodule, ExllamaLinear):
            submodule.post_init()

    torch.cuda.empty_cache()
    return model
    
def exllama_post_init(model, use_act_order: bool, max_input_length: Optional[int] = None):
    model_uses_exllama = False
    for name, submodule in model.named_modules():
        if isinstance(submodule, ExllamaLinear):
            model_uses_exllama = True
    if model_uses_exllama:
        if check_exllama_device(model):
            post_init(model, use_act_order, max_input_length)
        else:
            logging.warning("ExLlama cannot be initialized. ExLlama's device only accepts cuda.")
    return model
        
def exllama_set_max_input_length(model, max_input_length: int):
    """
    This method does not necessarily require `model` to inherit from BaseGPTQForCausalLM.

    When using the exllama backend with act-order, it is necessary to initialize a buffer that depends on the maximum expected input length. In case the
    default used (EXLLAMA_DEFAULT_MAX_INPUT_LENGTH) is too short, this method can be called to extend the buffer size without reloading the whole model.
    """
    if not model.quant_config.act_order:
        raise ValueError("The method exllama_set_max_input_length should be called only when using the exllama backend **with act-order**.")
    
    device_to_buffers_size = {}
    for device, buffers in model.device_to_buffers.items():
        device_to_buffers_size[device] = {"max_dq_buffer_size": buffers["max_dq_buffer_size"], "max_inner_outer_dim": buffers["max_inner_outer_dim"]}
    
    # For an unknown reason calling just `del model.device_to_buffers` raises an AttributeError.
    for key in list(model.device_to_buffers.keys()):
        del model.device_to_buffers[key]
    model.device_to_buffers = None
    del model.device_to_buffers

    gc.collect()
    torch.cuda.empty_cache()
    # cleanup_buffers_cuda()

    device_to_buffers = {}
    for device, buffers_size in device_to_buffers_size.items():
        # The temp_state buffer is required to reorder X in the act-order case.
        # The temp_dq buffer is required to dequantize weights when using cuBLAS, typically for the prefill.
        device_to_buffers[device] = {
            "temp_state": torch.zeros((max_input_length, buffers_size["max_inner_outer_dim"]), dtype=torch.float16, device=device),
            "temp_dq": torch.zeros((1, buffers_size["max_dq_buffer_size"]), dtype=torch.float16, device=device),
            "max_dq_buffer_size": buffers_size["max_dq_buffer_size"],
            "max_inner_outer_dim": buffers_size["max_inner_outer_dim"],
        }

        prepare_buffers(device, device_to_buffers[device]["temp_state"], device_to_buffers[device]["temp_dq"])

    # Buffers need to be persistent to avoid any bug.
    model.device_to_buffers = device_to_buffers

    return model
    
def make_sure_no_tensor_in_meta_device(model):
    for n, m in model.named_modules():
        if isinstance(submodule, ExllamaLinear) and m.bias.device == torch.device("meta"):
            m.register_buffer('bias', torch.zeros((m.out_features), dtype=torch.float16, device="cpu"))
auto_quant/modules/qlinear_exllama.py
qwopqwop200/AutoQuant
from .base import AWQConfig, BaseAWQForCausalLM
from .mpt import MptAWQForCausalLM
from .llama import LlamaAWQForCausalLM
from .opt import OPTAWQForCausalLM
from .falcon import FalconAWQForCausalLM
from .bloom import BloomAWQForCausalLM
from .gptj import GPTJAWQForCausalLM
from .auto import *
auto_quant/quant/awq/models/__init__.py
qwopqwop200/AutoQuant
from inspect import signature
from typing import Dict, Optional, Union

import torch
from transformers import AutoConfig
from auto_quant.quant.awq.models import *

AWQ_CAUSAL_LM_MODEL_MAP = {
    "mpt": MptAWQForCausalLM,
    "llama": LlamaAWQForCausalLM,
    "opt": OPTAWQForCausalLM,
    "RefinedWeb": FalconAWQForCausalLM,
    "RefinedWebModel": FalconAWQForCausalLM,
    "bloom": BloomAWQForCausalLM,
    "gptj": GPTJAWQForCausalLM
}

def check_and_get_model_type(model_dir, trust_remote_code=True):
    config = AutoConfig.from_pretrained(model_dir, trust_remote_code=trust_remote_code)
    if config.model_type not in AWQ_CAUSAL_LM_MODEL_MAP.keys():
        raise TypeError(f"{config.model_type} isn't supported yet. Only support: {list(AWQ_CAUSAL_LM_MODEL_MAP.keys())}")
    model_type = config.model_type
    return model_type

class AutoAWQForCausalLM:
    def __init__(self):
        raise EnvironmentError('You must instantiate AutoAWQForCausalLM with\n'
                               'AutoAWQForCausalLM.from_quantized or AutoAWQForCausalLM.from_pretrained')
    
    @classmethod
    def from_pretrained(
        self, 
        model_path, 
        quant_config: AWQConfig, 
        max_memory: Optional[dict] = None,
        torch_dtype: torch.dtype = torch.float16,
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseAWQForCausalLM:
        
        model_type = check_and_get_model_type(model_path, trust_remote_code)
        return AWQ_CAUSAL_LM_MODEL_MAP[model_type].from_pretrained(
            model_path, 
            quant_config=quant_config,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            **kwargs
        )
        
    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseAWQForCausalLM:
        model_type = check_and_get_model_type(quant_path, trust_remote_code)
        quant_func = AWQ_CAUSAL_LM_MODEL_MAP[model_type].from_quantized
        # A static list of kwargs needed for huggingface_hub
        huggingface_kwargs = [
            "cache_dir",
            "force_download",
            "proxies",
            "resume_download",
            "local_files_only",
            "use_auth_token",
            "revision",
            "subfolder",
            "_raise_exceptions_for_missing_entries",
            "_commit_hash"
        ]
        keywords = {
            key: kwargs[key]
            for key in list(signature(quant_func).parameters.keys()) + huggingface_kwargs
            if key in kwargs
        }
        return quant_func(
            quant_path,
            device=device,
            device_map=device_map,
            max_memory=max_memory,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype, 
            trust_remote_code=trust_remote_code, 
            **kwargs
        )

auto_quant/quant/awq/models/auto.py
qwopqwop200/AutoQuant
import os
import gc
import functools
import logging
from tqdm import tqdm
from collections import defaultdict
from dataclasses import dataclass, field, fields
from typing import Dict, List, Optional, Union

import torch
import torch.nn as nn
from torch import LongTensor
from torch.nn.utils.rnn import pad_sequence

from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoConfig
from transformers.modeling_utils import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_in_model
from accelerate.utils import get_balanced_memory

from auto_quant.models.base import BaseQuantConfig, BaseQuantForCausalLM
from auto_quant.modules.qlinear_exllama import ExllamaLinear, exllama_post_init, make_sure_no_tensor_in_meta_device
from auto_quant.utils.module import append_str_prefix, get_op_name, get_named_module, set_op_by_name, simple_dispatch_model, get_device
from ..quantize.quantizer import pseudo_quantize_tensor, ScaledActivation
from ..quantize.auto_clip import auto_clip_block, apply_clip
from ..quantize.auto_scale import auto_scale_block, apply_scale
@dataclass
class AWQConfig(BaseQuantConfig):
    bits: int = field(default=4, metadata={"choices": [4]})
    group_size: int = field(default=-1)
    zero_point: bool = field(default=True)
    auto_scale: bool = field(default=True)
    mse_range: bool = field(default=True)
    inplace: bool = field(default=False)
    quant_type = 'AWQ'
    
    def __post_init__(self):
        fields_info = fields(self)

        if self.bits not in fields_info[0].metadata["choices"]:
            raise ValueError(f"only support quantize to {fields_info[0].metadata['choices']} bits.")
        if self.group_size != -1 and self.group_size <= 0:
            raise ValueError("unless equal to -1, group_size must greater then 0.")
        if not self.zero_point:
            raise NotImplementedError("We only support zero_point quantization now.")
                
    def to_dict(self):
        return {
            "bits": self.bits,
            "group_size": self.group_size,
            "zero_point": self.zero_point,
            "auto_scale": self.auto_scale,
            "mse_range": self.mse_range,
            "inplace": self.inplace,
            "quant_type": self.quant_type,
        }

class BaseAWQForCausalLM(BaseQuantForCausalLM):
    def __init__(self, model, quant_config, is_quantized):
        super().__init__(model, quant_config, is_quantized)
        self.quant_config:AWQConfig = quant_config
        self.search_result = None
                
    @classmethod
    def check_quant_config(self, quant_config: AWQConfig):
        pass
    
    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs,
    ):
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", "")
        commit_hash = kwargs.pop("_commit_hash", None)

        snapshot_cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "revision": revision,
        }

        cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "use_auth_token": use_auth_token,
            "revision": revision,
            "subfolder": subfolder,
            "_raise_exceptions_for_missing_entries": False,
            "_commit_hash": commit_hash,
        }
    
        if not os.path.isdir(quant_path):
            quant_path = snapshot_download(quant_path,ignore_patterns=[ "*.msgpack" , "*.h5"], **snapshot_cached_file_kwargs)
            
        config = AutoConfig.from_pretrained(quant_path, trust_remote_code=trust_remote_code, **cached_file_kwargs)
        quant_config = AWQConfig.from_pretrained(quant_path, **cached_file_kwargs)

        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code)
        
        self._load_quantized_modules(self, model, quant_config)
        
        model.tie_weights()
        
        if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
            raise ValueError(
                "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
                "'sequential'."
            )
        if isinstance(device_map, dict):
            max_memory = None
        else:
            if device is None and not device_map and not max_memory:
                device_map = "auto"
            if device is not None:
                device = torch.device(device)
                if not max_memory and not device_map:
                    device_map = {"": device.index if device.type == "cuda" else device.type}
            if not isinstance(device_map, dict) and device_map != "sequential":
                max_memory = get_balanced_memory(
                    model=model,
                    max_memory=max_memory,
                    no_split_module_classes=[self.layer_type],
                    low_zero=(device_map == "balanced_low_0")
                )
        if not isinstance(device_map, dict):
            device_map = infer_auto_device_map(
                model,
                max_memory=max_memory,
                no_split_module_classes=[self.layer_type]
            )

        if low_cpu_mem_usage:
            make_sure_no_tensor_in_meta_device(model)
        if os.path.isfile(os.path.join(quant_path, WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, WEIGHTS_NAME)
        elif os.path.isfile(os.path.join(quant_path, SAFE_WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, SAFE_WEIGHTS_NAME)
        load_checkpoint_in_model(
            model,
            checkpoint=quant_path,
            device_map=device_map,
            offload_state_dict=True,
            offload_buffers=True
        )
        model = simple_dispatch_model(model, device_map)
        model = exllama_post_init(model, False)
        return self(model, quant_config, is_quantized=True)

    def _load_quantized_modules(self, model, quant_config):
        # Get blocks of model
        layers = self.get_model_layers(model)

        for i in tqdm(range(len(layers)), desc="Replacing layers..."):
            layer = layers[i]

            # Get every linear layer in a block
            named_linears = get_named_module(layer)

            # Replace activation functions
            self._scale_activations(self, layer)

            # Replace nn.Linear with ExllamaLinear
            for name, module in named_linears.items():
                q_linear = ExllamaLinear.from_linear(module, quant_config.bits, quant_config.group_size, False, True)
                q_linear.to(get_device(layer))
                set_op_by_name(layer, name, q_linear)
            torch.cuda.empty_cache()
            gc.collect()
            
    def _prepare_examples_for_quantization(self, examples: List[Union[List[int], torch.LongTensor]]):
        pad_token_id = self.model.config.pad_token_id
        if not pad_token_id:
            pad_token_id = self.config.eos_token_id
    
        if isinstance(examples,list):
            new_examples = []
            for tensor in examples:
                if tensor.dim() > 2:
                    raise Exception('examples must be 2D tensor or less.')
                else:
                    for i in range(tensor.dim(),2):
                        tensor.squeeze(0)
                        
                for i in range(tensor.shape[0]):
                    new_examples.append(tensor[i])
            examples = pad_sequence(new_examples, True, padding_value=pad_token_id)
            
        if examples.dim() > 2:
            raise Exception('examples must be 2D tensor or less.')
        else:
            for i in range(examples.dim(),2):
                examples.squeeze(0)
        return examples
            
    @torch.no_grad()
    def quantize(self, examples: List[Union[List[int], torch.LongTensor]]):
        examples = self._prepare_examples_for_quantization(examples)
        self._awq_search(examples)
        self._awq_quant()
        logging.warning("Exllama has not been post-initialized. inference doesn't work. Please save the model and then load it.")
        self.is_quantized = True
        
    def _awq_search(self, examples):
        device_map = self.hf_device_map
        layers = self.get_model_layers(self.model)
        inps = []
        layer_kwargs = {}

        layers[0] = layers[0].cuda()
        self.move_embed(self.model, "cuda")
        
        # get input and kwargs to layer 0
        # with_kwargs is only supported in PyTorch 2.0
        # use this Catcher hack for now
        class Catcher(nn.Module):
            def __init__(self, module):
                super().__init__()
                self.module = module

            def forward(self, hijacked_inputs, **kwargs):
                inps.append(hijacked_inputs)
                layer_kwargs.update(kwargs)
                raise ValueError  # early exit to break later inference

        # patch layer 0 to catch input and kwargs
        layers[0] = Catcher(layers[0])
        try:
            self.model(examples.to(get_device(self.model)))
        except ValueError:  # work with early exit
            pass
        del examples
        layers[0] = layers[0].module  # restore
        inps = inps[0]

        layers[0] = layers[0].cpu()
        self.move_embed(self.model, "cpu")
        
        gc.collect()
        torch.cuda.empty_cache()
        awq_results = {"scale": [],"clip": [],}

        # Run AWQ search layer by layer
        for i in tqdm(range(len(layers)), desc="AWQ Search"):
            layer = layers[i].cuda()
            named_linears = get_named_module(layer)

            # firstly, get input features of all linear layers
            def cache_input_hook(m, x, y, name, feat_dict):
                x = x[0]
                x = x.detach().cpu()
                feat_dict[name].append(x)

            input_feat = defaultdict(list)
            handles = []
            for name in named_linears:
                handles.append(named_linears[name].register_forward_hook(functools.partial(cache_input_hook, name=name,feat_dict=input_feat)))
            inps = inps.to(get_device(layer))  # in case multi-gpu
            # get output as next layer's input
            inps = layer(inps, **layer_kwargs)[0]
            for h in handles:
                h.remove()
            # now solve for scaling and clipping
            input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()}

            # Clear GPU memory
            torch.cuda.empty_cache()

            if self.quant_config.auto_scale:  # if it applies, we should also modify the input_feat with scales
                scales_list = auto_scale_block(
                    self,
                    layer,
                    layer_kwargs,
                    quant_config=self.quant_config,
                    input_feat=input_feat,
                )

                apply_scale(layers[i], scales_list, input_feat_dict=input_feat)

                # append prefix to make names global
                awq_results["scale"] += append_str_prefix(scales_list, get_op_name(self.model, layer) + ".")

            # Clear GPU memory
            torch.cuda.empty_cache()
            
            if self.quant_config.mse_range:
                clip_list = auto_clip_block(layer,quant_config=self.quant_config,input_feat=input_feat)
                apply_clip(layer, clip_list)
                # append prefix to make names global
                awq_results["clip"] += append_str_prefix(clip_list, get_op_name(self.model, layer) + ".")

            layer = layer.cpu()
            # Haotian: check activation replacement
            del input_feat
            gc.collect()
            torch.cuda.empty_cache()
        
        self.search_result = awq_results
        if device_map:
            self.model = simple_dispatch_model(self.model, device_map)

    def _awq_quant(self):
        layers = self.get_model_layers(self.model)

        # Run AWQ quantization
        for i in tqdm(range(len(layers)), desc="AWQ Quantization"):
            layer = layers[i]
            named_linears = get_named_module(layer)
            self._scale_activations(self, layer)

            for name, module in named_linears.items():
                device = get_device(module)
                module.cuda()
                
                module.weight.data, scales, zeros = pseudo_quantize_tensor(
                    module.weight.data, 
                    self.quant_config.bits, 
                    self.quant_config.group_size,
                    self.quant_config.zero_point, 
                    self.quant_config.inplace,
                    get_scale_zp=True, 
                )
                q_linear = ExllamaLinear.from_linear(
                    module, 
                    self.quant_config.bits, 
                    self.quant_config.group_size,
                    False,
                    False, 
                    scales, 
                    zeros
                )

                module.cpu()
                q_linear.to(device)
                set_op_by_name(layer, name, q_linear)
                torch.cuda.empty_cache()
                gc.collect()
            
    @staticmethod
    def _scale_activations(self, layer):
        scale_dict = self.get_act_for_scaling(layer)
        if scale_dict['is_scalable']:
            if not isinstance(scale_dict['scale_layer'], ScaledActivation):
                device = get_device(layer)
                
                # get activation scale
                scale_like = torch.ones(scale_dict['scale_shape'], dtype=param.dtype, device=device)
                
                # scale activation
                scaled_act = ScaledActivation(scale_dict['scale_layer'], scale_like)
                set_op_by_name(layer, scale_dict['scale_name'], scaled_act)
auto_quant/quant/awq/models/base.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.bloom.modeling_bloom import BloomForCausalLM, BloomBlock

class BloomAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "BloomBlock"
    @staticmethod
    def get_model_layers(model: BloomForCausalLM):
        return model.transformer.h
    
    @staticmethod
    def get_act_for_scaling(module: BloomBlock):
        return dict(
            is_scalable=True,
            scale_name="mlp.gelu_impl",
            scale_layer=module.mlp.gelu_impl,
            scale_shape=module.mlp.dense_h_to_4h.out_features
        )
    
    @staticmethod
    def move_embed(model: BloomForCausalLM, device: str):
        model.transformer.word_embeddings = model.transformer.word_embeddings.to(device)
        model.transformer.word_embeddings_layernorm = model.transformer.word_embeddings_layernorm.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: BloomBlock, input_feat, module_kwargs):
        layers = []

        # attention input
        layers.append(dict(
            prev_op=module.input_layernorm,
            layers=[module.self_attention.query_key_value],
            inp=input_feat['self_attention.query_key_value'],
            module2inspect=module, kwargs=module_kwargs,
        ))
        # attention out
        # Please refer to https://github.com/mit-han-lab/llm-awq/issues/2#issuecomment-1606297469
        """
        scales_list.append(_auto_get_scale(
            prev_op=module.self_attention.query_key_value,
            layers=[module.self_attention.dense],
            inp=input_feat['self_attention.dense'],
        ))
        """
        # linear 1
        layers.append(dict(
            prev_op=module.post_attention_layernorm,
            layers=[module.mlp.dense_h_to_4h],
            inp=input_feat['mlp.dense_h_to_4h'],
            module2inspect=module, kwargs=module_kwargs,
        ))
        # linear 2
        layers.append(dict(
            prev_op=module.mlp.gelu_impl,
            layers=[module.mlp.dense_4h_to_h],
            inp=input_feat['mlp.dense_4h_to_h'],
        ))

        return layers
auto_quant/quant/awq/models/bloom.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.falcon.modeling_falcon import FalconDecoderLayer, FalconForCausalLM

class FalconAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "FalconDecoderLayer"
    @staticmethod
    def get_model_layers(model: FalconForCausalLM):
        return model.transformer.h
    
    @staticmethod
    def get_act_for_scaling(module: FalconDecoderLayer):
        return dict(
            is_scalable=True,
            scale_name="mlp.act",
            scale_layer=module.mlp.act,
            scale_shape=module.mlp.dense_h_to_4h.out_features
        )
    
    @staticmethod
    def move_embed(model: FalconForCausalLM, device):
        model.transformer.word_embeddings = model.transformer.word_embeddings.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: FalconDecoderLayer, input_feat, module_kwargs):
        layers = []
        
        # Falcon 7B (older architecture)
        if module.config.num_attention_heads == 71:
            # linear 1 + attention
            layers.append(dict(
                prev_op=module.input_layernorm,
                layers=[module.mlp.dense_h_to_4h, module.self_attention.query_key_value],
                inp=input_feat['self_attention.query_key_value'],
                module2inspect=module,
                kwargs=module_kwargs,
            ))

        # Falcon 40B (newer architecture)
        else:
            # linear 1 + attention
            layers.append(dict(
                prev_op=module.ln_attn,
                layers=[module.self_attention.query_key_value],
                inp=input_feat['self_attention.query_key_value'],
                module2inspect=module,
                kwargs=module_kwargs,
            ))

            # linear 2
            layers.append(dict(
                prev_op=module.ln_mlp,
                layers=[module.mlp.dense_h_to_4h],
                inp=input_feat['mlp.dense_h_to_4h'],
                module2inspect=module,
                kwargs=module_kwargs,
            ))

        return layers
auto_quant/quant/awq/models/falcon.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.gptj.modeling_gptj import GPTJForCausalLM, GPTJBlock

class GPTJAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "GPTJBlock"
    @staticmethod
    def get_model_layers(model: GPTJForCausalLM):
        return model.transformer.h
    
    @staticmethod
    def get_act_for_scaling(module: GPTJBlock):
        return dict(
            is_scalable=True,
            scale_name="mlp.act",
            scale_layer=module.mlp.act,
            scale_shape=module.mlp.fc_in.out_features
        )
    
    @staticmethod
    def move_embed(model: GPTJForCausalLM, device: str):
        model.transformer.wte = model.transformer.wte.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: GPTJBlock, input_feat, module_kwargs):
        layers = []

        # attention input + linear 1
        layers.append(dict(
            prev_op=module.ln_1,
            layers=[module.attn.q_proj,
                    module.attn.k_proj, module.attn.v_proj, module.mlp.fc_in],
            inp=input_feat['attn.q_proj'],
            module2inspect=module,
            kwargs=module_kwargs
        ))

        # attention out
        layers.append(dict(
            prev_op=module.attn.v_proj,
            layers=[module.attn.out_proj],
            inp=input_feat['attn.out_proj'],
        ))

        # linear 2
        layers.append(dict(
            prev_op=module.mlp.act,
            layers=[module.mlp.fc_out],
            inp=input_feat['mlp.fc_out'],
        ))

        return layers
auto_quant/quant/awq/models/gptj.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaForCausalLM

class LlamaAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "LlamaDecoderLayer"    
    @staticmethod
    def get_model_layers(model: LlamaForCausalLM):
        return model.model.layers
    
    @staticmethod
    def get_act_for_scaling(module: LlamaDecoderLayer):
        return dict(
            is_scalable=False
        )
    
    @staticmethod
    def move_embed(model: LlamaForCausalLM, device: str):
        model.model.embed_tokens = model.model.embed_tokens.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: LlamaDecoderLayer, input_feat, module_kwargs):
        layers = []

        # attention input
        layers.append(dict(
            prev_op=module.input_layernorm,
            layers=[module.self_attn.q_proj,
                    module.self_attn.k_proj, module.self_attn.v_proj],
            inp=input_feat['self_attn.q_proj'],
            module2inspect=module.self_attn, kwargs=module_kwargs,
        ))

        # attention out
        # Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696
        if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape:
            layers.append(dict(
                prev_op=module.self_attn.v_proj,
                layers=[module.self_attn.o_proj],
                inp=input_feat['self_attn.o_proj'],
            ))
        
        # linear 1
        layers.append(dict(
            prev_op=module.post_attention_layernorm,
            layers=[module.mlp.gate_proj, module.mlp.up_proj],
            inp=input_feat['mlp.gate_proj'],
            module2inspect=module.mlp,
        ))

        # linear 2
        layers.append(dict(
            prev_op=module.mlp.up_proj,
            layers=[module.mlp.down_proj],
            inp=input_feat['mlp.down_proj'],
        ))

        return layers
auto_quant/quant/awq/models/llama.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.mpt.modeling_mpt import MptBlock, MptForCausalLM, MptMLP

class MptAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "MPTBlock"
    @staticmethod
    def get_model_layers(model: MptForCausalLM):
        return model.transformer.blocks
    
    @staticmethod
    def get_act_for_scaling(module: MptBlock):
        return dict(
            is_scalable=True,
            scale_name="ffn.act",
            scale_layer=module.ffn.act,
            scale_shape=module.ffn.up_proj.out_features
        )
    
    @staticmethod
    def move_embed(model: MptForCausalLM, device: str):
        model.transformer.wte = model.transformer.wte.to(device)
        model.transformer.emb_drop = model.transformer.emb_drop.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: MptBlock, input_feat, module_kwargs):
        layers = []

        # attention input
        layers.append(dict(
            prev_op=module.norm_1,
            layers=[module.attn.Wqkv],
            inp=input_feat['attn.Wqkv'],
            module2inspect=module.attn,
            kwargs=module_kwargs
        ))

        # attention output
        layers.append(dict(
            prev_op=module.attn.Wqkv,
            layers=[module.attn.out_proj],
            inp=input_feat['attn.out_proj']
        ))

        # linear 1
        layers.append(dict(
            prev_op=module.norm_2,
            layers=[module.ffn.up_proj],
            inp=input_feat['ffn.up_proj'],
            module2inspect=module.ffn
        ))

        # linear 2
        layers.append(dict(
            prev_op=module.ffn.act,
            layers=[module.ffn.down_proj],
            inp=input_feat['ffn.down_proj']
        ))

        return layers
auto_quant/quant/awq/models/mpt.py
qwopqwop200/AutoQuant
from .base import BaseAWQForCausalLM
from transformers.models.opt.modeling_opt import OPTForCausalLM, OPTDecoderLayer

class OPTAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "OPTDecoderLayer"
    @staticmethod
    def get_model_layers(model: OPTForCausalLM):
        return model.model.decoder.layers
    
    @staticmethod
    def get_act_for_scaling(module: OPTDecoderLayer):
        return dict(
            is_scalable=False
        )
    
    @staticmethod
    def move_embed(model: OPTForCausalLM, device: str):
        model.model.decoder.embed_tokens = model.model.decoder.embed_tokens.to(device)
        model.model.decoder.embed_positions = model.model.decoder.embed_positions.to(device)
    
    @staticmethod
    def get_layers_for_scaling(module: OPTDecoderLayer, input_feat, module_kwargs):
        layers = []

        # attention input
        layers.append(dict(
            prev_op=module.self_attn_layer_norm,
            layers=[module.self_attn.q_proj,
                    module.self_attn.k_proj, module.self_attn.v_proj],
            inp=input_feat['self_attn.q_proj'],
            module2inspect=module.self_attn, kwargs=module_kwargs,
        ))

        # attention out
        layers.append(dict(
            prev_op=module.self_attn.v_proj,
            layers=[module.self_attn.out_proj],
            inp=input_feat['self_attn.out_proj'],
        ))

        # linear 1
        layers.append(dict(
            prev_op=module.final_layer_norm,
            layers=[module.fc1],
            inp=input_feat['fc1'],
        ))

        # linear 2
        layers.append(dict(
            prev_op=module.fc1,
            layers=[module.fc2],
            inp=input_feat['fc2'],
        ))

        return layers
auto_quant/quant/awq/models/opt.py
qwopqwop200/AutoQuant
from .quantizer import pseudo_quantize_tensor, ScaledActivation
from .auto_clip import auto_clip_block, apply_clip
from .auto_scale import auto_scale_block, apply_scale
auto_quant/quant/awq/quantize/__init__.py
qwopqwop200/AutoQuant
import torch
import torch.nn as nn
from .quantizer import pseudo_quantize_tensor
import gc

# weight quantization
@torch.no_grad()
def auto_clip_layer(w, 
                    input_feat, 
                    quant_config,
                    n_grid=20,
                    max_shrink=0.5,
                    n_sample_token=512):
    assert w.dim() == 2
    org_w_shape = w.shape
    # w           [co, ci]      -> [co, 1, n_group, group size]
    # input_feat  [n_token, ci] -> [1, n_token, n_group, group size]
    group_size = quant_config.group_size if quant_config.group_size > 0 else w.shape[1]
    input_feat = input_feat.view(-1, input_feat.shape[-1])
    input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size)
    input_feat = input_feat[:, 0::input_feat.shape[1] // n_sample_token]
    w = w.reshape(w.shape[0], 1, -1, group_size)

    oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64  # prevent OOM
    assert w.shape[0] % oc_batch_size == 0
    w_all = w
    best_max_val_all = []

    for i_b in range(w.shape[0] // oc_batch_size):
        w = w_all[i_b * oc_batch_size: (i_b + 1) * oc_batch_size]

        org_max_val = w.abs().amax(dim=-1, keepdim=True)  # co, 1, n_group, 1

        best_max_val = org_max_val.clone()
        min_errs = torch.ones_like(org_max_val) * 1e9
        input_feat = input_feat.to(w.device)
        org_out = (input_feat * w).sum(dim=-1)  # co, n_token, n_group

        for i_s in range(int(max_shrink * n_grid)):
            max_val = org_max_val * (1 - i_s / n_grid)
            min_val = - max_val
            cur_w = torch.clamp(w, min_val, max_val)
            q_w = pseudo_quantize_tensor(cur_w, quant_config.bits, quant_config.group_size,
                                         quant_config.zero_point, quant_config.inplace)
            cur_out = (input_feat * q_w).sum(dim=-1)

            # co, 1, n_group, 1
            err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape)
            del cur_w
            del cur_out
            cur_best_idx = err < min_errs
            min_errs[cur_best_idx] = err[cur_best_idx]
            best_max_val[cur_best_idx] = max_val[cur_best_idx]
        best_max_val_all.append(best_max_val)

    best_max_val = torch.cat(best_max_val_all, dim=0)

    del input_feat
    del org_out
    gc.collect()
    torch.cuda.empty_cache()
    return best_max_val.squeeze(1)


@torch.no_grad()
def auto_clip_block(module,
                    quant_config,
                    input_feat):

    named_linears = {name: m for name,
                     m in module.named_modules() if isinstance(m, nn.Linear)}

    clip_list = []
    for name in named_linears:
        # due to qk bmm, it is hard to clip precisely
        if any([_ in name for _ in ["q_", "k_", "query", "key", "Wqkv"]]):
            continue
        named_linears[name].cuda()
        max_val = auto_clip_layer(named_linears[name].weight, input_feat[name], quant_config=quant_config)
        clip_list.append((name, max_val))
        named_linears[name].cpu()
    return clip_list


@torch.no_grad()
def apply_clip(module, clip_list):
    from auto_quant.utils.module import get_op_by_name
    for name, max_val in clip_list:
        layer = get_op_by_name(module, name)
        layer.cuda()
        max_val = max_val.to(layer.weight.device)
        org_shape = layer.weight.shape
        layer.weight.data = layer.weight.data.reshape(*max_val.shape[:2], -1)
        layer.weight.data = torch.clamp(layer.weight.data, -max_val, max_val)
        layer.weight.data = layer.weight.data.reshape(org_shape)
        layer.cpu()

auto_quant/quant/awq/quantize/auto_clip.py
qwopqwop200/AutoQuant
import gc
import torch
import torch.nn as nn
import logging

from transformers.models.bloom.modeling_bloom import BloomGelu
from transformers.models.llama.modeling_llama import LlamaRMSNorm
from transformers.activations import NewGELUActivation
from auto_quant.utils.module import get_op_by_name, get_op_name, set_op_by_name
from .quantizer import ScaledActivation
@torch.no_grad()
def get_weight_scale(weight, group_size=-1):
    org_shape = weight.shape
    if group_size > 0:
        weight = weight.view(-1, group_size)
    scale = weight.abs() / weight.abs().amax(dim=1, keepdim=True)
    scale = scale.view(org_shape)
    scale = scale.mean(0)
    return scale


@torch.no_grad()
def get_act_scale(x):
    return x.abs().view(-1, x.shape[-1]).mean(0)


@torch.no_grad()
def scale_ln_fcs(ln, fcs, scales):
    if not isinstance(fcs, list):
        fcs = [fcs]
    
    scales = scales.to(ln.weight.device)

    # debugging start even scales = 1 does not work?
    """
    scales = scales * 0
    scales = scales + 1
    """
    # debugging end

    ln.weight.div_(scales)
    if hasattr(ln, 'bias') and ln.bias is not None:
        ln.bias.div_(scales)

    for fc in fcs:
        fc.weight.mul_(scales.view(1, -1))

    for p in ln.parameters():
        assert torch.isnan(p).sum() == 0
    for fc in fcs:
        for p in fc.parameters():
            assert torch.isnan(p).sum() == 0


@torch.no_grad()
def scale_fc_fc(fc1, fc2, scales):
    assert isinstance(fc1, nn.Linear)
    assert isinstance(fc2, nn.Linear)
    # assert fc1.out_features == fc2.in_features
    
    scales = scales.to(fc1.weight.device)

    # fc1.weight.div_(scales.view(-1, 1))
    fc1.weight[-scales.size(0):].div_(scales.view(-1, 1))
    if fc1.bias is not None:
        fc1.bias.div_(scales.view(-1))

    fc2.weight.mul_(scales.view(1, -1))

    for p in fc1.parameters():
        assert torch.isnan(p).sum() == 0
    for p in fc2.parameters():
        assert torch.isnan(p).sum() == 0


@torch.no_grad()
def scale_gelu_fc(gelu, fc, scales):
    assert any(isinstance(gelu,t) for t in [nn.GELU, BloomGelu, NewGELUActivation])
    assert isinstance(fc, nn.Linear)

    fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))

    for p in fc.parameters():
        assert torch.isnan(p).sum() == 0
    

@torch.no_grad()
def auto_scale_block(awq_model,
                     module, 
                     module_kwargs,
                     quant_config,
                     input_feat):
    from .quantizer import pseudo_quantize_tensor
    # firstly, get the weight quantize function
    if quant_config.bits is not None:
        def w_quantize_func(p): return pseudo_quantize_tensor(p, quant_config.bits, quant_config.group_size,
                                                              quant_config.zero_point, quant_config.inplace).detach()
    else:
        def w_quantize_func(p): return p

    if "use_cache" in module_kwargs:
        module_kwargs.pop("use_cache")

    # find the best scale ratio
    def _search_module_scale(block, linears2scale: list, x, kwargs={}):
        # w: co, ci
        # x: n, ci
        weight = torch.cat([_m.weight for _m in linears2scale], dim=0)
        w_max = get_weight_scale(weight, group_size=quant_config.group_size)
        # Clear GPU memory
        del weight
        gc.collect()
        torch.cuda.empty_cache()

        x = x.to(next(block.parameters()).device)
        with torch.no_grad():
            org_out = block(x, **kwargs)
            if isinstance(org_out, tuple):
                org_out = org_out[0]

        x_max = get_act_scale(x)

        best_error = float('inf')
        best_ratio = -1
        best_scales = None

        n_grid = 20
        history = []

        org_sd = {k: v.cpu() for k, v in block.state_dict().items()}
        for ratio in range(n_grid):
            ratio = ratio * 1 / n_grid
            scales = (x_max.pow(ratio) / w_max.pow(1-ratio)
                      ).clamp(min=1e-4).view(-1)
            scales = scales / (scales.max() * scales.min()).sqrt()
            for fc in linears2scale:
                fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))
                fc.weight.data = w_quantize_func(fc.weight.data) / (scales.view(1, -1))
            out = block(x, **kwargs)
            if isinstance(out, tuple):
                out = out[0]

            loss = (org_out - out).float().pow(2).mean().item()  # float prevents overflow
            history.append(loss)
            is_best = loss < best_error
            if is_best:
                best_error = loss
                best_ratio = ratio
                best_scales = scales
            block.load_state_dict(org_sd)
        if best_ratio == -1:
            logging.debug(history)
            raise Exception
        best_scales = best_scales.view(-1)

        assert torch.isnan(best_scales).sum() == 0, best_scales
        return best_scales.detach()

    def _auto_get_scale(prev_op, layers, inp, module2inspect=None, kwargs={}):
        # module2inspect: if given, we will check the output diff of this module instead of layers
        if module2inspect is None:
            assert len(layers) == 1
            module2inspect = layers[0]

        scales = _search_module_scale(module2inspect, layers, inp, kwargs)
        scales = scales.detach().cpu()
        # prev_op_name, [layer_name], scale
        return (get_op_name(module, prev_op), tuple([get_op_name(module, m) for m in layers]), scales)

    layers: list[dict] = awq_model.get_layers_for_scaling(
        module, input_feat, module_kwargs
    )
    scales_list = [_auto_get_scale(**layer) for layer in layers]

    return scales_list

def apply_scale(module, scales_list, input_feat_dict=None):
    for prev_op_name, layer_names, scales in scales_list:
        prev_op = get_op_by_name(module, prev_op_name)
        layers = [get_op_by_name(module, name) for name in layer_names]

        prev_op.cuda()
        for layer in layers:
            layer.cuda()
        scales.cuda()
        
        if isinstance(prev_op, nn.Linear):
            assert len(layers) == 1
            scale_fc_fc(prev_op, layers[0], scales)

        elif any(isinstance(prev_op,t) for t in [nn.LayerNorm, LlamaRMSNorm]) \
             or 'rmsnorm' in str(prev_op.__class__).lower():
            scale_ln_fcs(prev_op, layers, scales)

        elif any(isinstance(prev_op,t) for t in [nn.GELU, BloomGelu, NewGELUActivation]):
            new_module = ScaledActivation(prev_op, scales)
            set_op_by_name(module, prev_op_name, new_module)
            scale_gelu_fc(prev_op, layers[0], scales)
            
        else:
            raise NotImplementedError(
                f"prev_op {type(prev_op)} not supported yet!")
            
        # apply the scaling to input feat if given; prepare it for clipping
        if input_feat_dict is not None:  
            for layer_name in layer_names:
                inp = input_feat_dict[layer_name]
                inp.div_(scales.view(1, -1).to(inp.device))

        prev_op.cpu()
        for layer in layers:
            layer.cpu()
        scales.cpu()

auto_quant/quant/awq/quantize/auto_scale.py
qwopqwop200/AutoQuant
import torch
import torch.nn as nn
# core quantization method (simulated quantization)
def pseudo_quantize_tensor(w, bits=4,
                           group_size=-1,
                           zero_point=True, 
                           inplace=False,
                           get_scale_zp=False):
    org_w_shape = w.shape
    if group_size > 0:
        assert org_w_shape[-1] % group_size == 0
        w = w.reshape(-1, group_size)
    assert w.dim() == 2
    if zero_point:
        max_val = w.amax(dim=1, keepdim=True)
        min_val = w.amin(dim=1, keepdim=True)
        max_int = 2 ** bits - 1
        min_int = 0
        scales = (max_val - min_val).clamp(min=1e-5) / max_int
        zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
    else:  # we actually never used this
        assert min_val is None
        max_val = w.abs().amax(dim=1, keepdim=True)
        max_val = max_val.clamp(min=1e-5)
        max_int = 2 ** (bits - 1) - 1
        min_int = - 2 ** (bits - 1)
        scales = max_val / max_int
        zeros = 0

    assert torch.isnan(scales).sum() == 0
    assert torch.isnan(w).sum() == 0

    if inplace:
        ((w.div_(scales).round_().add_(zeros)).clamp_(min_int, max_int).sub_(zeros)).mul_(scales)
    else:
        w = (torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros) * scales
    assert torch.isnan(w).sum() == 0

    w = w.reshape(org_w_shape)

    if get_scale_zp:
        return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1)
    else:
        return w

class ScaledActivation(nn.Module):
    def __init__(self, module, scales):
        super().__init__()
        self.act = module
        self.scales = nn.Parameter(scales.data)
    
    def forward(self, x):
        return self.act(x) / self.scales.view(1, 1, -1).to(x.device)
auto_quant/quant/awq/quantize/quantizer.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM, GPTQConfig
from .bloom import BloomGPTQForCausalLM
from .codegen import CodeGenGPTQForCausalLM
from .gpt_neox import GPTNeoXGPTQForCausalLM
from .gptj import GPTJGPTQForCausalLM
from .gpt2 import GPT2GPTQForCausalLM
from .llama import LlamaGPTQForCausalLM
from .moss import MOSSGPTQForCausalLM
from .opt import OPTGPTQForCausalLM
from .rw import RWGPTQForCausalLM
from .gpt_bigcode import GPTBigCodeGPTQForCausalLM
from .baichuan import BaiChuanGPTQForCausalLM
from .internlm import InternLMGPTQForCausalLM
from .qwen import QwenGPTQForCausalLM
from .mpt import MptGPTQForCausalLM
from .auto import *

auto_quant/quant/gptq/models/__init__.py
qwopqwop200/AutoQuant
from inspect import signature
from typing import Dict, Optional, Union

import torch
from transformers import AutoConfig
from auto_quant.quant.gptq.models import *

GPTQ_CAUSAL_LM_MODEL_MAP = {
    "bloom": BloomGPTQForCausalLM,
    "gpt_neox": GPTNeoXGPTQForCausalLM,
    "gptj": GPTJGPTQForCausalLM,
    "gpt2": GPT2GPTQForCausalLM,
    "llama": LlamaGPTQForCausalLM,
    "opt": OPTGPTQForCausalLM,
    "moss": MOSSGPTQForCausalLM,
    "gpt_bigcode": GPTBigCodeGPTQForCausalLM,
    "codegen": CodeGenGPTQForCausalLM,
    "RefinedWebModel": RWGPTQForCausalLM,
    "RefinedWeb": RWGPTQForCausalLM,
    "falcon": RWGPTQForCausalLM,
    "baichuan": BaiChuanGPTQForCausalLM,
    "internlm": InternLMGPTQForCausalLM,
    "qwen": QwenGPTQForCausalLM,
    "mpt": MptGPTQForCausalLM,
}

def check_and_get_model_type(model_dir, trust_remote_code=True):
    config = AutoConfig.from_pretrained(model_dir, trust_remote_code=trust_remote_code)
    if config.model_type not in GPTQ_CAUSAL_LM_MODEL_MAP.keys():
        raise TypeError(f"{config.model_type} isn't supported yet. Only support: {list(GPTQ_CAUSAL_LM_MODEL_MAP.keys())}")
    model_type = config.model_type
    return model_type

class AutoGPTQForCausalLM:
    def __init__(self):
        raise EnvironmentError(
            "AutoGPTQModelForCausalLM is designed to be instantiated\n"
            "using `AutoGPTQModelForCausalLM.from_pretrained` if want to quantize a pretrained model.\n"
            "using `AutoGPTQModelForCausalLM.from_quantized` if want to inference with quantized model."
        )

    @classmethod
    def from_pretrained(
        self, 
        model_path, 
        quant_config: GPTQConfig, 
        max_memory: Optional[dict] = None,
        torch_dtype: torch.dtype = torch.float16,
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseGPTQForCausalLM:
        
        model_type = check_and_get_model_type(model_path, trust_remote_code)
        return GPTQ_CAUSAL_LM_MODEL_MAP[model_type].from_pretrained(
            model_path, 
            quant_config=quant_config,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            **kwargs
        )

    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseGPTQForCausalLM:
        model_type = check_and_get_model_type(quant_path, trust_remote_code)
        quant_func = GPTQ_CAUSAL_LM_MODEL_MAP[model_type].from_quantized
        # A static list of kwargs needed for huggingface_hub
        huggingface_kwargs = [
            "cache_dir",
            "force_download",
            "proxies",
            "resume_download",
            "local_files_only",
            "use_auth_token",
            "revision",
            "subfolder",
            "_raise_exceptions_for_missing_entries",
            "_commit_hash"
        ]
        keywords = {
            key: kwargs[key]
            for key in list(signature(quant_func).parameters.keys()) + huggingface_kwargs
            if key in kwargs
        }
        return quant_func(
            quant_path,
            device=device,
            device_map=device_map,
            max_memory=max_memory,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype, 
            trust_remote_code=trust_remote_code, 
            **kwargs
        )
auto_quant/quant/gptq/models/auto.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class BaiChuanGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "DecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.W_pack"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"]
    ]
auto_quant/quant/gptq/models/baichuan.py
qwopqwop200/AutoQuant
import os
import gc
import logging
import copy
from tqdm import tqdm
from dataclasses import dataclass, field, fields
from typing import Dict, List, Optional, Union

import torch
import torch.nn as nn
from torch import LongTensor
from torch.nn.utils.rnn import pad_sequence

from huggingface_hub import snapshot_download
from transformers.modeling_utils import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model, load_checkpoint_in_model, cpu_offload_with_hook
from accelerate.utils import get_balanced_memory
from accelerate.hooks import remove_hook_from_module

from auto_quant.models.base import BaseQuantConfig, BaseQuantForCausalLM
from auto_quant.modules.qlinear_exllama import ExllamaLinear, exllama_post_init, make_sure_no_tensor_in_meta_device, check_exllama_can_save
from auto_quant.utils.module import get_device, move_to_device, get_named_module, get_module_by_name_prefix, get_module_by_name_suffix, set_op_by_name, simple_dispatch_model
from ..quantize import GPTQ

@dataclass
class GPTQConfig(BaseQuantConfig):
    bits: int = field(default=4, metadata={"choices": [4]})
    group_size: int = field(default=-1)
    damp_percent: float = field(default=0.01)
    act_order: bool = field(default=True)
    static_groups: bool = field(default=False)
    sym: bool = field(default=False)
    true_sequential: bool = field(default=True)
    quant_type = 'GPTQ'
    
    def __post_init__(self):
        fields_info = fields(self)

        if self.bits not in fields_info[0].metadata["choices"]:
            raise ValueError(f"only support quantize to {fields_info[0].metadata['choices']} bits.")
        if self.group_size != -1 and self.group_size <= 0:
            raise ValueError("unless equal to -1, group_size must greater then 0.")
        if not (0 < self.damp_percent < 1):
            raise ValueError("damp_percent must between 0 and 1.")

    def to_dict(self):
        return {
            "bits": self.bits,
            "group_size": self.group_size,
            "damp_percent": self.damp_percent,
            "act_order": self.act_order,
            "static_groups": self.static_groups,
            "sym": self.sym,
            "true_sequential": self.true_sequential,
            "quant_type": self.quant_type,
        }


class BaseGPTQForCausalLM(BaseQuantForCausalLM):
    layer_type: str = None
    layers_block_name: str = None
    outside_layer_modules: List[str] = None
    inside_layer_modules: List[List[str]] = None
    lm_head_name: str = "lm_head"
    def __init__(self, model, quant_config, is_quantized):
        super().__init__(model, quant_config, is_quantized)
        self.quant_config:GPTQConfig = quant_config
        self.search_result = None
                
    @classmethod
    def check_quant_config(self, quant_config: GPTQConfig):
        pass

    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        post_init: bool = True,
        **kwargs,
    ):
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", "")
        commit_hash = kwargs.pop("_commit_hash", None)

        snapshot_cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "revision": revision,
        }

        cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "use_auth_token": use_auth_token,
            "revision": revision,
            "subfolder": subfolder,
            "_raise_exceptions_for_missing_entries": False,
            "_commit_hash": commit_hash,
        }
    
        if not os.path.isdir(quant_path):
            quant_path = snapshot_download(quant_path,ignore_patterns=[ "*.msgpack" , "*.h5"], **snapshot_cached_file_kwargs)

        config = AutoConfig.from_pretrained(quant_path, trust_remote_code=trust_remote_code, **cached_file_kwargs)
        quant_config = GPTQConfig.from_pretrained(quant_path, **cached_file_kwargs)

        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code)
        
        self._load_quantized_modules(self, model, quant_config)
        
        model.tie_weights()
        
        if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
            raise ValueError(
                "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
                "'sequential'."
            )
        if isinstance(device_map, dict):
            max_memory = None
        else:
            if device is None and not device_map and not max_memory:
                device_map = "auto"
            if device is not None:
                device = torch.device(device)
                if not max_memory and not device_map:
                    device_map = {"": device.index if device.type == "cuda" else device.type}
            if not isinstance(device_map, dict) and device_map != "sequential":
                max_memory = get_balanced_memory(
                    model=model,
                    max_memory=max_memory,
                    no_split_module_classes=[self.layer_type],
                    low_zero=(device_map == "balanced_low_0")
                )
        if not isinstance(device_map, dict):
            device_map = infer_auto_device_map(
                model,
                max_memory=max_memory,
                no_split_module_classes=[self.layer_type]
            )

        if low_cpu_mem_usage:
            make_sure_no_tensor_in_meta_device(model)
        
        if os.path.isfile(os.path.join(quant_path, WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, WEIGHTS_NAME)
        elif os.path.isfile(os.path.join(quant_path, SAFE_WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, SAFE_WEIGHTS_NAME)
        load_checkpoint_in_model(
            model,
            checkpoint=quant_path,
            device_map=device_map,
            offload_state_dict=True,
            offload_buffers=True
        )
        model = dispatch_model(model, device_map)
        if not(quant_config.act_order) or post_init:
            model = exllama_post_init(model, quant_config.act_order)
            if not check_exllama_can_save(model):
                logging.warning("The model cannot be saved again. If you want to save the model again, set post_init=False. But in this case, inference doesn't work.")
        else:
            logging.warning("Exllama has not been post-initialized. inference doesn't work.")
        return self(model, quant_config, is_quantized=True)
        
    def _load_quantized_modules(self, model, quant_config):
        layers = get_named_module(model)
        ignore_layers = [self.lm_head_name] + self.outside_layer_modules
        for name in list(layers.keys()):
            if any([name.startswith(ignore_layer) for ignore_layer in ignore_layers]):
                del layers[name]

        for name in tqdm(layers.keys(), desc="Replacing layers..."):
            layer = layers[name]
            q_linear = ExllamaLinear.from_linear(
                    layer, 
                    quant_config.bits, 
                    quant_config.group_size,
                    quant_config.act_order,
                    True)
            q_linear.to(get_device(layer))
            set_op_by_name(model, name, q_linear)
            torch.cuda.empty_cache()
            gc.collect()
        
    @staticmethod
    def _resize_attention_mask(attention_mask: List[torch.LongTensor]):
        return attention_mask

    @staticmethod
    def _resize_position_ids(position_ids: List[torch.LongTensor]):
        return position_ids

    def _prepare_examples_for_quantization(self, examples: List[Union[List[int], torch.LongTensor]], batch_size: int = 1, pad_token_id: Optional[int] = None):
        if pad_token_id is None:
            pad_token_id = self.model.config.pad_token_id
            if not pad_token_id:
                pad_token_id = self.config.eos_token_id
    
        if isinstance(examples,list):
            new_examples = []
            for tensor in examples:
                if tensor.dim() > 2:
                    raise Exception('examples must be 2D tensor or less.')
                else:
                    for i in range(tensor.dim(),2):
                        tensor.squeeze(0)
                        
                for i in range(tensor.shape[0]):
                    new_examples.append(tensor[i])
            if pad_token_id is not None:
                examples = pad_sequence(new_examples, True, padding_value=pad_token_id)
            else:
                examples = torch.cat(new_examples,dim=0)
                
        if examples.dim() > 2:
            raise Exception('examples must be 2D tensor or less.')
        else:
            for i in range(examples.dim(),2):
                examples.squeeze(0)

        new_examples = []
        for start in range(0, examples.shape[0], batch_size):
            input_ids = examples[start: start + batch_size]
            new_examples.append({"input_ids": input_ids, "attention_mask": torch.ones_like(input_ids)})
        return new_examples

    @torch.inference_mode()
    def quantize(
        self,
        examples: List[Union[List[int], torch.LongTensor]],
        pad_token_id: Optional[int] = None,
        batch_size: int = 1,
        cache_examples_on_gpu: bool = True,
    ):
        device_map = self.hf_device_map
        if device_map:
            for name, device in device_map.items():
                if device == "cpu":
                    logging.info(f"truly offloading {name} to cpu with hook.")
                    module = get_module_by_name_suffix(self.model, name)
                    remove_hook_from_module(module, recurse=True)
                    cpu_offload_with_hook(module, 'cuda:0')

        layer_inputs = []
        attention_masks = []
        position_ids = []
        layer_input_kwargs = []
        layer_outputs = []

        examples = self._prepare_examples_for_quantization(examples, batch_size, pad_token_id)

        class LayerHijacker(nn.Module):
            """hijack layer's forward pass to cache data"""

            def __init__(self, m, device):
                super().__init__()
                self.module = m
                self.data_device = device if cache_examples_on_gpu else 'cpu'

            def forward(self, inp=None, **kwargs):
                if inp is None:  # some models use all key-value arguments in forward pass call
                    for kwarg_name in ["hidden_states"]:
                        if kwarg_name in kwargs:
                            inp = kwargs[kwarg_name]
                            break
                layer_inputs.append(move_to_device(inp, self.data_device))
                attention_masks.append(kwargs["attention_mask"].to(self.data_device))
                pos_ids = kwargs.get("position_ids", None)
                if pos_ids is not None:
                    position_ids.append(move_to_device(pos_ids, self.data_device))
                one_kwargs = dict()
                for k, v in kwargs.items():  # make sure other arguments also be captured
                    if k not in ["hidden_states", "attention_mask", "position_ids"]:
                        if isinstance(v, torch.Tensor):
                            one_kwargs[k] = move_to_device(v, self.data_device)
                        else:
                            one_kwargs[k] = v
                layer_input_kwargs.append(one_kwargs)
                raise ValueError

        forward_pass_use_cache = self.model.config.use_cache
        self.model.config.use_cache = False

        num_batches = len(examples)
        layers = get_module_by_name_prefix(self.model, self.layers_block_name)

        force_layer_back_to_cpu = False
        if get_device(layers[0]) == 'cpu':
            layers[0] = layers[0].to('cuda:0')
            force_layer_back_to_cpu = True

        cur_layer_device = get_device(layers[0])
        ori_outside_layer_module_devices = {}
        for module_name in self.outside_layer_modules:
            module = get_module_by_name_prefix(self.model, module_name)

            if module is None:
                continue

            ori_outside_layer_module_devices[module_name] = get_device(module)
            if module is not None:
                move_to_device(module, cur_layer_device)

        # get inputs for first layer
        layers[0] = LayerHijacker(layers[0], cur_layer_device)
        for example in examples:
            for k, v in example.items():
                if len(v.shape) == 1:
                    v = v.unsqueeze(0)
                example[k] = move_to_device(v, cur_layer_device)
            try:
                self.model(**example)
            except ValueError:
                pass
        layers[0] = layers[0].module

        move_to_device(layers[0], 'cpu' if force_layer_back_to_cpu else cur_layer_device)
        for module_name in self.outside_layer_modules:
            module = get_module_by_name_prefix(self.model, module_name)
            if module is not None:
                move_to_device(module, ori_outside_layer_module_devices[module_name])

        torch.cuda.empty_cache()

        # resize attention mask and position ids for some special models
        attention_masks = self._resize_attention_mask(attention_masks)
        position_ids = self._resize_position_ids(position_ids)

        inside_layer_modules = self.inside_layer_modules
        if not self.quant_config.true_sequential:
            inside_layer_modules = [sum(inside_layer_modules, [])]
        quantizers = {}
        for i in tqdm(range(len(layers)), desc="GPTQ Quantization"):
            layer = layers[i]
            force_layer_back_to_cpu = False
            if get_device(layer) == 'cpu':
                move_to_device(layer, 'cuda:0')
                force_layer_back_to_cpu = True
            cur_layer_device = get_device(layer)

            full = get_named_module(layer)
            for names in inside_layer_modules:
                subset = {n: full[n] for n in names}
                gptq = {}
                for name in subset:
                    gptq[name] = GPTQ(subset[name])
                    gptq[name].quantizer.configure(
                        self.quant_config.bits,
                        perchannel=True,
                        sym=self.quant_config.sym,
                        mse=False,
                    )

                def add_batch(name):
                    def tmp(_, inp, out):
                        gptq[name].add_batch(inp[0].data, out.data)

                    return tmp

                handles = []
                for name in subset:
                    handles.append(subset[name].register_forward_hook(add_batch(name)))
                for j in range(num_batches):
                    layer_input = move_to_device(layer_inputs[j], cur_layer_device)
                    layer_attention_mask = move_to_device(attention_masks[j], cur_layer_device)
                    additional_layer_inputs = {"attention_mask": layer_attention_mask}
                    layer_position_ids = None if not position_ids else move_to_device(position_ids[j], cur_layer_device)
                    if layer_position_ids is not None:
                        additional_layer_inputs["position_ids"] = layer_position_ids
                    for k, v in layer_input_kwargs[j].items():
                        if isinstance(v, torch.Tensor):
                            additional_layer_inputs[k] = move_to_device(v, cur_layer_device)
                        else:
                            additional_layer_inputs[k] = v
                    layer(layer_input, **additional_layer_inputs)
                for h in handles:
                    h.remove()

                for name in subset:
                    scale, zero, g_idx = gptq[name].fasterquant(
                        percdamp=self.quant_config.damp_percent,
                        group_size=self.quant_config.group_size,
                        actorder=self.quant_config.act_order,
                        static_groups=self.quant_config.static_groups
                    )

                    quantizers[f'{self.layers_block_name}.{i}.{name}'] = (
                        move_to_device(scale, 'cpu' if force_layer_back_to_cpu else cur_layer_device),
                        move_to_device(zero, 'cpu' if force_layer_back_to_cpu else cur_layer_device),
                        move_to_device(g_idx, 'cpu' if force_layer_back_to_cpu else cur_layer_device)
                    )
                    gptq[name].free()

            for j in range(num_batches):
                layer_input = move_to_device(layer_inputs[j], cur_layer_device)
                layer_attention_mask = move_to_device(attention_masks[j], cur_layer_device)
                additional_layer_inputs = {"attention_mask": layer_attention_mask}
                layer_position_ids = None if not position_ids else move_to_device(position_ids[j], cur_layer_device)
                if layer_position_ids is not None:
                    additional_layer_inputs["position_ids"] = layer_position_ids
                for k, v in layer_input_kwargs[j].items():
                    if isinstance(v, torch.Tensor):
                        additional_layer_inputs[k] = move_to_device(v, cur_layer_device)
                    else:
                        additional_layer_inputs[k] = v
                layer_output = move_to_device(
                    layer(layer_input, **additional_layer_inputs)[0],
                    cur_layer_device if cache_examples_on_gpu else 'cpu'
                )
                layer_outputs.append(layer_output)

            layers[i] = move_to_device(layer, 'cpu' if force_layer_back_to_cpu else cur_layer_device)
            del layer
            del gptq
            del layer_inputs
            layer_inputs, layer_outputs = layer_outputs, []
            torch.cuda.empty_cache()

        self.pack_model(quantizers=quantizers)
        if device_map:
            self.model = remove_hook_from_module(self.model, recurse=True)
            self.model = simple_dispatch_model(self.model, device_map)

        self.model.config.use_cache = forward_pass_use_cache
        logging.warning("Exllama has not been post-initialized. inference doesn't work. Please save the model and then load it.")

        self.is_quantized = True
        torch.cuda.empty_cache()
            
    def pack_model(self, quantizers):
        layers = get_named_module(self.model)
        layers = {n: layers[n] for n in quantizers}
        for name in tqdm(layers, desc="Pack Model"):
            scale, zero, g_idx = quantizers[name]
            layer_device = get_device(layers[name])
            layer, scale, zero, g_idx = layers[name].to('cpu'), scale.to('cpu'), zero.to('cpu'), g_idx.to('cpu')
            qlayer = ExllamaLinear.from_linear(
                    layer, 
                    self.quant_config.bits, 
                    self.quant_config.group_size,
                    self.quant_config.act_order,
                    False, 
                    scale, 
                    zero,
                    g_idx)
            qlayer.to(layer_device)
            set_op_by_name(self.model, name, qlayer)
            torch.cuda.empty_cache()
            gc.collect()
auto_quant/quant/gptq/models/base.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class BloomGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "BloomBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.word_embeddings", "transformer.word_embeddings_layernorm", "transformer.ln_f"]
    inside_layer_modules = [
        ["self_attention.query_key_value"],
        ["self_attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
auto_quant/quant/gptq/models/bloom.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class CodeGenGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "CodeGenBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.qkv_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]

auto_quant/quant/gptq/models/codegen.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class GPT2GPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "GPT2Block"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.wpe", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.c_fc"],
        ["mlp.c_proj"]
    ]
auto_quant/quant/gptq/models/gpt2.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class GPTBigCodeGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "GPTBigCodeBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = [
        "transformer.wpe", "transformer.wte", "transformer.ln_f"
    ]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.c_fc"],
        ["mlp.c_proj"]
    ]
auto_quant/quant/gptq/models/gpt_bigcode.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class GPTNeoXGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "GPTNeoXLayer"
    layers_block_name = "gpt_neox.layers"
    outside_layer_modules = ["gpt_neox.embed_in", "gpt_neox.final_layer_norm"]
    inside_layer_modules = [
        ["attention.query_key_value"],
        ["attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
    lm_head_name = "embed_out"
auto_quant/quant/gptq/models/gpt_neox.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class GPTJGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "GPTJBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.k_proj", "attn.v_proj", "attn.q_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]
auto_quant/quant/gptq/models/gptj.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class InternLMGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "InternLMDecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"],
    ]

auto_quant/quant/gptq/models/internlm.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class LlamaGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "LlamaDecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"]
    ]
auto_quant/quant/gptq/models/llama.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class MOSSGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "MossBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.qkv_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]
auto_quant/quant/gptq/models/moss.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class MptGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "MPTBlock"
    layers_block_name = "transformer.blocks"
    outside_layer_modules = [
        "transformer.wte",  "transformer.norm_f"
    ]

    inside_layer_modules = [
        ["attn.Wqkv"],
        ["attn.out_proj"],
        ["ffn.up_proj"],
        ["ffn.down_proj"]
    ]
auto_quant/quant/gptq/models/mpt.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class OPTGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "OPTDecoderLayer"
    layers_block_name = "model.decoder.layers"
    outside_layer_modules = [
        "model.decoder.embed_tokens", "model.decoder.embed_positions", "model.decoder.project_out",
        "model.decoder.project_in", "model.decoder.final_layer_norm"
    ]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.out_proj"],
        ["fc1"],
        ["fc2"]
    ]
auto_quant/quant/gptq/models/opt.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class QwenGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "QWenBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.wpe", "transformer.ln_f", "transformer.visual"]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.w1", "mlp.w2"],
        ["mlp.c_proj"]
    ]

auto_quant/quant/gptq/models/qwen.py
qwopqwop200/AutoQuant
from .base import BaseGPTQForCausalLM

class RWGPTQForCausalLM(BaseGPTQForCausalLM):
    layer_type = "DecoderLayer"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.word_embeddings", "transformer.ln_f"]
    inside_layer_modules = [
        ["self_attention.query_key_value"],
        ["self_attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
auto_quant/quant/gptq/models/rw.py
qwopqwop200/AutoQuant
from .gptq import GPTQ
from .quantizer import Quantizer, quantize

auto_quant/quant/gptq/quantize/__init__.py
qwopqwop200/AutoQuant
import math
import os
import time
import torch
import torch.nn as nn
import transformers

from .quantizer import Quantizer

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

class GPTQ:
    def __init__(self, layer):
        self.layer = layer
        self.dev = self.layer.weight.device
        W = layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.pytorch_utils.Conv1D):
            W = W.t()
        self.rows = W.shape[0]
        self.columns = W.shape[1]
        self.H = torch.zeros((self.columns, self.columns), device=self.dev)
        self.nsamples = 0
        self.quantizer = Quantizer()

    def add_batch(self, inp, out):
        if os.environ.get("DEBUG"):
            self.inp1 = inp
            self.out1 = out
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(self.layer, transformers.Conv1D):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        if isinstance(self.layer, nn.Conv2d):
            unfold = nn.Unfold(
                self.layer.kernel_size,
                dilation=self.layer.dilation,
                padding=self.layer.padding,
                stride=self.layer.stride
            )
            inp = unfold(inp)
            inp = inp.permute([1, 0, 2])
            inp = inp.flatten(1)
        self.H *= self.nsamples / (self.nsamples + tmp)
        self.nsamples += tmp
        # inp = inp.float()
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        # self.H += 2 / self.nsamples * inp.matmul(inp.t())
        self.H += inp.matmul(inp.t())

    def fasterquant(
        self, blocksize=128, percdamp=.01, group_size=-1, actorder=False, static_groups=False
    ):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        if not self.quantizer.ready():
            self.quantizer.find_params(W, weight=True)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        g_idx = []
        scale = []
        zero = []
        now_idx = 1

        if static_groups:
            import copy
            groups = []
            for i in range(0, self.columns, group_size):
                quantizer = copy.deepcopy(self.quantizer)
                quantizer.find_params(W[:, i:(i + group_size)], weight=True)
                scale.append(quantizer.scale)
                zero.append(quantizer.zero)
                groups.append(quantizer)

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]
            invperm = torch.argsort(perm)

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if group_size != -1:
                    if not static_groups:
                        if (i1 + i) % group_size == 0:
                            self.quantizer.find_params(W[:, (i1 + i):(i1 + i + group_size)], weight=True)
                            
                        if ((i1 + i) // group_size) - now_idx == -1:
                            scale.append(self.quantizer.scale)
                            zero.append(self.quantizer.zero)
                            now_idx += 1
                    else:
                        idx = i1 + i
                        if actorder:
                            idx = perm[idx]
                        self.quantizer = groups[idx // group_size]
                        
                q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

            if os.environ.get("DEBUG"):
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                logger.debug(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                logger.debug(torch.sum(Losses))

        torch.cuda.synchronize()

        group_size = group_size if group_size != -1 else self.columns
        if static_groups and actorder:
            g_idx = [perm[i] // group_size for i in range(self.columns)]
        else:
            g_idx = [i // group_size for i in range(self.columns)]
        g_idx = torch.tensor(g_idx, dtype=torch.int32, device=Q.device)
        if actorder:
            Q = Q[:, invperm]
            g_idx = g_idx[invperm]

        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        self.layer.weight.data = Q.reshape(self.layer.weight.shape).type_as(self.layer.weight.data)
        if os.environ.get("DEBUG"):
            logger.debug(torch.sum((self.layer(self.inp1) - self.out1) ** 2))

        if scale == []:
            scale.append(self.quantizer.scale)
            zero.append(self.quantizer.zero)
        scale = torch.cat(scale, dim=1)
        zero = torch.cat(zero, dim=1)
        return scale, zero, g_idx

    def free(self):
        if os.environ.get("DEBUG"):
            self.inp1 = None
            self.out1 = None
        self.H = None
        self.Losses = None
        self.Trace = None
        torch.cuda.empty_cache()
auto_quant/quant/gptq/quantize/gptq.py
qwopqwop200/AutoQuant
import torch
import torch.nn as nn

def quantize(x, scale, zero, maxq):
    if maxq < 0:
        return (x > scale / 2).float() * scale + (x < zero / 2).float() * zero
    q = torch.clamp(torch.round(x / scale) + zero, 0, maxq)
    return scale * (q - zero)


class Quantizer(nn.Module):

    def __init__(self, shape=1):
        super(Quantizer, self).__init__()
        self.register_buffer('maxq', torch.tensor(0))
        self.register_buffer('scale', torch.zeros(shape))
        self.register_buffer('zero', torch.zeros(shape))

    def configure(
        self,
        bits, perchannel=False, sym=True,
        mse=False, norm=2.4, grid=100, maxshrink=.8,
        trits=False
    ):

        self.maxq = torch.tensor(2 ** bits - 1)
        self.perchannel = perchannel
        self.sym = sym
        self.mse = mse
        self.norm = norm
        self.grid = grid
        self.maxshrink = maxshrink
        if trits:
            self.maxq = torch.tensor(-1)

    def find_params(self, x, weight=False):
        dev = x.device
        self.maxq = self.maxq.to(dev)

        shape = x.shape
        if self.perchannel:
            if weight:
                x = x.flatten(1)
            else:
                if len(shape) == 4:
                    x = x.permute([1, 0, 2, 3])
                    x = x.flatten(1)
                if len(shape) == 3:
                    x = x.reshape((-1, shape[-1])).t()
                if len(shape) == 2:
                    x = x.t()
        else:
            x = x.flatten().unsqueeze(0)

        tmp = torch.zeros(x.shape[0], device=dev)
        xmin = torch.minimum(x.min(1)[0], tmp)
        xmax = torch.maximum(x.max(1)[0], tmp)

        if self.sym:
            xmax = torch.maximum(torch.abs(xmin), xmax)
            tmp = xmin < 0
            if torch.any(tmp):
                xmin[tmp] = -xmax[tmp]
        tmp = (xmin == 0) & (xmax == 0)
        xmin[tmp] = -1
        xmax[tmp] = +1

        if self.maxq < 0:
            self.scale = xmax
            self.zero = xmin
        else:
            self.scale = (xmax - xmin) / self.maxq
            if self.sym:
                self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2)
            else:
                self.zero = torch.round(-xmin / self.scale)

        if self.mse:
            best = torch.full([x.shape[0]], float('inf'), device=dev)
            for i in range(int(self.maxshrink * self.grid)):
                p = 1 - i / self.grid
                xmin1 = p * xmin
                xmax1 = p * xmax
                scale1 = (xmax1 - xmin1) / self.maxq
                zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero
                q = quantize(x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq)
                q -= x
                q.abs_()
                q.pow_(self.norm)
                err = torch.sum(q, 1)
                tmp = err < best
                if torch.any(tmp):
                    best[tmp] = err[tmp]
                    self.scale[tmp] = scale1[tmp]
                    self.zero[tmp] = zero1[tmp]
        if not self.perchannel:
            if weight:
                tmp = shape[0]
            else:
                tmp = shape[1] if len(shape) != 3 else shape[2]
            self.scale = self.scale.repeat(tmp)
            self.zero = self.zero.repeat(tmp)

        if weight:
            shape = [-1] + [1] * (len(shape) - 1)
            self.scale = self.scale.reshape(shape)
            self.zero = self.zero.reshape(shape)
            return
        if len(shape) == 4:
            self.scale = self.scale.reshape((1, -1, 1, 1))
            self.zero = self.zero.reshape((1, -1, 1, 1))
        if len(shape) == 3:
            self.scale = self.scale.reshape((1, 1, -1))
            self.zero = self.zero.reshape((1, 1, -1))
        if len(shape) == 2:
            self.scale = self.scale.unsqueeze(0)
            self.zero = self.zero.unsqueeze(0)

    def quantize(self, x):
        if self.ready():
            return quantize(x, self.scale, self.zero, self.maxq)
        return x

    def enabled(self):
        return self.maxq > 0

    def ready(self):
        return torch.all(self.scale != 0)

auto_quant/quant/gptq/quantize/quantizer.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM, RTNConfig
from .bloom import BloomRTNForCausalLM
from .codegen import CodeGenRTNForCausalLM
from .gpt_neox import GPTNeoXRTNForCausalLM
from .gptj import GPTJRTNForCausalLM
from .gpt2 import GPT2RTNForCausalLM
from .llama import LlamaRTNForCausalLM
from .moss import MOSSRTNForCausalLM
from .opt import OPTRTNForCausalLM
from .rw import RWRTNForCausalLM
from .gpt_bigcode import GPTBigCodeRTNForCausalLM
from .baichuan import BaiChuanRTNForCausalLM
from .internlm import InternLMRTNForCausalLM
from .qwen import QwenRTNForCausalLM
from .mpt import MptRTNForCausalLM
from .auto import *

auto_quant/quant/rtn/models/__init__.py
qwopqwop200/AutoQuant
from inspect import signature
from typing import Dict, Optional, Union

import torch
from transformers import AutoConfig
from auto_quant.quant.rtn.models import *

RTN_CAUSAL_LM_MODEL_MAP = {
    "bloom": BloomRTNForCausalLM,
    "gpt_neox": GPTNeoXRTNForCausalLM,
    "gptj": GPTJRTNForCausalLM,
    "gpt2": GPT2RTNForCausalLM,
    "llama": LlamaRTNForCausalLM,
    "opt": OPTRTNForCausalLM,
    "moss": MOSSRTNForCausalLM,
    "gpt_bigcode": GPTBigCodeRTNForCausalLM,
    "codegen": CodeGenRTNForCausalLM,
    "RefinedWebModel": RWRTNForCausalLM,
    "RefinedWeb": RWRTNForCausalLM,
    "falcon": RWRTNForCausalLM,
    "baichuan": BaiChuanRTNForCausalLM,
    "internlm": InternLMRTNForCausalLM,
    "qwen": QwenRTNForCausalLM,
    "mpt": MptRTNForCausalLM,
}

def check_and_get_model_type(model_dir, trust_remote_code=True):
    config = AutoConfig.from_pretrained(model_dir, trust_remote_code=trust_remote_code)
    if config.model_type not in RTN_CAUSAL_LM_MODEL_MAP.keys():
        raise TypeError(f"{config.model_type} isn't supported yet. Only support: {list(RTN_CAUSAL_LM_MODEL_MAP.keys())}")
    model_type = config.model_type
    return model_type

class AutoRTNForCausalLM:
    def __init__(self):
        raise EnvironmentError(
            "AutoRTNForCausalLM is designed to be instantiated\n"
            "using `AutoRTNForCausalLM.from_pretrained` if want to quantize a pretrained model.\n"
            "using `AutoRTNForCausalLM.from_quantized` if want to inference with quantized model."
        )

    @classmethod
    def from_pretrained(
        self, 
        model_path, 
        quant_config: RTNConfig, 
        max_memory: Optional[dict] = None,
        torch_dtype: torch.dtype = torch.float16,
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseRTNForCausalLM:
        
        model_type = check_and_get_model_type(model_path, trust_remote_code)
        return RTN_CAUSAL_LM_MODEL_MAP[model_type].from_pretrained(
            model_path, 
            quant_config=quant_config,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            **kwargs
        )

    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs
    ) -> BaseRTNForCausalLM:
        model_type = check_and_get_model_type(quant_path, trust_remote_code)
        quant_func = RTN_CAUSAL_LM_MODEL_MAP[model_type].from_quantized
        # A static list of kwargs needed for huggingface_hub
        huggingface_kwargs = [
            "cache_dir",
            "force_download",
            "proxies",
            "resume_download",
            "local_files_only",
            "use_auth_token",
            "revision",
            "subfolder",
            "_raise_exceptions_for_missing_entries",
            "_commit_hash"
        ]
        keywords = {
            key: kwargs[key]
            for key in list(signature(quant_func).parameters.keys()) + huggingface_kwargs
            if key in kwargs
        }
        return quant_func(
            quant_path,
            device=device,
            device_map=device_map,
            max_memory=max_memory,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype, 
            trust_remote_code=trust_remote_code, 
            **kwargs
        )
auto_quant/quant/rtn/models/auto.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class BaiChuanRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "DecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.W_pack"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"]
    ]
auto_quant/quant/rtn/models/baichuan.py
qwopqwop200/AutoQuant
import os
import gc
import logging
import copy
from tqdm import tqdm
from dataclasses import dataclass, field, fields
from typing import Dict, List, Optional, Union

import torch
import torch.nn as nn
from torch import LongTensor
from torch.nn.utils.rnn import pad_sequence

from huggingface_hub import snapshot_download
from transformers.modeling_utils import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model, load_checkpoint_in_model, cpu_offload_with_hook
from accelerate.utils import get_balanced_memory
from accelerate.hooks import remove_hook_from_module

from auto_quant.models.base import BaseQuantConfig, BaseQuantForCausalLM
from auto_quant.modules.qlinear_exllama import ExllamaLinear, exllama_post_init, make_sure_no_tensor_in_meta_device, check_exllama_can_save
from auto_quant.utils.module import get_device, move_to_device, get_named_module, get_module_by_name_prefix, get_module_by_name_suffix, set_op_by_name, simple_dispatch_model
from ..quantize import RTN

@dataclass
class RTNConfig(BaseQuantConfig):
    bits: int = field(default=4, metadata={"choices": [4]})
    sym: bool = field(default=False)
    quant_type = 'RTN'
    
    def __post_init__(self):
        fields_info = fields(self)

        if self.bits not in fields_info[0].metadata["choices"]:
            raise ValueError(f"only support quantize to {fields_info[0].metadata['choices']} bits.")

    def to_dict(self):
        return {
            "bits": self.bits,
            "sym": self.sym,
            "quant_type": self.quant_type,
        }


class BaseRTNForCausalLM(BaseQuantForCausalLM):
    layer_type: str = None
    layers_block_name: str = None
    outside_layer_modules: List[str] = None
    inside_layer_modules: List[List[str]] = None
    lm_head_name: str = "lm_head"
    def __init__(self, model, quant_config, is_quantized):
        super().__init__(model, quant_config, is_quantized)
        self.quant_config:RTNConfig = quant_config
        self.search_result = None
                
    @classmethod
    def check_quant_config(self, quant_config: RTNConfig):
        pass

    @classmethod
    def from_quantized(
        self, 
        quant_path, 
        device: Optional[Union[str, int]] = None,
        device_map: Optional[Union[str, Dict[str, Union[int, str]]]] = None,
        max_memory: Optional[dict] = None,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.float16, 
        trust_remote_code: bool = False,
        **kwargs,
    ):
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_auth_token = kwargs.pop("use_auth_token", None)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", "")
        commit_hash = kwargs.pop("_commit_hash", None)

        snapshot_cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "revision": revision,
        }

        cached_file_kwargs = {
            "cache_dir": cache_dir,
            "force_download": force_download,
            "proxies": proxies,
            "resume_download": resume_download,
            "local_files_only": local_files_only,
            "use_auth_token": use_auth_token,
            "revision": revision,
            "subfolder": subfolder,
            "_raise_exceptions_for_missing_entries": False,
            "_commit_hash": commit_hash,
        }
    
        if not os.path.isdir(quant_path):
            quant_path = snapshot_download(quant_path,ignore_patterns=[ "*.msgpack" , "*.h5"], **snapshot_cached_file_kwargs)

        config = AutoConfig.from_pretrained(quant_path, trust_remote_code=trust_remote_code, **cached_file_kwargs)
        quant_config = RTNConfig.from_pretrained(quant_path, **cached_file_kwargs)

        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code)
        
        self._load_quantized_modules(self, model, quant_config)
        
        model.tie_weights()
        
        if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
            raise ValueError(
                "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
                "'sequential'."
            )
        if isinstance(device_map, dict):
            max_memory = None
        else:
            if device is None and not device_map and not max_memory:
                device_map = "auto"
            if device is not None:
                device = torch.device(device)
                if not max_memory and not device_map:
                    device_map = {"": device.index if device.type == "cuda" else device.type}
            if not isinstance(device_map, dict) and device_map != "sequential":
                max_memory = get_balanced_memory(
                    model=model,
                    max_memory=max_memory,
                    no_split_module_classes=[self.layer_type],
                    low_zero=(device_map == "balanced_low_0")
                )
        if not isinstance(device_map, dict):
            device_map = infer_auto_device_map(
                model,
                max_memory=max_memory,
                no_split_module_classes=[self.layer_type]
            )

        if low_cpu_mem_usage:
            make_sure_no_tensor_in_meta_device(model)
        
        if os.path.isfile(os.path.join(quant_path, WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, WEIGHTS_NAME)
        elif os.path.isfile(os.path.join(quant_path, SAFE_WEIGHTS_NAME)):
            quant_path = os.path.join(quant_path, SAFE_WEIGHTS_NAME)
        load_checkpoint_in_model(
            model,
            checkpoint=quant_path,
            device_map=device_map,
            offload_state_dict=True,
            offload_buffers=True
        )
        model = dispatch_model(model, device_map)
        model = exllama_post_init(model, False)
        return self(model, quant_config, is_quantized=True)
        
    def _load_quantized_modules(self, model, quant_config):
        layers = get_named_module(model)
        ignore_layers = [self.lm_head_name] + self.outside_layer_modules
        for name in list(layers.keys()):
            if any([name.startswith(ignore_layer) for ignore_layer in ignore_layers]):
                del layers[name]

        for name in tqdm(layers.keys(), desc="Replacing layers..."):
            layer = layers[name]
            q_linear = ExllamaLinear.from_linear(
                    layer, 
                    quant_config.bits, 
                    -1,
                    False,
                    True)
            q_linear.to(get_device(layer))
            set_op_by_name(model, name, q_linear)
            torch.cuda.empty_cache()
            gc.collect()
        
    @staticmethod
    def _resize_attention_mask(attention_mask: List[torch.LongTensor]):
        return attention_mask

    @staticmethod
    def _resize_position_ids(position_ids: List[torch.LongTensor]):
        return position_ids

    @torch.inference_mode()
    def quantize(self):
        device_map = self.hf_device_map
        layers = get_module_by_name_prefix(self.model, self.layers_block_name)
        inside_layer_modules = self.inside_layer_modules
        quantizers = {}
        for i in tqdm(range(len(layers)), desc="RTN Quantization"):
            layer = layers[i]
            force_layer_back_to_cpu = False
            if get_device(layer) == 'cpu':
                move_to_device(layer, 'cuda:0')
                force_layer_back_to_cpu = True
            cur_layer_device = get_device(layer)

            full = get_named_module(layer)
            for names in inside_layer_modules:
                subset = {n: full[n] for n in names}
                rtn = {}
                for name in subset:
                    rtn[name] = RTN(subset[name])
                    rtn[name].quantizer.configure(
                        self.quant_config.bits,
                        perchannel=True,
                        sym=self.quant_config.sym,
                        mse=False,
                    )
                    
                for name in subset:
                    scale, zero, g_idx = rtn[name].fasterquant()
                    quantizers[f'{self.layers_block_name}.{i}.{name}'] = (
                        move_to_device(scale, 'cpu' if force_layer_back_to_cpu else cur_layer_device),
                        move_to_device(zero, 'cpu' if force_layer_back_to_cpu else cur_layer_device),
                        g_idx
                    )

            layers[i] = move_to_device(layer, 'cpu' if force_layer_back_to_cpu else cur_layer_device)
            del layer
            del rtn
            torch.cuda.empty_cache()

        self.pack_model(quantizers=quantizers)
        if device_map:
            self.model = simple_dispatch_model(self.model, device_map)
        logging.warning("Exllama has not been post-initialized. inference doesn't work. Please save the model and then load it.")
        self.is_quantized = True
        torch.cuda.empty_cache()
            
    def pack_model(self, quantizers):
        layers = get_named_module(self.model)
        layers = {n: layers[n] for n in quantizers}
        for name in tqdm(layers, desc="Pack Model"):
            scale, zero, g_idx = quantizers[name]
            layer_device = get_device(layers[name])
            layer, scale, zero, g_idx = layers[name].to('cpu'), scale.to('cpu'), zero.to('cpu'), g_idx
            qlayer = ExllamaLinear.from_linear(
                    layer, 
                    self.quant_config.bits, 
                    -1,
                    False,
                    False, 
                    scale, 
                    zero,
                    g_idx)
            qlayer.to(layer_device)
            set_op_by_name(self.model, name, qlayer)
            torch.cuda.empty_cache()
            gc.collect()
auto_quant/quant/rtn/models/base.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class BloomRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "BloomBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.word_embeddings", "transformer.word_embeddings_layernorm", "transformer.ln_f"]
    inside_layer_modules = [
        ["self_attention.query_key_value"],
        ["self_attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
auto_quant/quant/rtn/models/bloom.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class CodeGenRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "CodeGenBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.qkv_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]

auto_quant/quant/rtn/models/codegen.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class GPT2RTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "GPT2Block"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.wpe", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.c_fc"],
        ["mlp.c_proj"]
    ]
auto_quant/quant/rtn/models/gpt2.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class GPTBigCodeRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "GPTBigCodeBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = [
        "transformer.wpe", "transformer.wte", "transformer.ln_f"
    ]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.c_fc"],
        ["mlp.c_proj"]
    ]
auto_quant/quant/rtn/models/gpt_bigcode.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class GPTNeoXRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "GPTNeoXLayer"
    layers_block_name = "gpt_neox.layers"
    outside_layer_modules = ["gpt_neox.embed_in", "gpt_neox.final_layer_norm"]
    inside_layer_modules = [
        ["attention.query_key_value"],
        ["attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
    lm_head_name = "embed_out"
auto_quant/quant/rtn/models/gpt_neox.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class GPTJRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "GPTJBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.k_proj", "attn.v_proj", "attn.q_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]
auto_quant/quant/rtn/models/gptj.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class InternLMRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "InternLMDecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"],
    ]

auto_quant/quant/rtn/models/internlm.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class LlamaRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "LlamaDecoderLayer"
    layers_block_name = "model.layers"
    outside_layer_modules = ["model.embed_tokens", "model.norm"]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.o_proj"],
        ["mlp.up_proj", "mlp.gate_proj"],
        ["mlp.down_proj"]
    ]
auto_quant/quant/rtn/models/llama.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class MOSSRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "MossBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.ln_f"]
    inside_layer_modules = [
        ["attn.qkv_proj"],
        ["attn.out_proj"],
        ["mlp.fc_in"],
        ["mlp.fc_out"]
    ]
auto_quant/quant/rtn/models/moss.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class MptRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "MPTBlock"
    layers_block_name = "transformer.blocks"
    outside_layer_modules = [
        "transformer.wte",  "transformer.norm_f"
    ]

    inside_layer_modules = [
        ["attn.Wqkv"],
        ["attn.out_proj"],
        ["ffn.up_proj"],
        ["ffn.down_proj"]
    ]
auto_quant/quant/rtn/models/mpt.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class OPTRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "OPTDecoderLayer"
    layers_block_name = "model.decoder.layers"
    outside_layer_modules = [
        "model.decoder.embed_tokens", "model.decoder.embed_positions", "model.decoder.project_out",
        "model.decoder.project_in", "model.decoder.final_layer_norm"
    ]
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.out_proj"],
        ["fc1"],
        ["fc2"]
    ]
auto_quant/quant/rtn/models/opt.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class QwenRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "QWenBlock"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.wte", "transformer.wpe", "transformer.ln_f", "transformer.visual"]
    inside_layer_modules = [
        ["attn.c_attn"],
        ["attn.c_proj"],
        ["mlp.w1", "mlp.w2"],
        ["mlp.c_proj"]
    ]

auto_quant/quant/rtn/models/qwen.py
qwopqwop200/AutoQuant
from .base import BaseRTNForCausalLM

class RWRTNForCausalLM(BaseRTNForCausalLM):
    layer_type = "DecoderLayer"
    layers_block_name = "transformer.h"
    outside_layer_modules = ["transformer.word_embeddings", "transformer.ln_f"]
    inside_layer_modules = [
        ["self_attention.query_key_value"],
        ["self_attention.dense"],
        ["mlp.dense_h_to_4h"],
        ["mlp.dense_4h_to_h"]
    ]
auto_quant/quant/rtn/models/rw.py
qwopqwop200/AutoQuant
from .rtn import RTN
auto_quant/quant/rtn/quantize/__init__.py
qwopqwop200/AutoQuant
import math
import scipy
import time

import torch
import torch.nn as nn
import transformers

from auto_quant.quant.gptq.quantize.quantizer import Quantizer

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

class RTN:
    def __init__(self, layer):
        self.layer = layer
        self.quantizer = Quantizer()

    def fasterquant(self,):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        if not self.quantizer.ready():
            self.quantizer.find_params(W, weight=True)
        
        Q = self.quantizer.quantize(W).to(self.layer.weight.data.dtype)
        
        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        self.layer.weight.data = Q
        return self.quantizer.scale, self.quantizer.zero, None
auto_quant/quant/rtn/quantize/rtn.py
qwopqwop200/AutoQuant
from .data import get_calib_dataset
from .lm_eval_adaptor import LMEvalAdaptor
auto_quant/utils/__init__.py
qwopqwop200/AutoQuant
import random
import numpy as np
import torch
from datasets import load_dataset

def get_calib_dataset(tokenizer, name='wikitext2', n_samples=128, seq_length=512, seed=42):
    DATASET_MAP = {
        'wikitext2': get_wikitext2,
        'ptb': get_ptb,
        'ptb_new':get_ptb_new,
        'c4':get_c4,
        'pile':get_pile
    }
    
    if name not in DATASET_MAP:
        raise NotImplementedError(f"{name} isn't supported yet. Only support: {list(DATASET_MAP.keys())}")
    
    return DATASET_MAP[name](tokenizer, n_samples, seq_length, seed)

def get_wikitext2(tokenizer, n_samples=128, seq_length=512, seed=42):
    random.seed(seed)
    
    data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train') 
    enc = tokenizer("\n\n".join(data['text']), return_tensors='pt')

    calib_data = []
    for _ in range(n_samples):
        i = random.randint(0, enc.input_ids.shape[1] - seq_length - 1)
        j = i + seq_length
        inp = enc.input_ids[:, i:j]
        calib_data.append(inp)
    return calib_data

def get_ptb(tokenizer, n_samples=128, seq_length=512, seed=42):
    random.seed(seed)
    
    data = load_dataset('ptb_text_only', 'penn_treebank', split='train')
    enc = tokenizer("\n\n".join(data['sentence']), return_tensors='pt')

    calib_data = []
    for _ in range(n_samples):
        i = random.randint(0, enc.input_ids.shape[1] - seq_length - 1)
        j = i + seq_length
        inp = enc.input_ids[:, i:j]
        calib_data.append(inp)
    return calib_data

def get_ptb_new(tokenizer, n_samples=128, seq_length=512, seed=42):
    random.seed(seed)
    
    data = load_dataset('ptb_text_only', 'penn_treebank', split='train')
    enc = tokenizer(" ".join(data['sentence']), return_tensors='pt')

    calib_data = []
    for _ in range(n_samples):
        i = random.randint(0, enc.input_ids.shape[1] - seq_length - 1)
        j = i + seq_length
        inp = enc.input_ids[:, i:j]
        calib_data.append(inp)
    return calib_data

def get_c4(tokenizer, n_samples=128, seq_length=512, seed=42):
    random.seed(seed)
    
    data = load_dataset('allenai/c4', 'allenai--c4', data_files={'train': 'en/c4-train.00000-of-01024.json.gz'}, split='train')

    calib_data = []
    for _ in range(n_samples):
        while True:
            i = random.randint(0, len(data) - 1)
            enc = tokenizer(data[i]['text'], return_tensors='pt')
            if enc.input_ids.shape[1] >= seq_length:
                break
        i = random.randint(0, enc.input_ids.shape[1] - seq_length - 1)
        j = i + seq_length
        inp = enc.input_ids[:, i:j]
        calib_data.append(inp)
    return calib_data

def get_pile(tokenizer, n_samples=128, seq_length=512, seed=42):
    dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
    dataset = dataset.shuffle(seed=seed)
    samples = []
    total_len = 0
    for data in dataset:
        line = data["text"]
        line = line.strip()
        line_encoded = tokenizer.encode(line)
        sample = torch.tensor([line_encoded])
        if sample.numel() == 0:
            continue
        total_len += sample.shape[1]
        samples.append(sample)
        if total_len // seq_length >= n_samples:
            break
    # now concatenate all samples and split according to block size
    cat_samples = torch.cat(samples, dim=1)
    n_split = cat_samples.shape[1] // seq_length
    calib_data = [cat_samples[:, i*seq_length:(i+1)*seq_length] for i in range(n_samples)]
    return calib_data
auto_quant/utils/data.py
qwopqwop200/AutoQuant
import logging
import torch
import transformers
from lm_eval.base import BaseLM

class LMEvalAdaptor(BaseLM):
    def __init__(self, model, tokenizer, batch_size=1, max_length=-1):
        super().__init__()
        assert isinstance(batch_size, int)
        self.model = model
        self.model.eval()
        self.tokenizer = tokenizer

        self.vocab_size = self.tokenizer.vocab_size
        self._batch_size = batch_size
        self._max_length = max_length

    @property
    def eot_token_id(self):
        # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
        return self.tokenizer.eos_token_id

    @property
    def max_length(self):
        if self._max_length != -1:
            return self._max_length
        if hasattr(self.model.config, 'n_ctx'):
            return self.model.config.n_ctx
        elif hasattr(self.model.config, 'max_position_embeddings'):
            return self.model.config.max_position_embeddings
        elif hasattr(self.model.config, 'n_positions'):
            return self.model.config.n_positions
        else:
            logger.warning("can't get model's sequence length from model config, will set to 2048.")
            return 2048

    @property
    def max_gen_toks(self):
        return 256

    @property
    def batch_size(self):
        return self._batch_size

    @property
    def device(self):
        return "cuda"

    def tok_encode(self, string: str):
        return self.tokenizer.encode(string, add_special_tokens=False)

    def tok_decode(self, tokens):
        return self.tokenizer.decode(tokens)

    def _model_call(self, inps):
        """
        inps: a torch tensor of shape [batch, sequence]
        the size of sequence may vary from call to call

        returns: a torch tensor of shape [batch, sequence, vocab] with the
        logits returned from the model
        """
        with torch.no_grad():
            if isinstance(self.model, transformers.models.t5.modeling_t5.T5ForConditionalGeneration):
                decoder_start_token_id = self.model.generation_config.decoder_start_token_id
                dec_inps = torch.cat([torch.tensor(decoder_start_token_id,).tile(len(inps), 1).to(inps),inps,],dim=1,)
             
                kwargs = {"decoder_input_ids": dec_inps,}
            else:
                kwargs = {}
            out = self.model(inps, **kwargs)[0]
            if self.model.config.model_type == "opt":  # there are a few extra tokens in opt, which we should omit
                return out[:, :, :50257]
            else:
                return out  # [:, :, :self.tokenizer.vocab_size]

    def _model_generate(self, context, max_length, eos_token_id):
        return self.model.generate(
            context,
            max_length=max_length,
            eos_token_id=eos_token_id,
            do_sample=False
        )


auto_quant/utils/lm_eval_adaptor.py
qwopqwop200/AutoQuant
from typing import Union, Optional

import torch
import torch.nn as nn
from torch import LongTensor
import transformers
from accelerate import cpu_offload_with_hook
from accelerate.hooks import add_hook_to_module, AlignDevicesHook
from accelerate.utils.modeling import find_tied_parameters, retie_parameters

def get_device(obj):
    if isinstance(obj, torch.Tensor):
        return obj.device.__str__()
    return next(obj.parameters()).device.__str__()

def move_to_device(obj: Union[torch.Tensor, nn.Module], device: torch.device):
    if get_device(obj) != device:
        obj = obj.to(device)
    return obj
    
def get_named_module(module, layers=None, name=''):
    if not layers:
        layers = (transformers.pytorch_utils.Conv1D, nn.Conv2d, nn.Linear)
    if isinstance(module,layers):
        return {name: module}
    res = {}
    for name1, child in module.named_children():
        res.update(get_named_module(child, layers=layers, name=name + '.' + name1 if name != '' else name1))
    return res

def get_op_name(module, op):
    # get the name of the op relative to the module
    for name, m in module.named_modules():
        if m is op:
            return name
    raise ValueError(f"Cannot find op {op} in module {module}")

def get_op_by_name(module, op_name):
    # get the op by its name relative to the module
    for name, m in module.named_modules():
        if name == op_name:
            return m
    raise ValueError(f"Cannot find op {op_name} in module {module}")

def get_module_by_name_prefix(model, module_name: str):
    for name, module in model.named_modules():
        if name.startswith(module_name):
            return module

def get_module_by_name_suffix(model, module_name: str):
    for name, module in model.named_modules():
        if name.endswith(module_name):
            return module

def set_op_by_name(layer, name, new_module):
    levels = name.split('.')
    if len(levels) > 1:
        mod_ = layer
        for l_idx in range(len(levels)-1):
            if levels[l_idx].isdigit():
                mod_ = mod_[int(levels[l_idx])]
            else:
                mod_ = getattr(mod_, levels[l_idx])
        setattr(mod_, levels[-1], new_module)
    else:
        setattr(layer, name, new_module)

def append_str_prefix(x, prefix):
    if isinstance(x, str):
        return prefix + x
    elif isinstance(x, tuple):
        return tuple([append_str_prefix(y, prefix) for y in x])
    elif isinstance(x, list):
        return [append_str_prefix(y, prefix) for y in x]
    else:
        return x
        
def simple_dispatch_model(model, device_map):
    if "" in device_map:
        d = device_map[""]
        model = model.to(torch.device(d))
        model.hf_device_map = device_map
        return model

    tied_params = find_tied_parameters(model)
    if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
        main_device = "cpu"
    else:
        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]

    cpu_offload_group = [(n, d) for n, d in device_map.items() if d == "cpu"]
    prev_hook = None
    for idx, (n, d) in enumerate(cpu_offload_group):
        m = get_module_by_name_suffix(model, n)
        _, prev_hook = cpu_offload_with_hook(m, execution_device=main_device, prev_module_hook=prev_hook)
    # set first cpu offload module's prev_module_hook to the last cpu offload module's hook
    if len(cpu_offload_group) > 1:
        get_module_by_name_suffix(model, cpu_offload_group[0][0])._hf_hook.prev_module_hook = prev_hook

    for n, d in device_map.items():
        m = get_module_by_name_suffix(model, n)
        if d != "cpu":
            d = torch.device(d)
            hook = AlignDevicesHook(d, io_same_device=True, place_submodules=True)
            add_hook_to_module(m, hook)
    retie_parameters(model, tied_params)
    model.hf_device_map = device_map
    return model
auto_quant/utils/module.py
qwopqwop200/AutoQuant
import os
import torch
from pathlib import Path
from setuptools import setup, find_packages
from distutils.sysconfig import get_python_lib
from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CUDAExtension

os.environ["CC"] = "g++"
os.environ["CXX"] = "g++"

common_setup_kwargs = {
    "version": "0.0.3",
    "name": "auto_quant",
    "author": "qwopqwop",
    "license": "MIT",
    "python_requires": ">=3.8.0",
    "description": "An easy-to-use LLMs quantization package",
    "long_description": (Path(__file__).parent / "README.md").read_text(encoding="UTF-8"),
    "long_description_content_type": "text/markdown",
    "url": "https://github.com/qwopqwop200/AutoQuant",
    "keywords": ["autoquant", "quantization", "transformers"],
    "platforms": ["linux", "windows"],
    "classifiers": [
        "Environment :: GPU :: NVIDIA CUDA :: 11.8",
        "Environment :: GPU :: NVIDIA CUDA :: 12",
        "License :: OSI Approved :: MIT License",
        "Natural Language :: English",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: C++",
    ]
}

requirements = [
    "torch>=2.0.0",
    "transformers>=4.32.0",
    "tokenizers>=0.12.1",
    "accelerate",
    "sentencepiece",
    "lm_eval",
]

include_dirs = []

conda_cuda_include_dir = os.path.join(get_python_lib(), "nvidia/cuda_runtime/include")
if os.path.isdir(conda_cuda_include_dir):
    include_dirs.append(conda_cuda_include_dir)

def check_dependencies():
    if CUDA_HOME is None:
        raise RuntimeError(
            f"Cannot find CUDA_HOME. CUDA must be available to build the package.")

def get_compute_capabilities():
    # Collect the compute capabilities of all available GPUs.
    compute_capabilities = set()
    for i in range(torch.cuda.device_count()):
        major, minor = torch.cuda.get_device_capability(i)
        if major < 8:
            raise RuntimeError("GPUs with compute capability less than 8.0 are not supported.")
        compute_capabilities.add(major * 10 + minor)

    # figure out compute capability
    compute_capabilities = {80, 86, 89, 90}

    capability_flags = []
    for cap in compute_capabilities:
        capability_flags += ["-gencode", f"arch=compute_{cap},code=sm_{cap}"]

    return capability_flags

check_dependencies()
arch_flags = get_compute_capabilities()

if os.name == "nt":
    # Relaxed args on Windows
    extra_compile_args={
        "nvcc": arch_flags
    }
else:
    extra_compile_args={
        "cxx": ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17"],
        "nvcc": ["-O3", "-std=c++17"] + arch_flags
    }

extensions = [
    CUDAExtension(
        "exllama_kernels",
        [
            "auto_quant/kernels/exllama/exllama_ext.cpp",
            "auto_quant/kernels/exllama/cuda_buffers.cu",
            "auto_quant/kernels/exllama/cuda_func/column_remap.cu",
            "auto_quant/kernels/exllama/cuda_func/q4_matmul.cu",
            "auto_quant/kernels/exllama/cuda_func/q4_matrix.cu"
        ], extra_compile_args=extra_compile_args
    )
]

additional_setup_kwargs = {
    "ext_modules": extensions,
    "cmdclass": {'build_ext': BuildExtension}
}

common_setup_kwargs.update(additional_setup_kwargs)

setup(
    packages=find_packages(),
    install_requires=requirements,
    include_dirs=include_dirs,
    **common_setup_kwargs
)
setup.py
qwopqwop200/AutoQuant