Upload 14 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+asserts/main-fig.png filter=lfs diff=lfs merge=lfs -text
+asserts/table-1.png filter=lfs diff=lfs merge=lfs -text
+asserts/table-2.png filter=lfs diff=lfs merge=lfs -text
+asserts/table-3.png filter=lfs diff=lfs merge=lfs -text
+asserts/table-4.png filter=lfs diff=lfs merge=lfs -text
+asserts/table-5.png filter=lfs diff=lfs merge=lfs -text

llama2-7b-w2a16/config.yaml

@@ -0,0 +1,39 @@
+model: model_zoo/Llama-2-7b-hf
+wbits: 2   
+abits: 16
+epochs: 60
+use_bfloat16: true
+eval_ppl: true
+quant_mode: lora_only  # [fp16,slider,lora_only] slider is for W4A4, lora_only is for WXA16.
+batch_size: 3
+low_memory: true
+nsamples: 128        # the num pf calibration data
+calib_dataset: wikitext2       
+test_datasets: wikitext2,c4
+
+# group layer
+num_layer: 4     # middle sliding window size
+sliding_layer: 2  # middle sliding stride
+fill_window_size: 4   # progressive sliding size
+
+
+# lora
+lora_rank: 4        # lora rank
+use_lora: true     # 默认
+
+
+# instraSliding
+quant_rate: 0.5     
+
+
+# target generation and loss
+inference_batch_size: 12   # the batchsize in get training data
+fp16_act: false           # use fp_act to reduce cpu memory
+
+
+# optimizer
+scale_lr: 0.00    #learn scaling学习率
+lora_lr: 0.0005     # lora 学习率
+lwc_lr: 0.01       # lwc学习率
+warmup_ratio: 0.0  # warnup占总训练iter的比例，范围：[0.0,1.0]
+wd: 0.0

llama2-7b-w4a4/config.yaml

@@ -0,0 +1,46 @@
+model: model_zoo/Llama-2-7b-hf
+wbits: 4 
+abits: 4
+epochs: 20
+use_bfloat16: true
+eval_ppl: true
+quant_mode: slider  # [fp16,slider,lora_only] slider is for W4A4, lora_only is for WXA16.
+batch_size: 3
+low_memory: true
+nsamples: 128        # the num pf calibration data
+calib_dataset: wikitext2       
+test_datasets: wikitext2,c4
+
+load_rotate_model_path: ./rotate_models/llama2-7B-rotate.pth  # Load the rotated FP16 model; if the file does not exist, a rotated model will be generated at the specified path. Applies only to weight-activation quantization.
+
+
+# group layer
+num_layer: 4          # middle sliding window size
+sliding_layer: 2      # middle sliding stride
+
+
+
+
+
+
+# lora
+lora_rank: 4        # lora rank
+use_lora: true     # 默认
+
+
+# instraSliding
+quant_step: 2
+quant_rate: 0.5     
+
+
+# target generation and loss
+inference_batch_size: 12   # the batchsize in get training data
+fp16_act: false           # use fp_act to reduce cpu memory
+
+
+# optimizer
+scale_lr: 0.001    #learn scaling学习率
+lora_lr: 0.0001     # lora 学习率
+lwc_lr: 0.01       # lwc学习率
+warmup_ratio: 0.0  # warnup占总训练iter的比例，范围：[0.0,1.0]
+wd: 0.0

models/LMClass.py

@@ -0,0 +1,108 @@
+import transformers
+import torch
+from .models_utils import BaseLM
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
+import torch.nn.functional as F
+import torch
+from tqdm import tqdm
+import pdb
+
+
+class LMClass(BaseLM):
+    def __init__(self, args,model_name=None):
+
+        super().__init__()
+
+        self.args = args
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if model_name is not None:
+            self.model_name = model_name
+        else:
+            self.model_name = args.model
+        self.batch_size_per_gpu = args.batch_size
+
+        self.model_config = self.model_name
+        config = AutoConfig.from_pretrained(
+            self.model_name, attn_implementation=args.attn_implementation,trust_remote_code=True
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False,legacy=False)
+        torch_dtype = torch.bfloat16 if args.use_bfloat16 is True else torch.float16
+
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_name, config=config, device_map='cpu',torch_dtype=torch_dtype, trust_remote_code=True)
+        
+        self.seqlen = self.model.config.max_position_embeddings
+        self.model.eval()
+        self.vocab_size = self.tokenizer.vocab_size
+        print("vocab size: ", self.vocab_size)
+
+    @property
+    def eot_token(self) -> str:
+        return self.tokenizer.eos_token
+
+    @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):
+        try:
+            return self.gpt2.config.n_ctx
+        except AttributeError:
+            # gptneoconfig doesn't have n_ctx apparently
+            return self.model.config.max_position_embeddings
+
+    @property
+    def max_gen_toks(self):
+        print("max_gen_toks fn")
+        return 256
+
+    @property
+    def batch_size(self):
+        # TODO: fix multi-gpu
+        return self.batch_size_per_gpu  # * gpus
+
+    @property
+    def device(self):
+        # TODO: fix multi-gpu
+        return self._device
+
+    def tok_encode(self, string: str):
+        return self.tokenizer.encode(string, add_special_tokens=False)
+
+    def tok_encode_batch(self, strings):
+        return self.tokenizer(
+            strings,
+            padding=True,
+            add_special_tokens=False,
+            return_tensors="pt",
+        )
+
+    def tok_decode(self, tokens):
+        return self.tokenizer.batch_decode(tokens, skip_special_tokens=True)
+
+    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():
+
+            return self.model(inps)["logits"]
+
+    def model_batched_set(self, inps):
+        dataset_logits = []
+        for batch in inps:
+            multi_logits = F.log_softmax(
+                self._model_call(batch), dim=-1
+            ).cpu()  # [batch, padding_length, vocab]
+            dataset_logits.append(multi_logits)
+        return dataset_logits
+
+    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
+        )

models/int_llama_layer.py

