Add files using upload-large-folder tool

55c92b3 verified 11 months ago

24.6 kB

from pathlib import Path
from typing import List, Literal, Optional, Tuple, Union

import torch
import transformers
from packaging import version
from peft import PeftModel
from peft import __version__ as PEFT_VERSION
from transformers import BitsAndBytesConfig

from lm_eval.models.huggingface import HFLM
from lm_eval import utils
from lm_eval.api.registry import register_model
from lm_eval.models.utils import get_dtype
import re
from pathlib import Path

from outliers.functional.quantization import quantize_blockwise

eval_logger = utils.eval_logger


SUPER_WEIGHTS_MAP = {
    "Mistral-7B-v0.1": [(1, 2070, 7310)],
    "llama-7B": [(2, 3968, 7003)],
    "llama-13B": [(2, 2231, 2278), (2, 2231, 6939)],
    "llama-30B": [(3, 5633, 12817), (3, 5633, 17439), (10, 5633, 14386)],
    "Meta-Llama-3-8B": [(1, 788, 2427), (1, 1384, 2427), (1, 4062, 2427)],
    "OLMo-1B-0724-hf": [(1, 1764, 1710), (2, 1764, 8041)],
    "OLMo-7B-0724-hf": [(1, 269, 7467), (2, 269, 8275), (7, 269, 453), (24, 269, 2300)],
    "Phi-3-mini-4k-instruct": [(2, 525, 808), (2, 1693, 808), (2, 1113, 808), (4, 525, 2723),  (4, 1113, 2723), (4, 1693, 2723)],
}


def get_layer_number(layer_name):
    # Define the pattern to extract the layer number
    match = re.search(r"layers\.(\d+)\.", layer_name)
    if match:
        return int(match.group(1))
    else:
        return None


def get_weight_type(name):
    for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]:
        if linear in name:
            return linear
    return "other"


def pack_4bit_to_int8(quantized_4bit_weights):
    # Ensure the quantized weights are uint8
    quantized_4bit_weights = quantized_4bit_weights.to(torch.uint8)
    
    # Reshape to ensure even number of elements per row
    if quantized_4bit_weights.size(1) % 2 != 0:
        quantized_4bit_weights = torch.cat((quantized_4bit_weights, torch.zeros((quantized_4bit_weights.size(0), 1), dtype=torch.uint8)), dim=1)
    
    # Packing two 4-bit values into one int8
    packed_weights = (quantized_4bit_weights[:, ::2] << 4) | (quantized_4bit_weights[:, 1::2] & 0xF)
    
    return packed_weights


def _get_accelerate_args(
    device_map_option: Optional[str] = "auto",
    max_memory_per_gpu: Optional[Union[int, str]] = None,
    max_cpu_memory: Optional[Union[int, str]] = None,
    offload_folder: Optional[str] = "./offload",
    gpus: Optional[int] = None,
) -> dict:
    """Returns the kwargs needed to apply `accelerate` in `AutoModel.from_pretrained`."""
    max_memory = {}
    if max_memory_per_gpu is not None:
        max_memory_per_gpu_map = {
            device_idx: max_memory_per_gpu for device_idx in range(gpus)
        }
        max_memory.update(max_memory_per_gpu_map)
    if max_cpu_memory is not None:
        max_memory["cpu"] = max_cpu_memory

    args = {}
    if max_memory:
        args["max_memory"] = max_memory
    args["device_map"] = device_map_option
    args["offload_folder"] = offload_folder
    return args


