add support to qwen2_moe and qwen3_next (#21)

Author: Qubitium

defuser/model_registry.py

@@ -7,6 +7,9 @@
     "mixtral": {
         "min_transformers_version": "5.0.0",
     },
+    "qwen2_moe": {
+        "min_transformers_version": "5.0.0",
+    },
     "qwen3_moe": {
         "min_transformers_version": "5.0.0",
     },
@@ -16,4 +19,7 @@
     "qwen3_5_moe_text": {
         "min_transformers_version": "5.2.0",
     },
+    "qwen3_next": {
+        "min_transformers_version": "5.0.0",
+    },
 }

pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "Defuser"
-version = "0.0.8"
+version = "0.0.9"
 description = "Model defuser helper for HF Transformers."
 readme = "README.md"
 requires-python = ">=3.9"

tests/test_convert_model.py

@@ -4,11 +4,45 @@
 # Contact: qubitium@modelcloud.ai, x.com/qubitium
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText
+from transformers.models.qwen2_moe.modeling_qwen2_moe import Qwen2MoeConfig, Qwen2MoeForCausalLM
+from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextConfig, Qwen3NextForCausalLM
 
 from defuser import convert_model
 from defuser.modeling.replace_modules import materialize_model
 
 
+def _tiny_moe_config(config_cls):
+    return config_cls(
+        num_hidden_layers=1,
+        hidden_size=64,
+        intermediate_size=128,
+        moe_intermediate_size=32,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        num_experts=4,
+        num_experts_per_tok=2,
+        vocab_size=128,
+    )
+
+
+def _assert_unfused_expert_module(experts):
+    assert hasattr(experts, "0")
+    expert0 = getattr(experts, "0")
+    assert hasattr(expert0, "gate_proj")
+    assert hasattr(expert0, "up_proj")
+    assert hasattr(expert0, "down_proj")
+
+
+def test_qwen2_moe():
+    model = Qwen2MoeForCausalLM(_tiny_moe_config(Qwen2MoeConfig))
+    assert model.config.model_type == "qwen2_moe"
+
+    converted = convert_model(model, max_layers=1)
+    assert converted
+
+    _assert_unfused_expert_module(model.model.layers[0].mlp.experts)
+
+
 def test_qwen3_moe():
     model_id = "Qwen/Qwen3-30B-A3B"
     config = AutoConfig.from_pretrained(model_id)
@@ -24,12 +58,17 @@ def test_qwen3_moe():
     converted = convert_model(model, max_layers=1)
     assert converted
 
-    experts = model.model.layers[0].mlp.experts
-    assert hasattr(experts, "0")
-    expert0 = getattr(experts, "0")
-    assert hasattr(expert0, "gate_proj")
-    assert hasattr(expert0, "up_proj")
-    assert hasattr(expert0, "down_proj")
+    _assert_unfused_expert_module(model.model.layers[0].mlp.experts)
+
+
+def test_qwen3_next():
+    model = Qwen3NextForCausalLM(_tiny_moe_config(Qwen3NextConfig))
+    assert model.config.model_type == "qwen3_next"
+
+    converted = convert_model(model, max_layers=1)
+    assert converted
+
+    _assert_unfused_expert_module(model.model.layers[0].mlp.experts)
 
 
 def test_qwen3_5_moe():
@@ -60,11 +99,8 @@ def test_qwen3_5_moe():
     moe_block = model.model.language_model.layers[0].mlp
     experts = moe_block.experts
 
-    assert hasattr(experts, "0")
+    _assert_unfused_expert_module(experts)
     expert0 = getattr(experts, "0")
-    assert hasattr(expert0, "gate_proj")
-    assert hasattr(expert0, "up_proj")
-    assert hasattr(expert0, "down_proj")
 
     materialize_model(model.model.language_model.layers[0])
 
@@ -102,14 +138,11 @@ def test_mixtral():
     moe_block = model.model.layers[0].mlp
     experts = moe_block.experts
 
-    assert hasattr(experts, "0")
+    _assert_unfused_expert_module(experts)
     expert0 = getattr(experts, "0")
-    assert hasattr(expert0, "gate_proj")
-    assert hasattr(expert0, "up_proj")
-    assert hasattr(expert0, "down_proj")
 
     materialize_model(model.model.layers[0])
 
     torch.testing.assert_close(expert0.gate_proj.weight, expected_gate)
     torch.testing.assert_close(expert0.up_proj.weight, expected_up)
-    torch.testing.assert_close(expert0.down_proj.weight, expected_down)
+    torch.testing.assert_close(expert0.down_proj.weight, expected_down)