Support Aquila models. (#123)

Signed-off-by: ldwang <ftgreat@gmail.com>
Co-authored-by: ldwang <ftgreat@gmail.com>
Co-authored-by: Casper Hansen <casperbh.96@gmail.com>

README.md

@@ -80,6 +80,8 @@ The detailed support list:
 | OPT      | 125m/1.3B/2.7B/6.7B/13B/30B |
 | Bloom    | 560m/3B/7B/                 |
 | GPTJ     | 6.7B                        |
+| Aquila   | 7B                          |
+| Aquila2  | 7B/34B                      |
 
 ## Usage
 

awq/models/__init__.py

@@ -7,3 +7,4 @@ from .gptj import GPTJAWQForCausalLM
 from .gpt_bigcode import GptBigCodeAWQForCausalLM
 from .mistral import MistralAWQForCausalLM
 from .gpt_neox import GPTNeoXAWQForCausalLM
+from .aquila import AquilaAWQForCausalLM

awq/models/aquila.py

+## Reference from llama.py
+from .base import BaseAWQForCausalLM
+from typing import Dict
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer as AquilaDecoderLayer,
+    LlamaForCausalLM as AquilaForCausalLM,
+    LlamaAttention as AquilaAttention,
+    LlamaRMSNorm as AquilaRMSNorm,
+    LlamaMLP as AquilaMLP
+)
+
+class AquilaAWQForCausalLM(BaseAWQForCausalLM):
+    layer_type = "AquilaDecoderLayer"
+    max_new_tokens_key = "max_position_embeddings"
+
+    @staticmethod
+    def fuse_layers(model: AquilaForCausalLM, quant_config: Dict):
+        fuser = AquilaFuser(model, quant_config)
+        fuser.fuse_attention()
+        fuser.fuse_rmsnorm()
+        fuser.fuse_mlp()
+
+    @staticmethod
+    def get_model_layers(model: AquilaForCausalLM):
+        return model.model.layers
+    
+    @staticmethod
+    def get_act_for_scaling(module: AquilaDecoderLayer):
+        return dict(
+            is_scalable=False
+        )
+    
+    @staticmethod
+    def move_embed(model: AquilaForCausalLM, device: str):
+        model.model.embed_tokens = model.model.embed_tokens.to(device)
+    
+    @staticmethod
+    def get_layers_for_scaling(module: AquilaDecoderLayer, input_feat, module_kwargs):
+        layers = []
+
+        # attention input
+        layers.append(dict(
+            prev_op=module.input_layernorm,
+            layers=[module.self_attn.q_proj,
+                    module.self_attn.k_proj, module.self_attn.v_proj],
+            inp=input_feat['self_attn.q_proj'],
+            module2inspect=module.self_attn, kwargs=module_kwargs,
+        ))
+
+        # attention out
+        # Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696
+        if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape:
+            layers.append(dict(
+                prev_op=module.self_attn.v_proj,
+                layers=[module.self_attn.o_proj],
+                inp=input_feat['self_attn.o_proj'],
+            ))
+        
+        # linear 1
+        layers.append(dict(
+            prev_op=module.post_attention_layernorm,
+            layers=[module.mlp.gate_proj, module.mlp.up_proj],
+            inp=input_feat['mlp.gate_proj'],
+            module2inspect=module.mlp,
+        ))
+
+        # linear 2
+        layers.append(dict(
+            prev_op=module.mlp.up_proj,
+            layers=[module.mlp.down_proj],
+            inp=input_feat['mlp.down_proj'],
+        ))
+
+        return layers
+
+import torch
+from typing import List, Tuple, Union
+from awq.utils.utils import set_module_name
+from awq.modules.fused.mlp import QuantLlamaMLP
+from awq.modules.fused.attn import QuantAttentionFused
+from awq.modules.fused.norm import FasterTransformerRMSNorm
+from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV
+
+class AquilaFuser:
+    def __init__(self, model, quant_config):
+        self.model = model
+        self.quant_config = quant_config
+
+        self.attention_modules: List[Tuple[str, AquilaAttention]] = [
+            (name, module) for name, module in self.model.named_modules()
+            if "AquilaAttention".lower() in module.__class__.__name__.lower()
+        ]
+
+        self.rmsnorm_modules: List[Tuple[str, AquilaRMSNorm]] = [
+            (name, module) for name, module in self.model.named_modules()
+            if "AquilaRMSNorm".lower() in module.__class__.__name__.lower()
+        ]
+        
+        self.mlp_modules: List[Tuple[str, AquilaMLP]] = [
+            (name, module) for name, module in self.model.named_modules()
+            if "AquilaMLP".lower() in module.__class__.__name__.lower()
+        ]
+    
+    def fuse_attention(self):
+        for name, module in self.attention_modules:
+            qkv_layer: Union[WQLinear_GEMM, WQLinear_GEMV] = self._fuse_qkv(module)
+            attn = QuantAttentionFused(
+                module.hidden_size,
+                module.num_heads,
+                module.num_key_value_heads,
+                qkv_layer, 
+                module.o_proj,
+                next(iter(qkv_layer.state_dict().values())).device,
+                self.model.config.max_new_tokens
+            )
+            set_module_name(self.model, name, attn)
+    
+    def _fuse_qkv(self, module: AquilaAttention):
+        q_proj, k_proj, v_proj = module.q_proj, module.k_proj, module.v_proj
+        bias = torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0) if q_proj.bias is not None else None
+
+        if isinstance(q_proj, WQLinear_GEMV):
+            q_linear = WQLinear_GEMV
+        else:
+            q_linear = WQLinear_GEMM
+
+        qkv_layer = q_linear(
+            q_proj.w_bit,
+            q_proj.group_size,
+            q_proj.in_features,
+            q_proj.out_features + k_proj.out_features + v_proj.out_features,
+            q_proj.bias is not None,
+            next(iter(module.state_dict().values())).device
+        )
+
+        if isinstance(qkv_layer, WQLinear_GEMV):
+            qkv_layer.qweight = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=0)
+            qkv_layer.qzeros = torch.cat([q_proj.qzeros, k_proj.qzeros, v_proj.qzeros], dim=0)
+            qkv_layer.scales = torch.cat([q_proj.scales, k_proj.scales, v_proj.scales], dim=0)
+            qkv_layer.split_k_iters = q_proj.split_k_iters
+        else:
+            qkv_layer.qweight = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=1)
+            qkv_layer.qzeros = torch.cat([q_proj.qzeros, k_proj.qzeros, v_proj.qzeros], dim=1)
+            qkv_layer.scales = torch.cat([q_proj.scales, k_proj.scales, v_proj.scales], dim=1)
+        
+        qkv_layer.bias = bias
+
+        return qkv_layer
+
+    def fuse_rmsnorm(self):
+        for name, module in self.rmsnorm_modules:
+            norm = FasterTransformerRMSNorm(module.weight, module.variance_epsilon)
+            set_module_name(self.model, name, norm)
+
+    def fuse_mlp(self):
+        for name, module in self.mlp_modules:
+            mlp = QuantLlamaMLP(module.gate_proj, module.down_proj, module.up_proj)
+            set_module_name(self.model, name, mlp)

awq/models/auto.py

@@ -15,6 +15,7 @@ AWQ_CAUSAL_LM_MODEL_MAP = {
     "gpt_bigcode": GptBigCodeAWQForCausalLM,
     "mistral": MistralAWQForCausalLM,
     "gpt_neox": GPTNeoXAWQForCausalLM,
+    "aquila": AquilaAWQForCausalLM,
 }
 
 def check_and_get_model_type(model_dir, trust_remote_code=True):