Add Qwen model (#182)

Co-authored-by: Casper Hansen <[email protected]>

awq/models/__init__.py

@@ -8,4 +8,5 @@
 from .mistral import MistralAWQForCausalLM
 from .gpt_neox import GPTNeoXAWQForCausalLM
 from .aquila import AquilaAWQForCausalLM
-from .yi import YiAWQForCausalLM
+from .yi import YiAWQForCausalLM
+from .qwen import QwenAWQForCausalLM

awq/models/auto.py

@@ -16,7 +16,8 @@
     "mistral": MistralAWQForCausalLM,
     "gpt_neox": GPTNeoXAWQForCausalLM,
     "aquila": AquilaAWQForCausalLM,
-    "Yi": YiAWQForCausalLM
+    "Yi": YiAWQForCausalLM,
+    "qwen": QwenAWQForCausalLM
 }
 
 def check_and_get_model_type(model_dir, trust_remote_code=True):

awq/models/qwen.py

@@ -0,0 +1,55 @@
+from .base import BaseAWQForCausalLM
+
+
+class QwenAWQForCausalLM(BaseAWQForCausalLM):
+    layer_type = "QWenBlock"
+    max_new_tokens_key = "seq_length"
+
+    @staticmethod
+    def get_model_layers(model):
+        return model.transformer.h
+
+    @staticmethod
+    def get_act_for_scaling(module):
+        return dict(is_scalable=False)
+
+    @staticmethod
+    def move_embed(model, device: str):
+        model.transformer.wte = model.transformer.wte.to(device)
+        model.transformer.rotary_emb = model.transformer.rotary_emb.to(device)
+
+    @staticmethod
+    def get_layers_for_scaling(module, input_feat, module_kwargs):
+        layers = []
+
+        # attention
+        layers.append(
+            dict(
+                prev_op=module.ln_1,
+                layers=[module.attn.c_attn, module.attn.c_proj],
+                inp=input_feat["attn.c_attn"],
+                module2inspect=module.attn,
+                kwargs=module_kwargs,
+            )
+        )
+
+        # mlp
+        layers.append(
+            dict(
+                prev_op=module.ln_2,
+                layers=[module.mlp.w2, module.mlp.w1],
+                inp=input_feat["mlp.w2"],
+                module2inspect=module.mlp,
+            )
+        )
+
+        # linear 2
+        layers.append(
+            dict(
+                prev_op=module.mlp.w1,
+                layers=[module.mlp.c_proj],
+                inp=input_feat["mlp.c_proj"],
+            )
+        )
+
+        return layers