extend model support

Author: Qubitium

defuser/model_registry.py

@@ -117,6 +117,10 @@ class PATCH(str, Enum):
             (
                 "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeThinkerTextSparseMoeBlock",
                 "defuser.modeling.unfused_moe.qwen3_omni_moe.LinearQwen3OmniMoeThinkerTextSparseMoeBlock",
+            ),
+            (
+                "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeTalkerTextSparseMoeBlock",
+                "defuser.modeling.unfused_moe.qwen3_omni_moe.LinearQwen3OmniMoeTalkerTextSparseMoeBlock",
             )
         ],
     },

defuser/modeling/model_patches.py

@@ -45,15 +45,18 @@ def decorator(func: Callable):
 def patch_qwen3_omni_text_class() -> list[str]:
     """Teach HF init code how to initialize unfused qwen3-omni thinker experts."""
     from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoePreTrainedModel
-    from defuser.modeling.unfused_moe.qwen3_omni_moe import LinearQwen3OmniMoeThinkerTextSparseMoeBlock
+    from defuser.modeling.unfused_moe.qwen3_omni_moe import (
+        LinearQwen3OmniMoeTalkerTextSparseMoeBlock,
+        LinearQwen3OmniMoeThinkerTextSparseMoeBlock,
+    )
     orig_init_weights = Qwen3OmniMoePreTrainedModel._init_weights
 
     def patched_init_weights(self, module):
         try:
             orig_init_weights(self, module)
         except AttributeError as e:
             # fallback for unfused experts
-            if isinstance(module, LinearQwen3OmniMoeThinkerTextSparseMoeBlock):
+            if isinstance(module, (LinearQwen3OmniMoeThinkerTextSparseMoeBlock, LinearQwen3OmniMoeTalkerTextSparseMoeBlock)):
                 std = self.config.initializer_range
                 experts = module.experts
 
@@ -63,9 +66,18 @@ def patched_init_weights(self, module):
                     torch.nn.init.normal_(experts.up_proj.weight, 0.0, std)
                 if hasattr(experts, "down_proj"):
                     torch.nn.init.normal_(experts.down_proj.weight, 0.0, std)
+                if isinstance(experts, torch.nn.ModuleList):
+                    for expert in experts:
+                        torch.nn.init.normal_(expert.gate_proj.weight, 0.0, std)
+                        torch.nn.init.normal_(expert.up_proj.weight, 0.0, std)
+                        torch.nn.init.normal_(expert.down_proj.weight, 0.0, std)
 
                 if hasattr(module, "gate"):
                     torch.nn.init.normal_(module.gate.weight, 0.0, std)
+                if hasattr(module, "shared_expert"):
+                    module.shared_expert._is_hf_initialized = True
+                if hasattr(module, "shared_expert_gate"):
+                    torch.nn.init.normal_(module.shared_expert_gate.weight, 0.0, std)
             else:
                 raise e
 

defuser/modeling/unfused_moe/qwen3_omni_moe.py

@@ -46,3 +46,47 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         )
         final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
         return final_hidden_states
+
+
+class LinearQwen3OmniMoeTalkerTextSparseMoeBlock(nn.Module):
+    """Text talker MoE block for qwen3-omni with explicit per-expert modules."""
+
+    def __init__(self, config):
+        super().__init__()
+        from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+            Qwen3OmniMoeTalkerTextMLP,
+            Qwen3OmniMoeTalkerTextTopKRouter,
+        )
+
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+
+        self.gate = Qwen3OmniMoeTalkerTextTopKRouter(config)
+        self.experts = nn.ModuleList(
+            [
+                Qwen3OmniMoeTalkerTextMLP(config, intermediate_size=config.moe_intermediate_size)
+                for _ in range(self.num_experts)
+            ]
+        )
+        self.shared_expert = Qwen3OmniMoeTalkerTextMLP(
+            config,
+            intermediate_size=config.shared_expert_intermediate_size,
+        )
+        self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_expert_output = self.shared_expert(hidden_states)
+        _, routing_weights, selected_experts = self.gate(hidden_states)
+        final_hidden_states = run_routed_experts(
+            self.experts,
+            hidden_states,
+            routing_weights.to(hidden_states.dtype),
+            selected_experts,
+            self.num_experts,
+        )
+        shared_expert_output = torch.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
+        final_hidden_states = final_hidden_states + shared_expert_output
+        return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)