@@ -0,0 +1,627 @@
+import torch
+from torch import nn
+from typing import Optional, Tuple, List
+from quantize.int_linear import QuantLinear
+from quantize.int_matmul import QuantMatMul
+import torch.nn.functional as F
+from quantize.slider_norm import SliderLlamaRMSNorm
+from quantize.slider_norm import RMSN
+from collections import OrderedDict
+import math
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding,apply_rotary_pos_emb,LlamaRMSNorm,repeat_kv
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.activations import ACT2FN
+import copy
+from models.transformation import *
+
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock
+
+
+from quantize.int_linear_lora import LoRAQuantLinear
+
+from models.hadamard_utils import random_hadamard_matrix
+
+from quantize.utils import cleanup_memory
+
+class QuantLlamaMLP(nn.Module):
+    def __init__(
+        self,
+        org_module: nn.Module,
+        hidden_act: str,
+        args=None,
+        use_lora=False,
+        lora_attr=None,
+    ):
+        super().__init__()
+        self.lora_attr = lora_attr
+        self.merged_down = False
+
+        self.args = args
+        if use_lora:
+            self.gate_proj = LoRAQuantLinear(
+                org_module.gate_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+            self.down_proj = LoRAQuantLinear(
+                org_module.down_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+            self.up_proj = LoRAQuantLinear(
+                org_module.up_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+        else:
+            self.gate_proj = QuantLinear(
+                org_module.gate_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+            )
+            self.down_proj = QuantLinear(
+                org_module.down_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+            )
+            self.up_proj = QuantLinear(
+                org_module.up_proj, args.weight_quant_params, args.act_quant_params
+            )
+        self.act_fn = ACT2FN[hidden_act]
+
+
+
+
+    def forward(self, x):
+
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+def get_quant_moe_mlp(org_module: nn.Module,args,config,use_lora,lora_attr):
+    if args.quant_gate is True or args.update_gate is True:
+        if use_lora:
+            weight_quant_params = copy.deepcopy(args.weight_quant_params)
+            act_quant_params = copy.deepcopy(args.act_quant_params)
+            
+            if args.update_gate is True:
+                weight_quant_params["n_bits"] = 16
+                act_quant_params["n_bits"] = 16
+
+            org_module.gate  = LoRAQuantLinear(
+                    org_module.gate,
+                    weight_quant_params,
+                    act_quant_params,
+                    r=args.lora_rank,
+                    lora_attr=lora_attr
+                )
+        else:
+            org_module.gate  = QuantLinear(
+                    org_module.gate,
+                    args.weight_quant_params,
+                    args.act_quant_params,
+                )
+
+
+    for i in range(len(org_module.experts)):
+        ori_mlp = org_module.experts[i]
+        org_module.experts[i] = QuantLlamaMLP(
+            org_module=ori_mlp,
+            hidden_act=config.hidden_act,
+            args=args,
+            use_lora=use_lora,
+            lora_attr=lora_attr
+        )
+    return org_module
+
+
+class QuantLlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, 
+                 org_module: nn.Module,
+                 config: LlamaConfig,
+                 args=None,
+                 use_lora=False,
+                 lora_attr=None,
+                 layer_id = None,
+                ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        # self.head_dim = self.hidden_size // self.num_heads
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.lora_attr = lora_attr
+        self.merge_v = False
+        self.layer_idx = layer_id
+        self.scaling = self.head_dim**-0.5
+        self.add_norm_model_list = ["qwen3" ,"qwen3_moe"]
+
+
+        if config.model_type in self.add_norm_model_list :
+            self.q_norm = copy.deepcopy(org_module.q_norm)
+            self.k_norm = copy.deepcopy(org_module.k_norm)
+
+        self.merged_vo = False
+
+        self.args = args
+        if (self.head_dim * self.num_heads) != self.hidden_size and self.config.model_type not in self.add_norm_model_list:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        # self.rotary_emb = copy.deepcopy(org_module.rotary_emb)
+        self.rotary_emb = LlamaRotaryEmbedding(self.config)
+
+        if use_lora:
+            self.k_proj = LoRAQuantLinear(
+                org_module.k_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+            self.v_proj = LoRAQuantLinear(
+                org_module.v_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+            self.q_proj = LoRAQuantLinear(
+                org_module.q_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+            self.o_proj = LoRAQuantLinear(
+                org_module.o_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                r=args.lora_rank,
+                lora_attr=self.lora_attr
+            )
+        else:
+            self.k_proj = QuantLinear(
+                org_module.k_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+            )
+            self.v_proj = QuantLinear(
+                org_module.v_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+            )
+            self.q_proj = QuantLinear(
+                org_module.q_proj,
+                args.weight_quant_params,
+                args.act_quant_params,
+                disable_input_quant=False,
+            )
+            self.o_proj = QuantLinear(
+                org_module.o_proj, args.weight_quant_params, args.act_quant_params
+            )
+
+        self.qkt_matmul = QuantMatMul(
+            args.q_quant_params, args.k_quant_params, matmul_func=torch.matmul
+        )
+        self.pv_matmul = QuantMatMul(
+            args.p_quant_params, args.v_quant_params, matmul_func=torch.matmul
+        )
+
+        self.use_weight_quant = False
+        self.use_act_quant = False
+
+
+
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        
+        if self.config.model_type in self.add_norm_model_list:
+            # import ipdb;ipdb.set_trace()
+            query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+            key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        else:
+            query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        
+
+        kv_seq_len = key_states.shape[-2]
+        
+        if past_key_value is not None:
+            # kv_seq_len += past_key_value[0].shape[-2]
+            if self.layer_idx is None:
+                    raise ValueError(
+                        f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                        "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                        "with a layer index."
+                    )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+
+
+        # [bsz, nh, t, hd]
+        
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        
+
+
+        query_states = self.qkt_matmul.quant_x1(query_states) # dont quant q
+        key_states = self.qkt_matmul.quant_x2(key_states)
+        attn_weights = self.qkt_matmul(query_states, key_states.transpose(2, 3)) * self.scaling
+
+
+        if attention_mask is not None:
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+            attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = self.pv_matmul.quant_x1(attn_weights) # dont quant p
+        value_states = self.pv_matmul.quant_x2(value_states)
+        attn_output = self.pv_matmul(attn_weights, value_states)
+
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+    
+                
+
+
+class QuantLlamaDecoderLayer(nn.Module):
+    def __init__(self, 
+                 config: LlamaConfig,
+                 ori_layer,
+                 layer_id,
+                 args,
+                 quant_mode="fp16",
+                 use_lora=False,
+                 lora_attr=None,
+                 ):
+        super().__init__()
+        self.use_lora = use_lora
+        self.hidden_size = config.hidden_size
+        self.lora_attr = lora_attr
+
+        self.self_attn = QuantLlamaAttention(
+            org_module=ori_layer.self_attn,
+            config=config,
+            args=args,
+            use_lora=self.use_lora,
+            lora_attr=self.lora_attr,
+            layer_id = layer_id,
+            )
+        if isinstance(ori_layer.mlp, Qwen3MoeSparseMoeBlock):
+            self.mlp = get_quant_moe_mlp(
+                org_module=ori_layer.mlp,
+                args=args,
+                config=config,
+                use_lora=self.use_lora,
+                lora_attr=self.lora_attr
+            )
+        else:
+            self.mlp = QuantLlamaMLP(
+                org_module=ori_layer.mlp,
+                hidden_act=config.hidden_act,
+                args=args,
+                use_lora=self.use_lora,
+                lora_attr=self.lora_attr
+            )
+
+        self.input_layernorm = SliderLlamaRMSNorm(ori_layer.input_layernorm,eps=ori_layer.input_layernorm.variance_epsilon)
+        self.post_attention_layernorm = SliderLlamaRMSNorm(ori_layer.post_attention_layernorm,eps=ori_layer.post_attention_layernorm.variance_epsilon)
+        self.eval_mode = False
+        self.quant_mode = quant_mode 
+        
+        assert self.quant_mode in ["slider","lora_only","fp16"],"only supprot quant_mode in [slider，fp16]"
+        self.finished_quant = False
+        self.revocer_act = False
+
+        self.layer_id = layer_id
+
+        self.massive_act_loc_mask = None
+        self.modify_massive_act = False
+        self.args = args
+        self.Q_merged = False
+        
+        self.update_quant_mode(self.quant_mode,args)
+    
+
+
+    def update_quant_mode(self,new_mode,args=None):
+        self.quant_mode = new_mode
+        support_list =  ["weight_merge","slider","lora_only","fp16"]
+        assert self.quant_mode in support_list,f"only supprot quant_mode in {support_list}"
+
+        if self.quant_mode == "fp16":
+            self.set_quant_state(weight_quant=False,act_quant=False,quant_rate=args.quant_rate)
+            for name, module in self.named_modules():
+                # if isinstance(module, QuantLinear):
+                if hasattr(module,"use_temporary_parameter"):
+                    module.use_temporary_parameter=False
+            self.clear_temp_variable()
+        elif self.quant_mode == "weight_merge":
+            self.set_quant_state(weight_quant=False,act_quant=bool(args.abits<16),quant_rate=args.quant_rate)
+            # import ipdb;ipdb.set_trace()
+            if self.finished_quant is False:
+                # import ipdb;ipdb.set_trace()
+                if args.quant_mode_layer_list[self.layer_id] in ["slider"]:
+                    self.smooth_inplace()
+                with torch.no_grad():
+                    for name, module in self.named_modules():
+                        if hasattr(module,"use_temporary_parameter"):
+                            module.use_temporary_parameter=False
+                        if isinstance(module, LoRAQuantLinear):
+                            module.merged = True
+                            # import ipdb;ipdb.set_trace()
+
+                            if args.lora_rank > 0 :
+                                for i in range(module.lora_iter_num):
+                                    if args.export_model_path and args.export_model_mode =="fp16":
+                                        module.weight = module.weight + module.lora_B[i] @ module.lora_A[i] * module.scaling
+                                    else:
+                                        module.weight = module.weight_quantizer(module.weight + module.lora_B[i] @ module.lora_A[i] * module.scaling)
+                            else:
+                                module.weight = module.weight_quantizer(module.weight)
+                        elif isinstance(module, QuantLinear):
+                            if args.export_model_path and args.export_model_mode =="fp16":
+                                module.weight = module.weight
+                            else:
+                                module.weight = module.weight_quantizer(module.weight)
+                            # module.use_temporary_parameter=False
+                        else:
+                            pass
+                self.finished_quant = True
+            self.clear_temp_variable()
+        else:
+            self.set_quant_state(weight_quant=True,act_quant=bool(args.abits<16),quant_rate=args.quant_rate)
+            self.clear_temp_variable()
+
+    def update_quant_parms(self,weight_quant_params):
+        for name, module in self.named_modules():
+            if isinstance(module, QuantLinear):
+                module.update_quant_parms(weight_quant_params)
+    
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_router_logits = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+
+
+        if not (self.quant_mode in ["fp16","weight_merge","lora_only"]):
+            if self.quant_mode in ["slider"]:
+                self.smooth_and_quant_temporary()
+            else:
+                raise NotImplementedError("only supprot quant_mode in [slider,lora_only,weight_merge,fp16]")
+        
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        
+
+        hidden_states = self.mlp(hidden_states)
+
+        if isinstance(hidden_states, tuple):
+            hidden_states, router_logits = hidden_states
+        else:
+            router_logits = None
+
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        if self.eval_mode:
+            self.clear_temp_variable()
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if output_router_logits:
+            outputs += (router_logits,)
+        return outputs        
+
+    def set_quant_state(self, weight_quant: bool = False, act_quant: bool = False, quant_rate:float = 1.0):
+        # setting weight quantization here does not affect actual forward pass
+        self.use_weight_quant = weight_quant
+        self.use_act_quant = act_quant
+        self.quant_rate = quant_rate
+        names = []
+        for name, m in self.named_modules():
+            if isinstance(m, (QuantLinear, QuantMatMul)):
+                names.append(name)
+                m.set_quant_state(weight_quant, act_quant, quant_rate)
+      
+    def smooth_and_quant_temporary(self):
+        with torch.no_grad():
+            for name, module in self.named_parameters():
+                if "smooth_scale" in name:
+                    module.data = truncate_number(module)
+
+
+
+        smooth_ln_fcs_temporary(self.input_layernorm,[self.self_attn.q_proj, self.self_attn.k_proj, self.self_attn.v_proj],
+                                self.qkv_smooth_scale,self.qkv_smooth_shift)
+        smooth_ln_fcs_temporary(self.post_attention_layernorm,[self.mlp.up_proj,self.mlp.gate_proj],
+                                self.fc1_smooth_scale,self.fc1_smooth_shift)
+        smooth_fc_fc_temporary(self.self_attn.v_proj,self.self_attn.o_proj,
+                        self.out_smooth_scale, self.out_smooth_shift,num_key_value_groups=self.self_attn.num_key_value_groups,head_dim=self.self_attn.head_dim,args=self.args)
+        smooth_q_k_temporary(self.self_attn.q_proj, self.self_attn.k_proj,
+                        self.qkt_smooth_scale,num_key_value_groups=self.self_attn.num_key_value_groups,head_dim=self.self_attn.head_dim,args=self.args)
+        if self.args.use_down_scale is True:
+            smooth_fc_fc_temporary(self.mlp.up_proj,self.mlp.down_proj,
+                                   self.fc2_smooth_scale,self.fc2_smooth_shift)
+        else:
+            self.mlp.down_proj.temp_weight = self.mlp.down_proj.weight
+    
+        
+        # quant
+        for name, module in self.named_modules():
+            if isinstance(module, QuantLinear):
+                if hasattr(module, "temp_weight"):
+                    module.temp_weight = module.temp_weight
+                else:
+                    module.temp_weight = module.weight
+                if not hasattr(module, "temp_bias"):
+                    module.temp_bias = module.bias
+                module.use_temporary_parameter=True
+
+    def clear_temp_variable(self):
+        for name, module in self.named_modules():
+            if hasattr(module,"temp_weight"):
+                del module.temp_weight
+            if hasattr(module,"temp_bias"):
+                del module.temp_bias
+
+
+    @torch.no_grad()
+    def smooth_inplace(self):
+        if self.let:
+            for name, module in self.named_parameters():
+                if "smooth_scale" in name:
+                    module.data = truncate_number(module)
+            smooth_ln_fcs_inplace(self.input_layernorm,[self.self_attn.q_proj, self.self_attn.k_proj, self.self_attn.v_proj],
+                                    self.qkv_smooth_scale,self.qkv_smooth_shift)
+            smooth_ln_fcs_inplace(self.post_attention_layernorm,[self.mlp.up_proj,self.mlp.gate_proj],
+                                    self.fc1_smooth_scale,self.fc1_smooth_shift)
+            smooth_fc_fc_inplace(self.self_attn.v_proj,self.self_attn.o_proj,
+                                self.out_smooth_scale, self.out_smooth_shift,self.self_attn.num_key_value_groups,head_dim=self.self_attn.head_dim,args=self.args)
+            smooth_q_k_inplace(self.self_attn.q_proj, self.self_attn.k_proj,
+                                self.qkt_smooth_scale,self.self_attn.num_key_value_groups,head_dim=self.self_attn.head_dim,args=self.args)
+            if self.args.use_down_scale is True:
+                smooth_fc_fc_inplace(self.mlp.up_proj,self.mlp.down_proj,
+                                    self.fc2_smooth_scale,self.fc2_smooth_shift)
+
+    def get_slider_parameters(self, use_list=["scale"]):
+        params = []
+        for n, m in self.named_parameters():
+            if any(n.find(t) > -1 for t in use_list):
+                params.append(m)
+        return iter(params)  
+
+    def get_lwc_parameters(self):
+        params = []
+        for n, m in self.named_parameters():
+            if n.find('bound_factor') > -1:
+                params.append(m)
+        return iter(params)  
+
+    def slider_parameters(self, use_list=["scale"]):
+        params = []
+        for n, m in self.named_parameters():
+            if n.find('bound_factor') > -1 or any(n.find(t) > -1 for t in use_list):
+                params.append(m)
+        return iter(params)  
+
+    def slider_state_dict(self, destination=None, prefix='', keep_vars=False):
+        if destination is None:
+            destination = OrderedDict()
+        for name, param in self.named_parameters():
+            if name.find('smooth') > -1 or name.find('bound_factor') > -1:
+                destination[prefix + name] = param if keep_vars else param.detach()
+        return destination
+    
+    def register_scales_and_zeros(self):
+        for name, module in self.named_modules():
+            if isinstance(module, QuantLinear):
+                module.weight_quantizer.register_scales_and_zeros()

models/models_utils.py

@@ -0,0 +1,534 @@
+import abc
+import torch
+import json
+import hashlib
+import collections
+from tqdm import tqdm
+from typing import Iterable
+from abc import abstractmethod
+from torch import nn
+import transformers
+import torch.nn.functional as F
+
+class TruncateFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, threshold):
+        truncated_tensor = input.clone()
+        truncated_tensor[truncated_tensor.abs() < threshold] = truncated_tensor[truncated_tensor.abs() < threshold].sign() * threshold
+        return truncated_tensor
+        
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input, None
+
+def truncate_number(number, threshold=1e-3):
+    # clamping scale to avoid overflow in the AMP training
+    return TruncateFunction.apply(number, threshold)
+
+def find_layers(module, layers=[nn.Conv2d, nn.Linear, transformers.Conv1D], name=""):
+    if type(module) in layers:
+        return {name: module}
+    res = {}
+    for name1, child in module.named_children():
+        res.update(
+            find_layers(
+                child, layers=layers, name=name + "." + name1 if name != "" else name1
+            )
+        )
+    return res
+
+
+class CacheHook:
+    def __init__(self, cachinglm):
+        if cachinglm is None:
+            self.dbdict = None
+            return
+
+        self.dbdict = cachinglm.dbdict
+
+    def add_partial(self, attr, req, res):
+        if self.dbdict is None:
+            return
+        hsh = hash_args(attr, req)
+        self.dbdict[hsh] = res
+
+
+class LM(abc.ABC):
+    def __init__(self):
+        self.cache_hook = CacheHook(None)
+
+    @abstractmethod
+    def loglikelihood(self, requests):
+        """Compute log-likelihood of generating a continuation from a context.
+        Downstream tasks should attempt to use loglikelihood instead of other
+        LM calls whenever possible.
+
+        :param requests: list
+            A list of pairs (context, continuation)
+            context: str
+                Context string. Implementations of LM must be able to handle an
+                empty context string.
+            continuation: str
+                The continuation over which log likelihood will be calculated. If
+                there is a word boundary, the space should be in the continuation.
+                For example, context="hello" continuation=" world" is correct.
+        :return: list
+            A list of pairs (logprob, isgreedy)
+            logprob: float
+                The log probability of `continuation`
+            isgreedy:
+                Whether `continuation` would be generated by greedy sampling from `context`
+        """
+        pass
+
+    @abstractmethod
+    def loglikelihood_rolling(self, requests):
+        """Compute full log-likelihood of a string, with no truncation, for perplexity computation
+        - We will use the full max context length of the model.
+        - For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
+        the max context length.
+        - IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
+          which may simply concatenate multiple documents together.
+        - IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
+          multiple chunks, the last input will still a full-sized context.
+          Example:
+            Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
+            Prefix: EOT
+            Max context length: 4
+            Resulting input/prediction pairs:
+
+                INPUT:  EOT   0   1   2
+                PRED:     0   1   2   3
+
+                INPUT:    3   4   5   6
+                PRED:     4   5   6   7
+
+                INPUT:    5   6   7   8
+                PRED:             8   9
+
+          Observe that:
+            1. Each token is predicted exactly once
+            2. For the last pair, we provide the full context, but only score the last two tokens
+
+        :param requests: list
+            A list of strings
+            string: str
+                String for which we are computing per-toke  loglikelihood
+        :return: list
+            A list of pairs (logprob, isgreedy)
+            logprob: float
+                The log probability of `continuation`
+            isgreedy:
+                Whether `continuation` would be generated by greedy sampling from `context`
+        """
+        pass
+
+    # TODO: Add an optional max length
+    @abstractmethod
+    def greedy_until(self, requests):
+        """Generate greedily until a stopping sequence
+
+        :param requests: list
+            A list of pairs (context, until)
+            context: str
+                Context string
+            until: [str]
+                The string sequences to generate until. These string sequences
+                may each span across multiple tokens, or may be part of one token.
+        :return: list
+            A list of strings continuation
+            continuation: str
+                The generated continuation.
+        """
+        pass
+
+    @classmethod
+    def create_from_arg_string(cls, additional_config=None):
+        additional_config = {} if additional_config is None else additional_config
+        args = {k: v for k, v in additional_config.items() if v is not None}
+        return cls(**args)
+
+    def set_cache_hook(self, cache_hook):
+        self.cache_hook = cache_hook
+
+
+class BaseLM(LM):
+    @property
+    @abstractmethod
+    def eot_token_id(self):
+        pass
+
+    @property
+    @abstractmethod
+    def max_length(self):
+        pass
+
+    @property
+    @abstractmethod
+    def max_gen_toks(self):
+        pass
+
+    @property
+    @abstractmethod
+    def batch_size(self):
+        pass
+
+    @property
+    @abstractmethod
+    def device(self):
+        pass
+
+    @abstractmethod
+    def tok_encode(self, string: str):
+        pass
+
+    @abstractmethod
+    def tok_decode(self, tokens: Iterable[int]):
+        pass
+
+    @abstractmethod
+    def _model_generate(self, context, max_length, eos_token_id):
+        pass
+
+    @abstractmethod
+    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
+        """
+        pass
+
+    # subclass must implement properties vocab_size, eot_token_id, max_gen_toks, batch_size, device, max_length.
+    # TODO: enforce this somehow
+
+    def loglikelihood(self, requests):
+        new_reqs = []
+        for context, continuation in requests:
+            if context == "":
+                # end of text as context
+                context_enc = [self.eot_token_id]
+            else:
+                context_enc = self.tok_encode(context)
+
+            continuation_enc = self.tok_encode(continuation)
+            new_reqs.append(((context, continuation), context_enc, continuation_enc))
+
+        return self._loglikelihood_tokens(new_reqs)
+
+    def loglikelihood_rolling(self, requests):
+        # TODO: Implement caching once we've confirmed the perplexity implementation
+        # TODO: automatic batch size detection for vectorization
+
+        loglikelihoods = []
+        for (string,) in tqdm(requests):
+            rolling_token_windows = list(
+                map(
+                    make_disjoint_window,
+                    get_rolling_token_windows(
+                        token_list=self.tok_encode(string),
+                        prefix_token=self.eot_token_id,
+                        max_seq_len=self.max_length,
+                        context_len=1,
+                    ),
+                )
+            )
+
+            rolling_token_windows = [(None,) + x for x in rolling_token_windows]
+
+            # TODO: extract out this call so it only gets called once and also somehow figure out partial caching for
+            # that
+            string_nll = self._loglikelihood_tokens(
+                rolling_token_windows, disable_tqdm=True
+            )
+
+            # discard is_greedy
+            string_nll = [x[0] for x in string_nll]
+
+            string_nll = sum(string_nll)
+            loglikelihoods.append(string_nll)
+
+        return loglikelihoods
+
+    def _loglikelihood_tokens(self, requests, disable_tqdm=False):
+        # TODO: implement some kind of efficient-request-middleware that lumps together requests with the same context
+        res = []
+        dataset_inps = []
+
+        def _collate(x):
+            # the negative sign on len(toks) sorts descending - this has a few advantages:
+            # - time estimates will always be over not underestimates, which is more useful for planning
+            # - to know the size of a batch when going through the list, you know the first one is always the batch
+            #   padded context length. this is useful to simplify the batching logic and more importantly to make
+            #   automatic adaptive batches much much easier to implement
+            # - any OOMs will happen right away rather than near the end
+
+            toks = x[1] + x[2]
+            return -len(toks), tuple(toks)
+
+        # TODO: automatic (variable) batch size detection for vectorization
+        re_ord = Reorderer(requests, _collate)
+        for chunk in chunks(
+            tqdm(re_ord.get_reordered(), disable=disable_tqdm), self.batch_size
+        ):
+            inps = []
+            cont_toks_list = []
+            inplens = []
+
+            padding_length = None
+
+            # because vectorizing is annoying, we first convert each (context, continuation) pair to padded
+            # tensors, then we pack them together into a batch, call the model, and then pick it all apart
+            # again because vectorizing is annoying
+
+            for _, context_enc, continuation_enc in chunk:
+                # sanity check
+                assert len(context_enc) > 0
+                assert len(continuation_enc) > 0
+                assert len(continuation_enc) <= self.max_length
+
+                # how this all works:
+                #          CTX      CONT
+                # inp    0 1 2 3|4 5 6 7 8 9   <- last token is deleted by inp[:, :-1]
+                # gpt2    \               \
+                # logits   1 2 3|4 5 6 7 8 9   <- the ctx half gets tossed out by the
+                # cont_toks      4 5 6 7 8 9      [:, -len(continuation_enc):, :self.vocab_size] slice
+
+                # when too long to fit in context, truncate from the left
+                inp = torch.tensor(
+                    (context_enc + continuation_enc)[-(self.max_length + 1) :][:-1],
+                    dtype=torch.long,
+                ).to(self.device)
+                (inplen,) = inp.shape
+
+                cont = continuation_enc
+
+                # since in _collate we make sure length is descending, the longest is always the first one.
+                padding_length = (
+                    padding_length if padding_length is not None else inplen
+                )
+
+                # pad length from seq to padding_length
+                inp = torch.cat(
+                    [
+                        inp,  # [seq]
+                        torch.zeros(padding_length - inplen, dtype=torch.long).to(
+                            inp.device
+                        ),  # [padding_length - seq]
+                    ],
+                    dim=0,
+                )
+
+                inps.append(inp.unsqueeze(0))  # [1, padding_length]
+                cont_toks_list.append(cont)
+                inplens.append(inplen)
+            # import pdb; pdb.set_trace()
+            batched_inps = torch.cat(inps, dim=0).to(
+                self.device
+            )  # [batch, padding_length
+
+            # self.model = self.model.to(self.device)
+            multi_logits = F.log_softmax(
+                self._model_call(batched_inps), dim=-1
+            ).cpu()  # [batch, padding_length, vocab]
+
+
+            for (cache_key, _, _), logits, inp, inplen, cont_toks in zip(
+                chunk, multi_logits, inps, inplens, cont_toks_list
+            ):
+
+                # Slice to original seq length
+                contlen = len(cont_toks)
+                logits = logits[inplen - contlen : inplen].unsqueeze(
+                    0
+                )  # [1, seq, vocab]
+
+                # Check if per-token argmax is exactly equal to continuation
+                greedy_tokens = logits.argmax(dim=-1)
+                cont_toks = torch.tensor(cont_toks, dtype=torch.long).unsqueeze(
+                    0
+                )  # [1, seq]
+                # import pdb; pdb.set_trace()
+                max_equal = (greedy_tokens == cont_toks).all()
+
+                # Obtain log-probs at the corresponding continuation token indices
+                # last_token_slice = logits[:, -1, :].squeeze(0).tolist()
+                logits = torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(
+                    -1
+                )  # [1, seq]
+
+                # Answer: (log prob, is-exact-match)
+                answer = (float(logits.sum()), bool(max_equal))
+
+                # partial caching
+                if cache_key is not None:
+                    self.cache_hook.add_partial("loglikelihood", cache_key, answer)
+                res.append(answer)
+
+        return re_ord.get_original(res)
+
+    def greedy_until(self, requests):
+        print("greedy utils in base...")
+        # TODO: implement fully general `until` that handles until that are
+        #       multiple tokens or that span multiple tokens correctly
+
+        # TODO: extract to TokenizedLM?
+        res = []
+
+        def _collate(x):
+            toks = self.tok_encode(x[0])
+            return len(toks), x[0]
+
+        re_ord = Reorderer(requests, _collate)
+
+        for context, until in tqdm(re_ord.get_reordered()):
+            if isinstance(until, str):
+                until = [until]
+
+            (primary_until,) = self.tok_encode(until[0])
+
+            context_enc = torch.tensor(
+                [self.tok_encode(context)[self.max_gen_toks - self.max_length :]]
+            ).to(self.device)
+
+            cont = self._model_generate(
+                context_enc, context_enc.shape[1] + self.max_gen_toks, primary_until
+            )
+
+            s = self.tok_decode(cont[0].tolist()[context_enc.shape[1] :])
+
+            for term in until:
+                s = s.split(term)[0]
+
+            # partial caching
+            self.cache_hook.add_partial("greedy_until", (context, until), s)
+
+            res.append(s)
+
+        return re_ord.get_original(res)
+
+
+def make_disjoint_window(pair):
+    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
+    a, b = pair
+    return a[: len(a) - (len(b) - 1)], b
+
+
+def hash_args(attr, args):
+    dat = json.dumps([attr] + list(args))
+    return hashlib.sha256(dat.encode("utf-8")).hexdigest()
+
+
+def simple_parse_args_string(args_string):
+    """
+    Parses something like
+        args1=val1,arg2=val2
+    Into a dictionary
+    """
+    args_string = args_string.strip()
+    if not args_string:
+        return {}
+    arg_list = args_string.split(",")
+    args_dict = {}
+    for arg in arg_list:
+        k, v = arg.split("=")
+        args_dict[k] = v
+    return args_dict
+
+
+def get_rolling_token_windows(token_list, prefix_token, max_seq_len, context_len):
+    """
+    - context_len allows for a rolling window context, allowing each prediction window to potentially
+      condition on some context
+
+    :param token_list: list
+        List of tokens to be PREDICTED
+    :param max_seq_len: int
+        max_seq_len of model (or max_seq_len we want to use)
+    :param context_len: int
+        Amount of desired token context for prediction. Needs to be at least 1.
+    :param prefix_token: token
+        Dummy token like <eos> so the first token has something to condition on
+    :return: generator
+        Generator of tuples
+            (input_tokens, pred_tokens)
+        Note: Score only the last len(pred_tokens) logits of the LM
+    """
+    assert 1 <= context_len <= max_seq_len
+    if not token_list:
+        return
+    # +1 offset, going from input->preds
+    pred_len = max_seq_len - context_len + 1
+    predicted = 0
+
+    # Special handling for first window: predict all tokens
+    first_seq_len = min(max_seq_len, len(token_list))
+    yield ([prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len])
+    predicted += first_seq_len
+
+    while predicted < len(token_list):
+        window_pred_len = min(len(token_list) - predicted, pred_len)
+        window_end = predicted + window_pred_len
+
+        yield (
+            token_list[window_end - max_seq_len - 1 : window_end - 1],
+            token_list[window_end - window_pred_len : window_end],
+        )
+        predicted += window_pred_len
+
+
+class Reorderer:
+    def __init__(self, arr, fn):
+        self.size = len(arr)
+        arr = list(enumerate(arr))
+        arr = group(arr, lambda x: fn(x[1]))
+        arr = [([y[0] for y in x], x[0][1]) for x in arr]
+        arr.sort(key=lambda x: fn(x[1]))
+
+        self.arr = arr
+
+    def get_reordered(self):
+        return [x[1] for x in self.arr]
+
+    def get_original(self, newarr):
+        res = [None] * self.size
+        cov = [False] * self.size
+
+        for (inds, _), v in zip(self.arr, newarr):
+            for ind in inds:
+                res[ind] = v
+                cov[ind] = True
+
+        assert all(cov)
+
+        return res
+
+
+def join_iters(iters):
+    for iter in iters:
+        yield from iter
+
+
+def chunks(iter, n):
+    arr = []
+    for x in iter:
+        arr.append(x)
+        if len(arr) == n:
+            yield arr
+            arr = []
+
+    if arr:
+        yield arr
+
+
+def group(arr, fn):
+    res = collections.defaultdict(list)
+
+    for ob in arr:
+        res[fn(ob)].append(ob)
+
+    return list(res.values())

models/rotation_utils.py

@@ -0,0 +1,336 @@
+import torch
+import typing
+from quantize import utils
+import transformers
+import tqdm, math
+from models.hadamard_utils import random_hadamard_matrix
+from dataclasses import dataclass
+
+OPT_MODEL = transformers.models.opt.modeling_opt.OPTForCausalLM
+OPT_LAYER = transformers.models.opt.modeling_opt.OPTDecoderLayer
+LLAMA_MODEL = transformers.models.llama.modeling_llama.LlamaForCausalLM
+LLAMA_LAYER = transformers.models.llama.modeling_llama.LlamaDecoderLayer
+QWEN_MODEL =  transformers.models.qwen2.modeling_qwen2.Qwen2ForCausalLM
+QWEN_LATER = transformers.models.qwen2.modeling_qwen2.Qwen2DecoderLayer
+
+
+DEV = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+
+class RMSN(torch.nn.Module):
+    """
+    This class implements the Root Mean Square Normalization (RMSN) layer.
+    We use the implementation from LLAMARMSNorm here:
+    https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L75
+    """
+
+    def __init__(self, mean_dim: int,ori_norm, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.mean_dim = mean_dim
+        if ori_norm:
+            self.weight = torch.nn.Parameter(torch.ones_like(ori_norm.weight))
+            if hasattr(ori_norm,"bias") and ori_norm.bias is not None:
+                self.bias = torch.nn.Parameter(torch.zeros_like(ori_norm.bias))
+        else:
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.bias = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return (self.weight * x).to(input_dtype)
+
+
+
+def replace_modules(
+    root: torch.nn.Module,
+    type_to_replace,
+    new_module_factory,
+    replace_layers: bool,
+) -> None:
+    """Replace modules of given type using the supplied module factory.
+
+    Perform a depth-first search of a module hierarchy starting at root
+    and replace all instances of type_to_replace with modules created by
+    new_module_factory. Children of replaced modules are not processed.
+
+    Args:
+        root: the root of the module hierarchy where modules should be replaced
+        type_to_replace: a type instances of which will be replaced
+        new_module_factory: a function that given a module that should be replaced
+            produces a module to replace it with.
+    """
+    for name, module in root.named_children():
+        new_module = None
+        if isinstance(module, type_to_replace):
+            if replace_layers:  # layernorm_fusion.replace_layers case where transformer layers are replaced
+                new_module = new_module_factory(module, int(name))
+            else:  # layernorm_fusion.fuse_modules case where layernorms are fused
+                new_module = new_module_factory(module)
+        elif len(list(module.children())) > 0:
+            replace_modules(module, type_to_replace, new_module_factory, replace_layers)
+
+        if new_module is not None:
+            setattr(root, name, new_module)
+
+def get_model_type(model):
+    if isinstance(model, OPT_MODEL):
+        model_type = OPT_MODEL
+    elif isinstance(model, LLAMA_MODEL):
+        model_type = LLAMA_MODEL
+    elif isinstance(model, QWEN_MODEL):
+        model_type = QWEN_MODEL
+    else:
+        raise ValueError(f'Unknown model type {model}')
+    return model_type
+
+
+def model_type_extractor(model):
+    if isinstance(model, LLAMA_MODEL):
+        return LLAMA_MODEL
+    elif isinstance(model, QWEN_MODEL):
+        return QWEN_MODEL
+    elif isinstance(model, OPT_MODEL):
+        return OPT_MODEL
+    else:
+        raise ValueError(f'Unknown model type {model}')
+
+def get_embeddings(model, model_type) -> list[torch.nn.Module]:
+    if model_type == LLAMA_MODEL or  model_type == QWEN_MODEL :
+        return [model.model.embed_tokens]
+    elif model_type == OPT_MODEL:
+        return [model.model.decoder.embed_tokens, model.model.decoder.embed_positions]
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+
+
+def get_transformer_layers(model, model_type):
+    if model_type == LLAMA_MODEL  or  model_type == QWEN_MODEL:
+        return [layer for layer in model.model.layers]
+    elif model_type == OPT_MODEL:
+        return [layer for layer in model.model.decoder.layers]
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+
+
+def get_lm_head(model, model_type):
+    if model_type == LLAMA_MODEL or  model_type == QWEN_MODEL:
+        return model.lm_head
+    elif model_type == OPT_MODEL:
+        return model.lm_head
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+
+
+
+
+def get_pre_head_layernorm(model, model_type):
+    if model_type == LLAMA_MODEL  or  model_type == QWEN_MODEL :
+        pre_head_layernorm = model.model.norm
+        assert isinstance(pre_head_layernorm,
+                          (transformers.models.llama.modeling_llama.LlamaRMSNorm,transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm))
+    elif model_type == OPT_MODEL:
+        pre_head_layernorm = model.model.decoder.final_layer_norm
+        assert pre_head_layernorm is not None
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+    return pre_head_layernorm
+
+def fuse_ln_linear(layernorm: torch.nn.Module, linear_layers: typing.Iterable[torch.nn.Linear]) -> None:
+    """
+    fuse the linear operations in Layernorm into the adjacent linear blocks.
+    """
+    for linear in linear_layers:
+        linear_dtype = linear.weight.dtype
+
+        # Calculating new weight and bias
+        W_ = linear.weight.data.double()
+        linear.weight.data = (W_ * layernorm.weight.double()).to(linear_dtype)
+
+        if hasattr(layernorm, 'bias'):
+            if linear.bias is None:
+                linear.bias = torch.nn.Parameter(torch.zeros(linear.out_features, dtype=torch.float64))
+            linear.bias.data = linear.bias.data.double() + torch.matmul(W_, layernorm.bias.double())
+            linear.bias.data = linear.bias.data.to(linear_dtype)
+            
+def bake_mean_into_linear(linear: torch.nn.Linear) -> None:
+    """
+    This function takes a linear layer and subtracts the means from the
+    weights and biases. This will result in the linear layer performing
+    the mean substitution which is usually done inside layernorm.
+    """
+    linear_dtype = linear.weight.dtype
+    W_ = linear.weight.data.double()
+    linear.weight.data = W_ - W_.mean(dim=-2, keepdim=True)
+    linear.weight.data = linear.weight.data.to(linear_dtype)
+    if linear.bias is not None:
+        b_ = linear.bias.data.double()
+        linear.bias.data = b_ - b_.mean()
+        linear.bias.data = linear.bias.data.to(linear_dtype)
+
+         
+            
+def fuse_layer_norms(model):
+    
+    model_type = get_model_type(model)
+    
+    kwargs = {'model': model, 'model_type': model_type}
+    
+    # Embedding fusion
+    for W in get_embeddings(**kwargs):
+        W_ = W.weight.data.double()
+        W.weight.data = (W_ - W_.mean(dim=-1, keepdim=True)).to(W.weight.data.dtype)
+        
+    layers = get_transformer_layers(**kwargs)
+    
+    # Fuse the linear operations in Layernorm into the adjacent linear blocks.
+    for layer in layers:
+        
+        # fuse the input layernorms into the linear layers
+        if model_type == LLAMA_MODEL  or  model_type == QWEN_MODEL :
+            fuse_ln_linear(layer.post_attention_layernorm, [layer.mlp.up_proj, layer.mlp.gate_proj])    
+            fuse_ln_linear(layer.input_layernorm, [layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj])
+        elif model_type == OPT_MODEL:
+            fuse_ln_linear(layer.self_attn_layer_norm, [layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj])
+            fuse_ln_linear(layer.final_layer_norm, [layer.fc1])
+        else:
+            raise ValueError(f'Unknown model type {model_type}')
+            
+            
+    
+        if model_type == OPT_MODEL:
+            bake_mean_into_linear(layer.self_attn.out_proj)
+            bake_mean_into_linear(layer.fc2)
+                    
+    
+    fuse_ln_linear(get_pre_head_layernorm(**kwargs), [get_lm_head(**kwargs)])
+    
+    model_type2_NORM = {
+        LLAMA_MODEL: transformers.models.llama.modeling_llama.LlamaRMSNorm,
+        QWEN_MODEL: transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm,
+        OPT_MODEL: torch.nn.LayerNorm
+    }
+    
+    replace_modules(
+        model,
+        model_type2_NORM[model_type],
+        lambda ori_norm: RMSN(model.config.hidden_size,ori_norm),
+        replace_layers=False,
+    )
+    
+
+
+def get_orthogonal_matrix(size, mode, device=DEV):
+    return random_hadamard_matrix(size, device)
+    
+
+def rotate_embeddings(model, Q: torch.Tensor) -> None:
+    # Rotate the embeddings.
+    model_type = model_type_extractor(model)
+    for W in get_embeddings(model, model_type):
+        dtype = W.weight.data.dtype
+        W_ = W.weight.data.to(device=DEV, dtype=torch.float64)
+        W.weight.data = torch.matmul(W_, Q).to(device="cpu", dtype=dtype)
+
+    
+def rotate_attention_inputs(layer, Q, model_type) -> None:
+    # Rotate the WQ, WK and WV matrices of the self-attention layer.
+    for W in [layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj]:
+        dtype = W.weight.dtype
+        W_ = W.weight.to(device=DEV, dtype=torch.float64)
+        W.weight.data = torch.matmul(W_, Q).to(device="cpu", dtype=dtype)
+
+def rotate_attention_output(layer, Q, model_type) -> None:
+    # Rotate output matrix of the self-attention layer.
+    if model_type == LLAMA_MODEL or  model_type == QWEN_MODEL:
+        W = layer.self_attn.o_proj
+    elif model_type == OPT_MODEL:
+        W = layer.self_attn.out_proj
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+
+    dtype = W.weight.data.dtype
+    W_ = W.weight.data.to(device=DEV, dtype=torch.float64)
+    W.weight.data = torch.matmul(Q.T, W_).to(device="cpu", dtype=dtype)
+    if W.bias is not None:
+        b = W.bias.data.to(device=DEV, dtype=torch.float64)
+        W.bias.data = torch.matmul(Q.T, b).to(device="cpu", dtype=dtype)
+
+def rotate_mlp_input(layer, Q, model_type):
+    # Rotate the MLP input weights.
+    if model_type == LLAMA_MODEL or  model_type == QWEN_MODEL:
+        mlp_inputs = [layer.mlp.up_proj, layer.mlp.gate_proj]
+    elif model_type == OPT_MODEL:
+        mlp_inputs = [layer.fc1]
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+    for W in mlp_inputs:
+        dtype = W.weight.dtype
+        W_ = W.weight.data.to(device=DEV, dtype=torch.float64)
+        W.weight.data = torch.matmul(W_, Q).to(device="cpu", dtype=dtype)
+    
+def rotate_mlp_output(layer, Q, model_type,add_online_rotate=True):
+    # Rotate the MLP output weights and bias.
+    if model_type == LLAMA_MODEL or  model_type == QWEN_MODEL:
+        W = layer.mlp.down_proj
+    elif model_type == OPT_MODEL:
+        W = layer.fc2
+    else:
+        raise ValueError(f'Unknown model type {model_type}')
+    dtype = W.weight.data.dtype
+    W_ = W.weight.data.to(device=DEV, dtype=torch.float64)
+    W.weight.data = torch.matmul(Q.T, W_).to(device="cpu", dtype=dtype)
+    if W.bias is not None:
+        b = W.bias.data.to(device=DEV, dtype=torch.float64)
+        W.bias.data = torch.matmul(Q.T, b).to(device="cpu", dtype=dtype)
+
+
+
+def rotate_head(model, Q: torch.Tensor) -> None:
+    # Rotate the head.
+    W = get_lm_head(model, model_type=model_type_extractor(model))
+    dtype = W.weight.data.dtype
+    W_ = W.weight.data.to(device=DEV, dtype=torch.float64)
+    W.weight.data = torch.matmul(W_, Q).to(device="cpu", dtype=dtype)
+
+@torch.inference_mode()
+def rotate_model(model, args,add_online_rotate=True):
+    Q = get_orthogonal_matrix(model.config.hidden_size,
+                                                args.rotate_mode)
+    config = model.config
+    num_heads = config.num_attention_heads
+    model_dim = config.hidden_size
+    head_dim = model_dim // num_heads
+
+
+    model_type = model_type_extractor(model)
+    rotate_embeddings(model, Q)
+    rotate_head(model, Q)
+    utils.cleanup_memory()
+    layers = get_transformer_layers(model, model_type=model_type)
+    for idx, layer in enumerate(tqdm.tqdm(layers, unit="layer", desc="Rotating")):
+        rotate_attention_inputs(layers[idx], Q, model_type)
+        rotate_attention_output(layers[idx], Q, model_type)
+        rotate_mlp_input(layers[idx], Q, model_type)
+        rotate_mlp_output(layers[idx], Q, model_type,add_online_rotate)
+
+
+
+def get_rotate_model(model,save_path):
+    @dataclass
+    class parm():
+        rotate_mode:str
+        save_qmodel_path:str
+        fp32_had:bool
+    args = parm("hadamard",save_path,False)
+    model.to("cpu")
+    model.eval()
+    fuse_layer_norms(model)
+    rotate_model(model, args,add_online_rotate=False)
+    save_dict = {}
+    save_dict["model"] = model.state_dict()
+    torch.save(save_dict, args.save_qmodel_path)
+    return save_dict

models/transformation.py

@@ -0,0 +1,214 @@
+
+import torch
+
+
+
+class TruncateFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, threshold):
+        ctx.save_for_backward(input)
+        ctx.threshold = threshold
+        truncated_tensor = input.clone()
+        truncated_tensor = torch.where(truncated_tensor.abs() < threshold, truncated_tensor.sign() * threshold, truncated_tensor)
+        return truncated_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, = ctx.saved_tensors
+        threshold = ctx.threshold
+        grad_input = grad_output.clone()
+        grad_input[input.abs() < threshold] = 0
+        return grad_input, None
+
+def truncate_number(number, threshold=1e-2):
+    # avoid overflow with AMP training
+    return TruncateFunction.apply(number, threshold)
+
+
+
+def smooth_ln_fcs_temporary(ln, fcs, scales,shifts,weight_in_cpu=False):
+    ln.use_temporary_parameter = True
+    if weight_in_cpu is True:
+        ln.weight = ln.weight.to(scales.device)
+        for fc in fcs:
+            fc.weight = fc.weight.to(scales.device)
+
+
+    if not isinstance(fcs, list):
+        fcs = [fcs]
+    if hasattr(ln, 'bias') and ln.bias is not None:
+        ln.temp_bias = (ln.bias - shifts) / scales
+    else:
+        ln.temp_bias = (-1*shifts) / scales
+
+    ln.temp_weight = ln.weight / scales
+
+    for fc in fcs:
+        fc.use_temporary_parameter = True
+        if hasattr(fc, 'bias') and fc.bias is not None:
+            fc.temp_bias = fc.bias + torch.matmul(fc.weight,shifts)
+        else:
+            fc.temp_bias = torch.matmul(fc.weight,shifts)
+        fc.temp_weight = fc.weight * scales.view(1,-1)
+    
+    if weight_in_cpu is True:
+        ln.weight = ln.weight.cpu()
+        for fc in fcs:
+            fc.weight = fc.weight.cpu()
+    
+
+def smooth_ln_fcs_inplace(ln, fcs, scales,shifts):
+    ln.use_temporary_parameter = False
+    if not isinstance(fcs, list):
+        fcs = [fcs]
+    if hasattr(ln, 'bias') and ln.bias is not None:
+        ln.bias.sub_(shifts)
+        ln.bias.div_(scales)
+    else:
+        del ln.bias
+        ln.register_buffer('bias',(-1*shifts)/scales)
+
+    ln.weight.div_(scales)
+    for fc in fcs:
+        fc.use_temporary_parameter = False
+        if hasattr(fc, 'bias') and fc.bias is not None:
+            fc.bias.add_(fc.weight@shifts)
+        else:
+            del fc.bias
+            # import ipdb;ipdb.set_trace()
+            fc.register_buffer('bias',fc.weight@shifts)
+        fc.weight.mul_(scales.view(1,-1))
+
+
+
+def smooth_fc_fc_temporary(fc1, fc2, scales,shifts=None,num_key_value_groups=1,head_dim=128,weight_in_cpu=False,args=None):
+    # only support for v_proj and out_proh now.
+    fc1.use_temporary_parameter = True
+    fc2.use_temporary_parameter = True
+    if weight_in_cpu is True:
+        fc1.weight = fc1.weight.to(scales.device)
+        fc2.weight = fc2.weight.to(scales.device)
+
+    # import ipdb;ipdb.set_trace()
+    if num_key_value_groups > 1:
+        # import ipdb;ipdb.set_trace()
+        if args.gqa_scales == "copy":
+            kv_scales = scales
+            kv_shift = shifts
+            scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+            shifts = shifts.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+        elif args.gqa_scales == "mean":
+            kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+            kv_shift = shifts.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+        else:
+            raise NotImplementedError("Only implemented copy and mean for gqa")
+    else:
+        kv_scales = scales
+        kv_shift = shifts
+
+    if hasattr(fc1, 'temp_weight'):
+        fc1.temp_bias = fc1.temp_bias - kv_shift
+        fc1.temp_bias = fc1.temp_bias/kv_scales.view(-1)
+        fc1.temp_weight = fc1.temp_weight/kv_scales.view(-1,1)
+    else:
+        fc1.temp_bias = fc1.bias/kv_scales.view(-1)
+        fc1.temp_weight = fc1.weight/kv_scales.view(-1,1)
+    
+    if hasattr(fc2, 'bias') and fc2.bias is not None:
+        fc2.temp_bias = fc2.bias + fc2.weight@shifts
+    else:
+        fc2.temp_bias = fc2.weight@shifts
+    fc2.temp_weight = fc2.weight * scales.view(1,-1)
+
+
+    if weight_in_cpu is True:
+        fc1.weight = fc1.weight.cpu()
+        fc2.weight = fc2.weight.cpu()
+
+def smooth_fc_fc_inplace(fc1, fc2, scales,shifts=None,num_key_value_groups=1,head_dim=128,args=None):
+    # only support for v_proj and out_proh now.
+    fc1.use_temporary_parameter = False
+    fc2.use_temporary_parameter = False
+
+    if num_key_value_groups > 1:
+        # import ipdb;ipdb.set_trace()
+        if args.gqa_scales == "copy":
+            kv_scales = scales
+            kv_shift = shifts
+            scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+            shifts = shifts.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+        elif args.gqa_scales == "mean":
+            kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+            kv_shift = shifts.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+        else:
+           raise NotImplementedError("Only implemented copy and mean for gqa")
+    else:
+        kv_scales = scales
+        kv_shift = shifts
+
+
+    fc1.bias.sub_(kv_shift)
+    fc1.bias.div_(kv_scales.view(-1))
+    fc1.weight.div_(kv_scales.view(-1,1))
+    
+
+    if hasattr(fc2, 'bias') and fc2.bias is not None:
+        fc2.bias.add_(fc2.weight@shifts)
+    else:
+        del fc2.bias
+        fc2.register_buffer('bias',fc2.weight@shifts)
+    fc2.weight.mul_(scales.view(1,-1))
+
+
+
+def smooth_q_k_temporary(q_proj, k_proj,scales,num_key_value_groups=1,head_dim=128,weight_in_cpu=False,args=None):
+    q_proj.use_temporary_parameter = True
+    k_proj.use_temporary_parameter = True
+
+    if weight_in_cpu is True:
+        q_proj.weight = q_proj.weight.to(scales.device)
+        k_proj.weight = k_proj.weight.to(scales.device)
+
+    if num_key_value_groups > 1:
+        # import ipdb;ipdb.set_trace()
+        if args.gqa_scales == "copy":
+            kv_scales = scales
+            scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+        elif args.gqa_scales == "mean":
+            kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+        else:
+           raise NotImplementedError("Only implemented copy and mean for gqa")
+    else:
+        kv_scales = scales
+
+    q_proj.temp_weight = q_proj.temp_weight/scales.view(-1,1)
+    q_proj.temp_bias = q_proj.temp_bias/scales.view(-1)
+
+    k_proj.temp_weight = k_proj.temp_weight*kv_scales.view(-1,1)
+    k_proj.temp_bias = k_proj.temp_bias*kv_scales.view(-1)
+
+    if weight_in_cpu is True:
+        q_proj.weight = q_proj.weight.cpu()
+        k_proj.weight = k_proj.weight.cpu()
+
+
+def smooth_q_k_inplace(q_proj, k_proj, scales,num_key_value_groups=1,head_dim=128,args=None):
+    q_proj.use_temporary_parameter = False
+    k_proj.use_temporary_parameter = False
+
+    if num_key_value_groups > 1:
+        # import ipdb;ipdb.set_trace()
+        if args.gqa_scales == "copy":
+            kv_scales = scales
+            scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
+        elif args.gqa_scales == "mean":
+            kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
+        else:
+           raise NotImplementedError("Only implemented copy and mean for gqa")
+    else:
+        kv_scales = scales
+
+    q_proj.weight.div_(scales.view(-1,1))
+    q_proj.bias.div_(scales.view(-1))
+    k_proj.weight.mul_(kv_scales.view(-1,1))
+    k_proj.bias.mul_(kv_scales.view(-1))