fix: update qwen3_omni with transformers>=5.0 (#296)

Author: ali-88123

angelslim/compressor/quant/core/config.py

@@ -66,6 +66,7 @@ def __init__(self, config, global_config=None):
         kv_cache_quant_method = quantization_args.quant_method.get("kv_cache", None)
         self.cpu_convert = quantization_args.cpu_convert
         self.save_name = quantization_args.save_name
+        self.quant_talker = getattr(quantization_args, "quant_talker", False)
 
         if global_config:
             self.max_seq_length = global_config.max_seq_length

angelslim/data/dataloader.py

@@ -45,6 +45,7 @@ def create_data_loader(
         model_name: str = None,
         quantization_config: str = None,
         is_sft_data: bool = False,
+        dtype=None,
     ) -> DataLoader:
         """
         Create appropriate DataLoader based on data source
@@ -114,6 +115,7 @@ def create_data_loader(
                 data_source=data_source,
                 is_hf_dataset=not os.path.isfile(data_source),
                 use_audio_in_video=use_audio_in_video,
+                dtype=dtype,
             )
         elif data_type == "AudioDataset":
             dataset = AudioDataset(

angelslim/data/omni_dataset.py

@@ -35,10 +35,12 @@ def __init__(
         data_source: Union[str, Dict] = None,
         is_hf_dataset: bool = False,
         use_audio_in_video: bool = False,
+        dtype=None,
     ):
         super().__init__(processor, device, max_length)
         self.is_hf_dataset = is_hf_dataset
         self.use_audio_in_video = use_audio_in_video
+        self.dtype = dtype
 
         self._load_file_based_dataset(data_source, num_samples)
 
@@ -112,10 +114,11 @@ def _process_and_append(self, messages: List[Dict]):
         inputs = self.processor(
             text=text,
             images=images,
-            audios=audios,
+            audio=audios,
             videos=videos,
             padding=True,
             return_tensors="pt",
             use_audio_in_video=self.use_audio_in_video,
         )
+        inputs = inputs.to(self.device).to(self.dtype)
         self.data.append(inputs)

angelslim/engine.py

@@ -160,6 +160,7 @@ def prepare_data(
         model_name=None,
         quantization_config=None,
         is_sft_data=False,
+        dtype=None,
     ) -> Optional[Any]:
         """Prepare compression dataset"""
         if custom_dataloader is not None:
@@ -187,6 +188,7 @@ def prepare_data(
             model_name=model_name,
             quantization_config=quantization_config,
             is_sft_data=is_sft_data,
+            dtype=dtype,
         )
         self.max_seq_length = max_length
 

angelslim/models/omni/qwen3_omni.py

@@ -19,10 +19,17 @@
     AutoTokenizer,
     Qwen3OmniMoeForConditionalGeneration,
 )
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeTalkerTextExperts,
+    Qwen3OmniMoeTalkerTextTopKRouter,
+    Qwen3OmniMoeThinkerTextExperts,
+    Qwen3OmniMoeThinkerTextTopKRouter,
+)
 
 from ...compressor.quant.core import PTQVLMSaveVllmHF
-from ...utils import find_layers, print_info
+from ...utils import find_layers, find_parent_layer_and_sub_name, print_info
 from ..base_model import BaseLLMModel
+from ..llm.qwen import QwenMoeExpertsWithLinear
 from ..model_factory import SlimModelFactory
 
 
@@ -38,7 +45,92 @@ def __init__(
             deploy_backend=deploy_backend,
         )
         self.modal_type = "Omni"
-        self.block_name = ["thinker.model.layers", "talker.model.layers"]
+        self.thinker_block_name = "thinker.model.layers"
+        self.talker_block_name = "talker.model.layers"
+        self.observer_layer_classes = [
+            torch.nn.Linear,
+            Qwen3OmniMoeThinkerTextTopKRouter,
+            Qwen3OmniMoeTalkerTextTopKRouter,
+        ]
+        self.observed_names = [
+            "k_proj",
+            "v_proj",
+            "q_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ]
+
+    def replace_moe(self):
+        for name, module in self.model.thinker.named_modules():
+            if isinstance(module, Qwen3OmniMoeThinkerTextExperts) and not isinstance(
+                module, QwenMoeExpertsWithLinear
+            ):
+                print(name)
+                parent_layer, sub_name = find_parent_layer_and_sub_name(self.model.thinker, name)
+                moe_linear = QwenMoeExpertsWithLinear(module)
+                del module
+                setattr(parent_layer, sub_name, moe_linear)
+
+        for name, module in self.model.talker.named_modules():
+            if isinstance(module, Qwen3OmniMoeTalkerTextExperts) and not isinstance(
+                module, QwenMoeExpertsWithLinear
+            ):
+                print(name)
+                parent_layer, sub_name = find_parent_layer_and_sub_name(self.model.talker, name)
+                moe_linear = QwenMoeExpertsWithLinear(module)
+                del module
+                setattr(parent_layer, sub_name, moe_linear)
+
+    def _patch_inputs_embeds_generate_device(self, module):
+        if module is None or getattr(module, "_angelslim_generate_device_patch", False):
+            return
+
+        original_generate = module.generate
+        skip_keys = {"past_key_values", "encoder_outputs"}
+
+        def move_to_target_device(value, target_device):
+            if isinstance(value, torch.Tensor):
+                if value.device.type == "meta" or value.device == target_device:
+                    return value
+                return value.to(target_device)
+            if isinstance(value, tuple):
+                return tuple(move_to_target_device(item, target_device) for item in value)
+            if isinstance(value, list):
+                return [move_to_target_device(item, target_device) for item in value]
+            if isinstance(value, dict):
+                return {
+                    key: item if key in skip_keys else move_to_target_device(item, target_device)
+                    for key, item in value.items()
+                }
+            return value
+
+        def generate_on_module_device(*args, **kwargs):
+            inputs_embeds = kwargs.get("inputs_embeds")
+            if inputs_embeds is not None:
+                target_device = getattr(module, "device", inputs_embeds.device)
+                if target_device.type == "meta":
+                    target_device = inputs_embeds.device
+
+                kwargs = {
+                    key: value if key in skip_keys else move_to_target_device(value, target_device)
+                    for key, value in kwargs.items()
+                }
+
+            return original_generate(*args, **kwargs)
+
+        module.generate = generate_on_module_device
+        module._angelslim_generate_device_patch = True
+
+    def _patch_omni_generate_devices(self):
+        talker = getattr(self.model, "talker", None)
+        self._patch_inputs_embeds_generate_device(talker)
+        self._patch_inputs_embeds_generate_device(getattr(talker, "code_predictor", None))
+
+    def init_ptq(self, slim_config):
+        super().init_ptq(slim_config)
+        self.replace_moe()
 
     def from_pretrained(
         self,
@@ -63,6 +155,7 @@ def from_pretrained(
                 device_map=device_map,
                 attn_implementation=attn_implementation,
             )
+        self._patch_omni_generate_devices()
 
         # Load tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(
@@ -74,24 +167,21 @@ def from_pretrained(
             model_path, trust_remote_code=trust_remote_code
         )
 
-    def get_observer_layers(self):
-        names = [
-            "k_proj",
-            "v_proj",
-            "q_proj",
-            "o_proj",
-            "up_proj",
-            "gate_proj",
-            "down_proj",
-        ]
+    def _get_quant_block_names(self):
+        block_names = [self.thinker_block_name]
+        if getattr(self.quant_config, "quant_talker", True):
+            block_names.append(self.talker_block_name)
+        return block_names
 
+    def get_observer_layers(self):
         observer_layers_dict = {}
         layers_dict = find_layers(self.model, layers=self.observer_layer_classes)
+        block_names = self._get_quant_block_names()
 
         ignore_layers = self.skip_layer_names()
         for name, module in layers_dict.items():
-            block_condition = any(name.startswith(block) for block in self.block_name)
-            if block_condition and name.split(".")[-1] in names:
+            block_condition = any(name.startswith(block) for block in block_names)
+            if block_condition and name.split(".")[-1] in self.observed_names:
                 observer_layers_dict[name] = module
             else:
                 ignore_layers.append(name)
@@ -106,10 +196,11 @@ def get_observer_layers(self):
 
     def get_kvcache_observer_layers_names(self, observe_names):
         names = ["self_attn.k_proj", "self_attn.v_proj"]
+        block_names = self._get_quant_block_names()
         return [
             k
             for k in observe_names
-            if any(k.startswith(block) for block in self.block_name)
+            if any(k.startswith(block) for block in block_names)
             and k.split(".")[-2] + "." + k.split(".")[-1] in names
         ]
 
@@ -129,10 +220,9 @@ def model_forward(self, dataloader, **kwargs):
         if dataloader is not None:
             with torch.no_grad():
                 for batch in tqdm(dataloader, desc="calibrating...", total=len(dataloader)):
-                    inputs = {k: v.to(device) for k, v in batch.items()}
                     try:
                         text_ids, audio = self.model.generate(
-                            **inputs, use_audio_in_video=self.use_audio_in_video
+                            **batch, use_audio_in_video=self.use_audio_in_video
                         )
                         calibrated_cnt += 1
                     except ValueError:

angelslim/utils/config_parser.py

@@ -200,6 +200,7 @@ class QuantizationConfig:
         quant_method: Algorithm used for quantization
         modules_to_quantize: List of module types to quantize
         ignore_layers: List of layer names to skip
+        quant_talker: Whether to quantize Qwen3-Omni talker LLM module
     """
 
     name: str = field(default="fp8_dynamic")
@@ -222,6 +223,7 @@ class QuantizationConfig:
     ignore_layers: List[str] = field(default_factory=list)
     quant_analyse: bool = field(default=False)
     quant_vit: bool = field(default=False)
+    quant_talker: bool = field(default=False)
     # DAQ-specific fields
     base_model_path: Optional[str] = field(default=None)
     base_is_fp8: bool = field(default=False)

configs/qwen3_omni/fp8_dynamic/qwen3_omni_fp8_dynamic.yaml

@@ -23,3 +23,4 @@ compression:
     quant_method:
       weight: "per-tensor"
       activation: "per-tensor"
+    quant_talker: false  # Whether to quantize Qwen3-Omni talker LLM module

configs/qwen3_omni/fp8_static/qwen3_omni_fp8_static.yaml

@@ -23,6 +23,7 @@ compression:
     quant_method:
       weight: "per-tensor"
       activation: "per-tensor"
+    quant_talker: false  # Whether to quantize Qwen3-Omni talker LLM module
 
 # Dataset for calibration
 dataset:

dataset/omni_fake_data/audios/0.wav

@@ -1,3 +1,3 @@
-{"messages": [{"role": "user", "content": "What happens after the text disappears from the screen?"}], "video_path": "./videos/0.mp4"}
-{"messages": [{"role": "user", "content": "How many food item is shown in the bar graph?"}], "image_path": "./images/0.png"}
-{"messages": [{"role": "user", "content": "Why is the speech described as rich in frequency content?"}], "audio_path": "./audios/0.png"}
+{"messages": [{"role": "user", "content": "描述这个视频的内容。"}], "video_path": "./videos/0.mp4"}
+{"messages": [{"role": "user", "content": "请描述这张图片的内容。"}], "image_path": "./images/0.png"}
+{"messages": [{"role": "user", "content": "请将这段语音转写成文字。"}], "audio_path": "./audios/0.wav"}