Add Qwen2.5-Omni support. (#759)

Signed-off-by: wenhu.xwh <wenhu.xwh@alibaba-inc.com>
Co-authored-by: wenhu.xwh <wenhu.xwh@alibaba-inc.com>

Author: tiger-of-shawn

awq/models/__init__.py

@@ -31,4 +31,5 @@
 from .internlm2 import InternLM2AWQForCausalLM
 from .minicpm3 import MiniCPM3AWQForCausalLM
 from .qwen2vl import Qwen2VLAWQForCausalLM
-from .qwen2_5_vl import Qwen2_5_VLAWQForCausalLM
+from .qwen2_5_vl import Qwen2_5_VLAWQForCausalLM
+from .qwen2_5_omni import Qwen2_5_OmniAWQForConditionalGeneration

awq/models/auto.py

@@ -43,6 +43,7 @@
     "minicpm3": MiniCPM3AWQForCausalLM,
     "qwen2_vl": Qwen2VLAWQForCausalLM,
     "qwen2_5_vl": Qwen2_5_VLAWQForCausalLM,
+    "qwen2_5_omni": Qwen2_5_OmniAWQForConditionalGeneration
 }
 
 
@@ -79,7 +80,6 @@ def from_pretrained(
         model_type = check_and_get_model_type(
             model_path, trust_remote_code, **model_init_kwargs
         )
-
         return AWQ_CAUSAL_LM_MODEL_MAP[model_type].from_pretrained(
             model_path,
             model_type,

awq/models/base.py

@@ -90,6 +90,7 @@
     "internlm2": "AutoModelForCausalLM",
     "qwen2_vl": "AutoModelForVision2Seq",
     "qwen2_5_vl": "AutoModelForVision2Seq",
+    "qwen2_5_omni": "AutoModelForTextToWaveform",
 }
 
 
@@ -377,12 +378,11 @@ def from_pretrained(
         target_cls = getattr(transformers, target_cls_name)
 
         processor = None
-        if target_cls_name == "AutoModelForVision2Seq":
+        if target_cls_name == "AutoModelForVision2Seq" or target_cls_name == "AutoModelForTextToWaveform":
             processor = AutoProcessor.from_pretrained(model_weights_path)
-
         if model_init_kwargs.get("low_cpu_mem_usage") is None:
             model_init_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
-        if model_init_kwargs.get("use_cache") is None and target_cls_name != "AutoModelForVision2Seq":
+        if model_init_kwargs.get("use_cache") is None and not ((target_cls_name == "AutoModelForVision2Seq") or (target_cls_name == "AutoModelForTextToWaveform")):
             model_init_kwargs["use_cache"] = use_cache
 
         # If not quantized, must load with AutoModelForCausalLM

awq/models/qwen2_5_omni.py

@@ -0,0 +1,86 @@
+from .base import BaseAWQForCausalLM
+from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+        Qwen2_5OmniDecoderLayer
+    )
+    from transformers import Qwen2_5OmniForConditionalGeneration
+
+
+class Qwen2_5_OmniAWQForConditionalGeneration(BaseAWQForCausalLM):
+    layer_type = "Qwen2_5OmniDecoderLayer"
+    max_seq_len_key = "max_position_embeddings"
+    modules_to_not_convert = ["visual"]
+    @staticmethod
+    def get_model_layers(model: "Qwen2_5OmniForConditionalGeneration"):
+        return model.thinker.model.layers
+
+    @staticmethod
+    def get_act_for_scaling(module: "Qwen2_5OmniForConditionalGeneration"):
+        return dict(is_scalable=False)
+
+    @staticmethod
+    def move_embed(model: "Qwen2_5OmniForConditionalGeneration", device: str):
+        model.thinker.model.embed_tokens = model.thinker.model.embed_tokens.to(device)
+        model.thinker.visual = model.thinker.visual.to(device)
+        model.thinker.audio_tower = model.thinker.audio_tower.to(device)
+        
+        model.thinker.visual.rotary_pos_emb = model.thinker.visual.rotary_pos_emb.to(device)
+        model.thinker.model.rotary_emb = model.thinker.model.rotary_emb.to(device)
+        
+        for layer in model.thinker.model.layers:
+            layer.self_attn.rotary_emb = layer.self_attn.rotary_emb.to(device)
+        
+    @staticmethod
+    def get_layers_for_scaling(
+        module: "Qwen2_5OmniDecoderLayer", 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

awq/quantize/quantizer.py

@@ -163,6 +163,10 @@ def quantize(self):
                     self.inps, self.module_kwargs["position_ids"]
                 )
 
+            if (transformers.__version__ >= "4.48.0"
+                and self.module_kwargs.get('attention_mask') is None):
+                self.module_kwargs['attention_mask'] = None
+
             for k, v in self.module_kwargs.items():
                 # position embeddings found in tuple
                 if isinstance(v, tuple):