@register_model("hf-outlier")
class HFOutlierLM(HFLM):
    def _create_model(
        self,
        pretrained: str,
        revision: Optional[str] = "main",
        dtype: Optional[Union[str, torch.dtype]] = "auto",
        trust_remote_code: Optional[bool] = False,
        # arguments used for splitting a model across GPUs naively.
        # only used if `parallelize=True`.
        # (accelerate naive PP (device_map) options)
        parallelize: Optional[bool] = False,
        gpus: Optional[int] = None,
        device_map_option: Optional[str] = "auto",
        max_memory_per_gpu: Optional[Union[int, str]] = None,
        max_cpu_memory: Optional[Union[int, str]] = None,
        offload_folder: Optional[str] = "./offload",
        # PEFT, delta weights and quantization options
        peft: Optional[str] = None,
        delta: Optional[str] = None,
        autogptq: Optional[Union[bool, str]] = False,
        outlier_method: Optional[str] = None, # added argument for outlier experiments
        manual_quantize: Optional[str] = None,
        restore_and_scale_GO: Optional[Union[bool, float]] = False,
        bnb_quantize_args: Optional[str] = None,
        clip_and_save_args: Optional[str] = None,
        load_quantized_args: Optional[str] = None,
        adjust_attn_args: Optional[str] = None,
        **kwargs,
    ) -> None:
        """
        Initializes an HF or HF-compatible PreTrainedModel from scratch
        inside HFLM, using the kwargs passed into self.__init__().
        Also handles functionality such as AutoGPTQ usage and PEFT wrapping.
        For future similar extensions to AutoGPTQ that are not core to HF's ecosystem,
        (such as PyTorch models that are nearly, but not quite, fully mirroring
        HF's public interface relied on in this HFLM class)
        please consider subclassing HFLM and overriding this and other methods as needed.
        """

        model_kwargs = kwargs if kwargs else {}

        if parallelize:
            model_kwargs.update(
                _get_accelerate_args(
                    device_map_option,  # TODO: phase out device_map_option?
                    max_memory_per_gpu,
                    max_cpu_memory,
                    offload_folder,
                    gpus,
                )
            )
        elif "device_map" not in model_kwargs:
            # set a device_map to initialize model on the right GPU.
            # this is needed because it seems that the default behavior
            # for quantized models now seems to be device_map="auto"
            # which breaks data-parallel mode.
            if hasattr(self, "accelerator"):
                model_kwargs.update({"device_map": {"": f"{self.accelerator.device}"}})
            else:
                model_kwargs.update({"device_map": {"": str(self.device)}})

        if not autogptq:
            if model_kwargs.get("load_in_4bit", None):
                assert (
                    transformers.__version__ >= "4.30.0"
                ), "load_in_4bit requires transformers >= 4.30.0"
            if transformers.__version__ >= "4.30.0":
                if model_kwargs.get("load_in_4bit", None):
                    if model_kwargs.get("bnb_4bit_compute_dtype", None):
                        model_kwargs["bnb_4bit_compute_dtype"] = get_dtype(
                            model_kwargs["bnb_4bit_compute_dtype"]
                        )         


            if bnb_quantize_args:
                load_in_4bit, bnb_4bit_quant_type, blocksize, clip_percentage = bnb_quantize_args.split('_') # True_fp4_512_1e-3
                load_in_4bit = True if load_in_4bit == "True" else False
                blocksize = int(blocksize)
                print("bnb args", load_in_4bit, blocksize, bnb_4bit_quant_type)
                quantization_config = BitsAndBytesConfig(
                    load_in_4bit=load_in_4bit,
                    bnb_4bit_quant_type=bnb_4bit_quant_type,
                    blocksize=blocksize
                )
                model_kwargs["load_in_4bit"] = False
                model_kwargs["quantization_config"] = quantization_config    



            self._model = self.AUTO_MODEL_CLASS.from_pretrained(
                pretrained,
                revision=revision,
                torch_dtype=get_dtype(dtype),
                trust_remote_code=trust_remote_code,
                **model_kwargs,
            )
            # Record SW value
            self.record_GO(pretrained)   
            if adjust_attn_args: # Now is working for Llama-7B only
                import math
                from outliers.functional.utils import ScaledLlamaSdpaAttention
                tau = float(adjust_attn_args)
                scale = 1 / (math.sqrt(self._model.model.layers[3].self_attn.head_dim) * tau) 

                def remove_massive_activation_llama(module, input, output):
                    # print(get_module_name(model, module))
                    hidden_states = output[0][0]
                    # Modify the hidden states here

                    hidden_states[0, 3968] *= tau
                    modified_output = hidden_states
                    output[0][0] = modified_output
                    return output


                self._model.model.layers[3].self_attn = ScaledLlamaSdpaAttention(config=self._model.config, layer_idx=3)

                all_hooks = []
                remove_massive_activation_hook = self._model.model.layers[2].register_forward_hook(remove_massive_activation_llama)
                all_hooks.append(remove_massive_activation_hook)


                self._model = self._model.to('cuda').to(torch.bfloat16) # hard code dtype for now. might need to adjust to arguments


            if manual_quantize:
                # Parse the quantization method and parameters from the string
                quantize_method, bits, block_size_arg, clip_method, clip_threshold, scale_shift, use_normal_float = manual_quantize.split('_') # minmax_4_128_bp_1e-4
                bits = int(bits)

                block_size_arg = int(block_size_arg) if block_size_arg not in ["channel", "tensor"] else block_size_arg
                clip_threshold = float(clip_threshold) if clip_threshold != "None" else None 
                clip_method_map = {"bp": "block_percentage", "tp": "tensor_percentage", "z": "zscore", "iqr": "iqr",  "no": "no"}
                clip_method = clip_method_map[clip_method]
                scale_shift = True if scale_shift == "True" else False
                use_normal_float = True if use_normal_float == "True" else False

                if quantize_method == "minmax":
                    print(f"Running {bits} bits manual quantization ...")
                    for name, param in self._model.model.named_parameters():
                        if not name.endswith("weight"):
                            continue
                        if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
                            continue
                        weight = param.data
                        shape = weight.shape
                        if block_size_arg == "channel":
                            block_size = shape[-1]
                        elif block_size_arg == "tensor":
                            block_size = weight.numel()
                        elif isinstance(block_size_arg, int):
                            block_size = block_size_arg
                        else:
                            raise ValueError("Block size must be int or 'tensor' or 'channel'")
                        quantized_weight, _num_outliers = quantize_blockwise(weight, bits, block_size, "no", 0, scale_shift, use_normal_float)
                        param.data = quantized_weight    
                
                elif quantize_method == "clip":
                    # clip outliers but restore super outliers
                    # quantize with min max
                    num_outliers = []
                    original_value_ranges = []
                    clipped_value_ranges = []

                    print(f"Running {bits} bits manual quantization ...")
                    print(f"Clipping parameters", clip_method, clip_threshold)
                    print(f"Restore GO?", restore_and_scale_GO)
                    for name, param in self._model.model.named_parameters():
                        if not name.endswith("weight"):
                            continue
                        if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
                            continue

                        weight = param.data
                        shape = weight.shape
                        if block_size_arg == "channel":
                            block_size = shape[-1]
                        elif block_size_arg == "tensor":
                            block_size = weight.numel()
                        elif isinstance(block_size_arg, int):
                            block_size = block_size_arg
                        else:
                            raise ValueError("Block size must be int or 'tensor' or 'channel'") 
                        if weight.numel() % (block_size) != 0:
                            print(name,  weight.numel(), block_size)
                        quantized_weight, _num_outliers = quantize_blockwise(weight, bits, block_size, clip_method, clip_threshold, scale_shift, use_normal_float)
                        param.data = quantized_weight    
                        num_outliers.append(_num_outliers)
                    # print statistics
                    print(f"Number of outliers. Max: {max(num_outliers)}, Min: {min(num_outliers)}, Sum: {sum(num_outliers)}, Mean {sum(num_outliers) / len(num_outliers)}")


            if outlier_method:
                def get_weight_type(model_id, name):
                    if model_id.startswith("facebook/opt"):
                        for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "fc1", "fc2"]:
                            if linear in name:
                                return linear   
                    else:

                        for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]:
                            if linear in name:
                                return linear
                    return "other"
                
                def get_layer_number(layer_name):
                    # Define the pattern to extract the layer number
                    match = re.search(r"layers\.(\d+)\.", layer_name)
                    if match:
                        return int(match.group(1))
                    else:
                        return None
                if outlier_method.startswith("manual_scaling_SO"): # for example, manual_scaling_SO_0.1
                    print("manual scaling SO")
                    scaling_factor = outlier_method.split('_')[-1]
                    scaling_factor = float(scaling_factor)
                    print("Original SO:", self.GO_values)
                    self.restore_GO(pretrained, scaling_factor)

                if outlier_method.startswith("removeSW_restoreSA"):
                    if pretrained == "huggyllama/llama-7B":
                        def hook_fn(module, input, output):
                            SW_channel = 3968
                            SW_row = 7003
                            # Record original hidden states
                            X1 = output.detach().clone()
                            # print(f"Shape of input/output: {output.shape}, {input[0].shape}")
                            
                            # Find the largest magnitude in channel SA and its position
                            channel_SA = X1[..., SW_channel]
                            # print(f"Shape of channel_SA: {channel_SA.shape}")
                            max_magnitude_indices = torch.argmax(torch.abs(channel_SA), dim=1)
                            # print(f"Max magnitude indices: {max_magnitude_indices}")
                            max_magnitudes = torch.gather(channel_SA, 1, max_magnitude_indices.unsqueeze(1)).squeeze(1)

                            
                            # Temporarily set SW to 0
                            weight = module.weight.clone()
                            original_weight = weight.data[SW_channel, SW_row].item()
                            weight.data[SW_channel, SW_row] = 0.0
                            
                            # Recompute hidden states
                            with torch.no_grad():
                                X2 = torch.nn.functional.linear(input[0], weight, module.bias)
                                # print(f"Shape of recomputed X2: {X2.shape}")
                            
                                batch_indices = torch.arange(X2.size(0), device=X2.device)
                                # print("New magnitude on SA position:", X2[batch_indices, max_magnitude_indices, 3968])
                                X2[batch_indices, max_magnitude_indices, SW_channel] = max_magnitudes

                            # Restore the original weight
                            module.weight.data[SW_channel, SW_row] = original_weight
                            return X2

                        # Register the hook to the specific layer
                        self._model.model.layers[2].mlp.down_proj.register_forward_hook(hook_fn)
                        print("Hook registered for removing SW and restoring SA")

                elif outlier_method.startswith("search"):
                    _, criterion, search_threhold, layers, outlier_weight_type = outlier_method.split("_") # "search_percentage_1e-4_0-26_downproj"
                    start_layer, end_layer = layers.split('-')
                    layers = range(int(start_layer), int(end_layer))
                    search_threhold = float(search_threhold)
                    # print(layers, search_threhold, criterion)
                    num_selected_elements = []
                    for name, param in self._model.model.named_parameters():
                        if not name.endswith("weight"):
                            continue
                        if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
                            continue
                        if outlier_weight_type != "all":
                            if outlier_weight_type == "down":
                                outlier_weight_type = "down_proj"
                            weight_type = get_weight_type(pretrained, name)
                            if weight_type != outlier_weight_type:
                                # print("Skip for unmatched weight type")
                                continue
                        if layers is not None:
                            layer_number = get_layer_number(name)
                            if layer_number not in layers:
                                # print("Skip for unmatched layer")
                                continue
                        weight = param.data
                        if criterion == "percentage":
                            num_top_elements = int(weight.numel() * search_threhold)
                            # print("# Total params:", weight.numel(), "# top params", num_top_elements)

                            if num_top_elements < 1: # too few elements to apply on
                                continue 
                            threshold = torch.topk(weight.view(-1).abs(), num_top_elements).values[-1]
                            mask = weight.abs() >= threshold
                            true_indices = mask.nonzero(as_tuple=False)
                            num_selected_elements.append(len(true_indices))
                            # print(name, len(true_indices))
                            
                            weight[mask] = 0.
                        param = torch.nn.Parameter(weight) 
                    if len(num_selected_elements) > 0:
                        print(f"Number of removed elements. Max: {max(num_selected_elements)}, Min: {min(num_selected_elements)}, Sum: {sum(num_selected_elements)}, Mean {sum(num_selected_elements) / len(num_selected_elements)}")


            if restore_and_scale_GO:
                self.restore_GO(pretrained, restore_and_scale_GO)
            else:
                print("Not restoring or scaling GO...")

        else:
            try:
                from auto_gptq import AutoGPTQForCausalLM
            except ModuleNotFoundError:
                raise Exception(
                    "Tried to load auto_gptq, but auto-gptq is not installed ",
                    "please install auto-gptq via pip install lm-eval[gptq] or pip install -e .[gptq]",
                )

            self._model = AutoGPTQForCausalLM.from_quantized(
                pretrained,
                trust_remote_code=trust_remote_code,
                model_basename=None if autogptq is True else Path(autogptq).stem,
                use_safetensors=True
                if autogptq is True
                else autogptq.endswith(".safetensors"),
                **model_kwargs,
            )

        if peft and delta:
            raise ValueError(
                "Cannot use both 'peft' and 'delta' options at the same time."
            )

        if peft:
            if model_kwargs.get("load_in_4bit", None):
                if version.parse(PEFT_VERSION) < version.parse("0.4.0"):
                    raise AssertionError("load_in_4bit requires peft >= 0.4.0")
            if self._model.config.vocab_size != len(self.tokenizer):
                # resize model for LoRAs with added tokens
                self._model.resize_token_embeddings(len(self.tokenizer))
                eval_logger.info(
                    f"Model config indicates vocab_size='{self._model.config.vocab_size}', but found tokenizer with vocab size '{len(self.tokenizer)}'. Resizing model embedding layer..."
                )
            self._model = PeftModel.from_pretrained(
                self._model, peft, revision=revision
            )
        elif delta:
            if autogptq:
                eval_logger.warning(
                    "Delta weights might trigger unexpected behavior when used with AutoGPTQ."
                )
            _model_delta = self.AUTO_MODEL_CLASS.from_pretrained(
                delta,
                revision=revision,
                torch_dtype=get_dtype(dtype),
                trust_remote_code=trust_remote_code,
                **model_kwargs,
            )
            for name, param in self._model.state_dict().items():
                try:
                    param.data += _model_delta.state_dict()[name]
                except KeyError:
                    raise KeyError(f"Delta model is missing weights for layer: {name}")
                except Exception as e:
                    raise RuntimeError(
                        f"Failed to add delta weights to layer {name}. Error: {e}"
                    )

            del _model_delta

        return None


    def record_GO(self, pretrained):
        '''Record GO values for original models'''
        def _record_GO(GO_map, layer, row, col):
            if pretrained in [
                "tiiuae/falcon-7b",
            ]:
                GO_map[(layer, row, col)] = self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col].item()
            else:
                GO_map[(layer, row, col)] = self._model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()

        GO_values = {}
        for model_name, coordinates in SUPER_WEIGHTS_MAP.items():
            if model_name in pretrained:
                for layer, row, col in coordinates:
                    _record_GO(GO_values, layer, row, col)
                break
        self.GO_values = GO_values


    def restore_GO(self, pretrained, scaling_factor):
        if pretrained in [
            "huggyllama/llama-30B", 
            "huggyllama/llama-13B", 
            "huggyllama/llama-7B", 
            "mistralai/Mistral-7B-v0.1",
            "meta/Meta-Llama-3-8B",
            "allenai/OLMo-1B-0724-hf",
            "allenai/OLMo-7B-0724-hf",
            "google/gemma-7b",
            "microsoft/Phi-3-mini-4k-instruct"
            ]:
            if getattr(self._model, "hf_quantizer", None):
                self._model.dequantize()
            for (layer, row, col), value in self.GO_values.items():  
                old_value = self._model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()
                new_value = value * scaling_factor
                self._model.model.layers[layer].mlp.down_proj.weight.data[row, col] = new_value
                print(f"Layer {layer}, Index [{row}, {col}], Old value: {old_value}, New value: {new_value}")

        elif pretrained in [
            "tiiuae/falcon-7b"
            ]:
            if getattr(self._model, "hf_quantizer", None):
                self._model.dequantize()
            for (layer, row, col), value in self.GO_values.items():  
                old_value = self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col].item()
                new_value = value * scaling_factor
                self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col] = new_value
                print(f"Layer {layer}, Index [{row}, {col}], Old value: {old_value}, New value: {new_value}")

	from pathlib import Path
	from typing import List, Literal, Optional, Tuple, Union

	import torch
	import transformers
	from packaging import version
	from peft import PeftModel
	from peft import __version__ as PEFT_VERSION
	from transformers import BitsAndBytesConfig

	from lm_eval.models.huggingface import HFLM
	from lm_eval import utils
	from lm_eval.api.registry import register_model
	from lm_eval.models.utils import get_dtype
	import re
	from pathlib import Path

	from outliers.functional.quantization import quantize_blockwise

	eval_logger = utils.eval_logger


	SUPER_WEIGHTS_MAP = {
	"Mistral-7B-v0.1": [(1, 2070, 7310)],
	"llama-7B": [(2, 3968, 7003)],
	"llama-13B": [(2, 2231, 2278), (2, 2231, 6939)],
	"llama-30B": [(3, 5633, 12817), (3, 5633, 17439), (10, 5633, 14386)],
	"Meta-Llama-3-8B": [(1, 788, 2427), (1, 1384, 2427), (1, 4062, 2427)],
	"OLMo-1B-0724-hf": [(1, 1764, 1710), (2, 1764, 8041)],
	"OLMo-7B-0724-hf": [(1, 269, 7467), (2, 269, 8275), (7, 269, 453), (24, 269, 2300)],
	"Phi-3-mini-4k-instruct": [(2, 525, 808), (2, 1693, 808), (2, 1113, 808), (4, 525, 2723), (4, 1113, 2723), (4, 1693, 2723)],
	}


	def get_layer_number(layer_name):
	# Define the pattern to extract the layer number
	match = re.search(r"layers\.(\d+)\.", layer_name)
	if match:
	return int(match.group(1))
	else:
	return None


	def get_weight_type(name):
	for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]:
	if linear in name:
	return linear
	return "other"


	def pack_4bit_to_int8(quantized_4bit_weights):
	# Ensure the quantized weights are uint8
	quantized_4bit_weights = quantized_4bit_weights.to(torch.uint8)

	# Reshape to ensure even number of elements per row
	if quantized_4bit_weights.size(1) % 2 != 0:
	quantized_4bit_weights = torch.cat((quantized_4bit_weights, torch.zeros((quantized_4bit_weights.size(0), 1), dtype=torch.uint8)), dim=1)

	# Packing two 4-bit values into one int8
	packed_weights = (quantized_4bit_weights[:, ::2] << 4) \| (quantized_4bit_weights[:, 1::2] & 0xF)

	return packed_weights


	def _get_accelerate_args(
	device_map_option: Optional[str] = "auto",
	max_memory_per_gpu: Optional[Union[int, str]] = None,
	max_cpu_memory: Optional[Union[int, str]] = None,
	offload_folder: Optional[str] = "./offload",
	gpus: Optional[int] = None,
	) -> dict:
	"""Returns the kwargs needed to apply `accelerate` in `AutoModel.from_pretrained`."""
	max_memory = {}
	if max_memory_per_gpu is not None:
	max_memory_per_gpu_map = {
	device_idx: max_memory_per_gpu for device_idx in range(gpus)
	}
	max_memory.update(max_memory_per_gpu_map)
	if max_cpu_memory is not None:
	max_memory["cpu"] = max_cpu_memory

	args = {}
	if max_memory:
	args["max_memory"] = max_memory
	args["device_map"] = device_map_option
	args["offload_folder"] = offload_folder
	return args


	@register_model("hf-outlier")
	class HFOutlierLM(HFLM):
	def _create_model(
	self,
	pretrained: str,
	revision: Optional[str] = "main",
	dtype: Optional[Union[str, torch.dtype]] = "auto",
	trust_remote_code: Optional[bool] = False,
	# arguments used for splitting a model across GPUs naively.
	# only used if `parallelize=True`.
	# (accelerate naive PP (device_map) options)
	parallelize: Optional[bool] = False,
	gpus: Optional[int] = None,
	device_map_option: Optional[str] = "auto",
	max_memory_per_gpu: Optional[Union[int, str]] = None,
	max_cpu_memory: Optional[Union[int, str]] = None,
	offload_folder: Optional[str] = "./offload",
	# PEFT, delta weights and quantization options
	peft: Optional[str] = None,
	delta: Optional[str] = None,
	autogptq: Optional[Union[bool, str]] = False,
	outlier_method: Optional[str] = None, # added argument for outlier experiments
	manual_quantize: Optional[str] = None,
	restore_and_scale_GO: Optional[Union[bool, float]] = False,
	bnb_quantize_args: Optional[str] = None,
	clip_and_save_args: Optional[str] = None,
	load_quantized_args: Optional[str] = None,
	adjust_attn_args: Optional[str] = None,
	**kwargs,
	) -> None:
	"""
	Initializes an HF or HF-compatible PreTrainedModel from scratch
	inside HFLM, using the kwargs passed into self.__init__().

	Also handles functionality such as AutoGPTQ usage and PEFT wrapping.

	For future similar extensions to AutoGPTQ that are not core to HF's ecosystem,
	(such as PyTorch models that are nearly, but not quite, fully mirroring
	HF's public interface relied on in this HFLM class)
	please consider subclassing HFLM and overriding this and other methods as needed.
	"""

	model_kwargs = kwargs if kwargs else {}

	if parallelize:
	model_kwargs.update(
	_get_accelerate_args(
	device_map_option, # TODO: phase out device_map_option?
	max_memory_per_gpu,
	max_cpu_memory,
	offload_folder,
	gpus,
	)
	)
	elif "device_map" not in model_kwargs:
	# set a device_map to initialize model on the right GPU.
	# this is needed because it seems that the default behavior
	# for quantized models now seems to be device_map="auto"
	# which breaks data-parallel mode.
	if hasattr(self, "accelerator"):
	model_kwargs.update({"device_map": {"": f"{self.accelerator.device}"}})
	else:
	model_kwargs.update({"device_map": {"": str(self.device)}})

	if not autogptq:
	if model_kwargs.get("load_in_4bit", None):
	assert (
	transformers.__version__ >= "4.30.0"
	), "load_in_4bit requires transformers >= 4.30.0"
	if transformers.__version__ >= "4.30.0":
	if model_kwargs.get("load_in_4bit", None):
	if model_kwargs.get("bnb_4bit_compute_dtype", None):
	model_kwargs["bnb_4bit_compute_dtype"] = get_dtype(
	model_kwargs["bnb_4bit_compute_dtype"]
	)


	if bnb_quantize_args:
	load_in_4bit, bnb_4bit_quant_type, blocksize, clip_percentage = bnb_quantize_args.split('_') # True_fp4_512_1e-3
	load_in_4bit = True if load_in_4bit == "True" else False
	blocksize = int(blocksize)
	print("bnb args", load_in_4bit, blocksize, bnb_4bit_quant_type)
	quantization_config = BitsAndBytesConfig(
	load_in_4bit=load_in_4bit,
	bnb_4bit_quant_type=bnb_4bit_quant_type,
	blocksize=blocksize
	)
	model_kwargs["load_in_4bit"] = False
	model_kwargs["quantization_config"] = quantization_config



	self._model = self.AUTO_MODEL_CLASS.from_pretrained(
	pretrained,
	revision=revision,
	torch_dtype=get_dtype(dtype),
	trust_remote_code=trust_remote_code,
	**model_kwargs,
	)
	# Record SW value
	self.record_GO(pretrained)
	if adjust_attn_args: # Now is working for Llama-7B only
	import math
	from outliers.functional.utils import ScaledLlamaSdpaAttention
	tau = float(adjust_attn_args)
	scale = 1 / (math.sqrt(self._model.model.layers[3].self_attn.head_dim) * tau)

	def remove_massive_activation_llama(module, input, output):
	# print(get_module_name(model, module))
	hidden_states = output[0][0]
	# Modify the hidden states here

	hidden_states[0, 3968] *= tau
	modified_output = hidden_states
	output[0][0] = modified_output
	return output


	self._model.model.layers[3].self_attn = ScaledLlamaSdpaAttention(config=self._model.config, layer_idx=3)

	all_hooks = []
	remove_massive_activation_hook = self._model.model.layers[2].register_forward_hook(remove_massive_activation_llama)
	all_hooks.append(remove_massive_activation_hook)


	self._model = self._model.to('cuda').to(torch.bfloat16) # hard code dtype for now. might need to adjust to arguments


	if manual_quantize:
	# Parse the quantization method and parameters from the string
	quantize_method, bits, block_size_arg, clip_method, clip_threshold, scale_shift, use_normal_float = manual_quantize.split('_') # minmax_4_128_bp_1e-4
	bits = int(bits)

	block_size_arg = int(block_size_arg) if block_size_arg not in ["channel", "tensor"] else block_size_arg
	clip_threshold = float(clip_threshold) if clip_threshold != "None" else None
	clip_method_map = {"bp": "block_percentage", "tp": "tensor_percentage", "z": "zscore", "iqr": "iqr", "no": "no"}
	clip_method = clip_method_map[clip_method]
	scale_shift = True if scale_shift == "True" else False
	use_normal_float = True if use_normal_float == "True" else False

	if quantize_method == "minmax":
	print(f"Running {bits} bits manual quantization ...")
	for name, param in self._model.model.named_parameters():
	if not name.endswith("weight"):
	continue
	if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
	continue
	weight = param.data
	shape = weight.shape
	if block_size_arg == "channel":
	block_size = shape[-1]
	elif block_size_arg == "tensor":
	block_size = weight.numel()
	elif isinstance(block_size_arg, int):
	block_size = block_size_arg
	else:
	raise ValueError("Block size must be int or 'tensor' or 'channel'")
	quantized_weight, _num_outliers = quantize_blockwise(weight, bits, block_size, "no", 0, scale_shift, use_normal_float)
	param.data = quantized_weight

	elif quantize_method == "clip":
	# clip outliers but restore super outliers
	# quantize with min max
	num_outliers = []
	original_value_ranges = []
	clipped_value_ranges = []

	print(f"Running {bits} bits manual quantization ...")
	print(f"Clipping parameters", clip_method, clip_threshold)
	print(f"Restore GO?", restore_and_scale_GO)
	for name, param in self._model.model.named_parameters():
	if not name.endswith("weight"):
	continue
	if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
	continue

	weight = param.data
	shape = weight.shape
	if block_size_arg == "channel":
	block_size = shape[-1]
	elif block_size_arg == "tensor":
	block_size = weight.numel()
	elif isinstance(block_size_arg, int):
	block_size = block_size_arg
	else:
	raise ValueError("Block size must be int or 'tensor' or 'channel'")
	if weight.numel() % (block_size) != 0:
	print(name, weight.numel(), block_size)
	quantized_weight, _num_outliers = quantize_blockwise(weight, bits, block_size, clip_method, clip_threshold, scale_shift, use_normal_float)
	param.data = quantized_weight
	num_outliers.append(_num_outliers)
	# print statistics
	print(f"Number of outliers. Max: {max(num_outliers)}, Min: {min(num_outliers)}, Sum: {sum(num_outliers)}, Mean {sum(num_outliers) / len(num_outliers)}")


	if outlier_method:
	def get_weight_type(model_id, name):
	if model_id.startswith("facebook/opt"):
	for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "fc1", "fc2"]:
	if linear in name:
	return linear
	else:

	for linear in ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]:
	if linear in name:
	return linear
	return "other"

	def get_layer_number(layer_name):
	# Define the pattern to extract the layer number
	match = re.search(r"layers\.(\d+)\.", layer_name)
	if match:
	return int(match.group(1))
	else:
	return None
	if outlier_method.startswith("manual_scaling_SO"): # for example, manual_scaling_SO_0.1
	print("manual scaling SO")
	scaling_factor = outlier_method.split('_')[-1]
	scaling_factor = float(scaling_factor)
	print("Original SO:", self.GO_values)
	self.restore_GO(pretrained, scaling_factor)

	if outlier_method.startswith("removeSW_restoreSA"):
	if pretrained == "huggyllama/llama-7B":
	def hook_fn(module, input, output):
	SW_channel = 3968
	SW_row = 7003
	# Record original hidden states
	X1 = output.detach().clone()
	# print(f"Shape of input/output: {output.shape}, {input[0].shape}")

	# Find the largest magnitude in channel SA and its position
	channel_SA = X1[..., SW_channel]
	# print(f"Shape of channel_SA: {channel_SA.shape}")
	max_magnitude_indices = torch.argmax(torch.abs(channel_SA), dim=1)
	# print(f"Max magnitude indices: {max_magnitude_indices}")
	max_magnitudes = torch.gather(channel_SA, 1, max_magnitude_indices.unsqueeze(1)).squeeze(1)


	# Temporarily set SW to 0
	weight = module.weight.clone()
	original_weight = weight.data[SW_channel, SW_row].item()
	weight.data[SW_channel, SW_row] = 0.0

	# Recompute hidden states
	with torch.no_grad():
	X2 = torch.nn.functional.linear(input[0], weight, module.bias)
	# print(f"Shape of recomputed X2: {X2.shape}")

	batch_indices = torch.arange(X2.size(0), device=X2.device)
	# print("New magnitude on SA position:", X2[batch_indices, max_magnitude_indices, 3968])
	X2[batch_indices, max_magnitude_indices, SW_channel] = max_magnitudes

	# Restore the original weight
	module.weight.data[SW_channel, SW_row] = original_weight
	return X2

	# Register the hook to the specific layer
	self._model.model.layers[2].mlp.down_proj.register_forward_hook(hook_fn)
	print("Hook registered for removing SW and restoring SA")

	elif outlier_method.startswith("search"):
	_, criterion, search_threhold, layers, outlier_weight_type = outlier_method.split("_") # "search_percentage_1e-4_0-26_downproj"
	start_layer, end_layer = layers.split('-')
	layers = range(int(start_layer), int(end_layer))
	search_threhold = float(search_threhold)
	# print(layers, search_threhold, criterion)
	num_selected_elements = []
	for name, param in self._model.model.named_parameters():
	if not name.endswith("weight"):
	continue
	if "layernorm" in name or "norm" in name or "embed_tokens" in name or "lm_head" in name:
	continue
	if outlier_weight_type != "all":
	if outlier_weight_type == "down":
	outlier_weight_type = "down_proj"
	weight_type = get_weight_type(pretrained, name)
	if weight_type != outlier_weight_type:
	# print("Skip for unmatched weight type")
	continue
	if layers is not None:
	layer_number = get_layer_number(name)
	if layer_number not in layers:
	# print("Skip for unmatched layer")
	continue
	weight = param.data
	if criterion == "percentage":
	num_top_elements = int(weight.numel() * search_threhold)
	# print("# Total params:", weight.numel(), "# top params", num_top_elements)

	if num_top_elements < 1: # too few elements to apply on
	continue
	threshold = torch.topk(weight.view(-1).abs(), num_top_elements).values[-1]
	mask = weight.abs() >= threshold
	true_indices = mask.nonzero(as_tuple=False)
	num_selected_elements.append(len(true_indices))
	# print(name, len(true_indices))

	weight[mask] = 0.
	param = torch.nn.Parameter(weight)
	if len(num_selected_elements) > 0:
	print(f"Number of removed elements. Max: {max(num_selected_elements)}, Min: {min(num_selected_elements)}, Sum: {sum(num_selected_elements)}, Mean {sum(num_selected_elements) / len(num_selected_elements)}")


	if restore_and_scale_GO:
	self.restore_GO(pretrained, restore_and_scale_GO)
	else:
	print("Not restoring or scaling GO...")

	else:
	try:
	from auto_gptq import AutoGPTQForCausalLM
	except ModuleNotFoundError:
	raise Exception(
	"Tried to load auto_gptq, but auto-gptq is not installed ",
	"please install auto-gptq via pip install lm-eval[gptq] or pip install -e .[gptq]",
	)

	self._model = AutoGPTQForCausalLM.from_quantized(
	pretrained,
	trust_remote_code=trust_remote_code,
	model_basename=None if autogptq is True else Path(autogptq).stem,
	use_safetensors=True
	if autogptq is True
	else autogptq.endswith(".safetensors"),
	**model_kwargs,
	)

	if peft and delta:
	raise ValueError(
	"Cannot use both 'peft' and 'delta' options at the same time."
	)

	if peft:
	if model_kwargs.get("load_in_4bit", None):
	if version.parse(PEFT_VERSION) < version.parse("0.4.0"):
	raise AssertionError("load_in_4bit requires peft >= 0.4.0")
	if self._model.config.vocab_size != len(self.tokenizer):
	# resize model for LoRAs with added tokens
	self._model.resize_token_embeddings(len(self.tokenizer))
	eval_logger.info(
	f"Model config indicates vocab_size='{self._model.config.vocab_size}', but found tokenizer with vocab size '{len(self.tokenizer)}'. Resizing model embedding layer..."
	)
	self._model = PeftModel.from_pretrained(
	self._model, peft, revision=revision
	)
	elif delta:
	if autogptq:
	eval_logger.warning(
	"Delta weights might trigger unexpected behavior when used with AutoGPTQ."
	)
	_model_delta = self.AUTO_MODEL_CLASS.from_pretrained(
	delta,
	revision=revision,
	torch_dtype=get_dtype(dtype),
	trust_remote_code=trust_remote_code,
	**model_kwargs,
	)
	for name, param in self._model.state_dict().items():
	try:
	param.data += _model_delta.state_dict()[name]
	except KeyError:
	raise KeyError(f"Delta model is missing weights for layer: {name}")
	except Exception as e:
	raise RuntimeError(
	f"Failed to add delta weights to layer {name}. Error: {e}"
	)

	del _model_delta

	return None


	def record_GO(self, pretrained):
	'''Record GO values for original models'''
	def _record_GO(GO_map, layer, row, col):
	if pretrained in [
	"tiiuae/falcon-7b",
	]:
	GO_map[(layer, row, col)] = self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col].item()
	else:
	GO_map[(layer, row, col)] = self._model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()

	GO_values = {}
	for model_name, coordinates in SUPER_WEIGHTS_MAP.items():
	if model_name in pretrained:
	for layer, row, col in coordinates:
	_record_GO(GO_values, layer, row, col)
	break
	self.GO_values = GO_values


	def restore_GO(self, pretrained, scaling_factor):
	if pretrained in [
	"huggyllama/llama-30B",
	"huggyllama/llama-13B",
	"huggyllama/llama-7B",
	"mistralai/Mistral-7B-v0.1",
	"meta/Meta-Llama-3-8B",
	"allenai/OLMo-1B-0724-hf",
	"allenai/OLMo-7B-0724-hf",
	"google/gemma-7b",
	"microsoft/Phi-3-mini-4k-instruct"
	]:
	if getattr(self._model, "hf_quantizer", None):
	self._model.dequantize()
	for (layer, row, col), value in self.GO_values.items():
	old_value = self._model.model.layers[layer].mlp.down_proj.weight.data[row, col].item()
	new_value = value * scaling_factor
	self._model.model.layers[layer].mlp.down_proj.weight.data[row, col] = new_value
	print(f"Layer {layer}, Index [{row}, {col}], Old value: {old_value}, New value: {new_value}")

	elif pretrained in [
	"tiiuae/falcon-7b"
	]:
	if getattr(self._model, "hf_quantizer", None):
	self._model.dequantize()
	for (layer, row, col), value in self.GO_values.items():
	old_value = self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col].item()
	new_value = value * scaling_factor
	self._model.transformer.h[layer].mlp.dense_4h_to_h.weight.data[row, col] = new_value
	print(f"Layer {layer}, Index [{row}, {col}], Old value: {old_value}, New value: {new_value}")