diff --git a/README.md b/README.md index 795178c..45b301c 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ Defuser currently supports the following `transformers>=5.3.0` `model_type` valu | Pattern | Supported model types | Defused op performed ⚙️ | | --- |-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| --- | | Standard routed expert tensors 🧱 | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `laguna`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open`, `solar_open2` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. | -| Mixed sparse and shared experts | `deepseek_v3`, `deepseek_v32`, `deepseek_v4`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text` | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. | +| Mixed sparse and shared experts | `deepseek_v3`, `deepseek_v32`, `deepseek_v4`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text`, `qwen4_exp`, `qwen4_exp_text` | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. | | Transposed or packed expert tensors | `gpt_oss`, `phimoe` | Splits transposed fused expert `gate_up_proj` tensors into per-expert `gate_proj` + `up_proj`, preserves expert bias when present, and converts expert tensors into numbered expert `nn.Linear` modules. | | Flattened expert layout | `dbrx` | Rebuilds the flattened DBRX expert FFN weights into numbered expert `gate_proj`, `up_proj`, and `down_proj` `nn.Linear` modules. | | Batched expert-input execution | `llama4` | Runtime expert tensor defusion plus preservation of the llama4 batched expert-input execution contract. | diff --git a/defuser/model_registry.py b/defuser/model_registry.py index a9d14a0..39a620d 100644 --- a/defuser/model_registry.py +++ b/defuser/model_registry.py @@ -146,6 +146,12 @@ class PATCH(str, Enum): ), ], }, + "qwen4_exp": { + "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, + }, + "qwen4_exp_text": { + "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, + }, "glm4_moe": { "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, PATCH.REPLACE_MODULE: [ diff --git a/tests/test_candidate_coverage.py b/tests/test_candidate_coverage.py index 5208ddc..716bb12 100644 --- a/tests/test_candidate_coverage.py +++ b/tests/test_candidate_coverage.py @@ -603,6 +603,8 @@ def _standard_hidden(case: dict) -> torch.Tensor: "qwen3_next", "qwen3_omni_moe", "qwen3_vl_moe", + "qwen4_exp", + "qwen4_exp_text", "solar_open", "solar_open2", "zamba2", diff --git a/tests/test_convert_model.py b/tests/test_convert_model.py index 19b40da..ec2e6b6 100644 --- a/tests/test_convert_model.py +++ b/tests/test_convert_model.py @@ -38,6 +38,8 @@ Qwen3NextSparseMoeBlock, ) from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeConfig +from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig +from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpForCausalLM from transformers.models.gpt_oss.modeling_gpt_oss import GptOssConfig, GptOssForCausalLM from transformers.models.phimoe.modeling_phimoe import PhimoeConfig, PhimoeForCausalLM from transformers.models.llama4.modeling_llama4 import Llama4Config, Llama4ForConditionalGeneration @@ -176,6 +178,39 @@ def _tiny_qwen3_5_moe_config(): ) +def _tiny_qwen4_exp_text_config(): + return Qwen4ExpTextConfig( + vocab_size=64, + hidden_size=16, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + moe_intermediate_size=8, + shared_expert_intermediate_size=8, + num_experts=3, + num_experts_per_tok=2, + max_position_embeddings=32, + layer_types=["linear_attention", "qwen_sparse_attention"], + linear_num_key_heads=2, + linear_num_value_heads=4, + linear_key_head_dim=4, + linear_value_head_dim=4, + linear_conv_kernel_dim=2, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=4, + indexer_budget=8, + indexer_compress_ratio=2, + ) + + def _tiny_mixtral_config(): return MixtralConfig( vocab_size=128, @@ -739,6 +774,24 @@ def test_qwen3_5_moe(): torch.testing.assert_close(expert0.down_proj.weight, expected_down) +def test_qwen4_exp_text_preserves_forward(): + model = Qwen4ExpForCausalLM(_tiny_qwen4_exp_text_config()).eval() + input_ids = torch.tensor([[1, 7, 8, 2]]) + + with torch.inference_mode(): + expected = model(input_ids=input_ids, use_cache=False).logits + + assert convert_model(model, cleanup_original=False) + assert model.config.model_type == "qwen4_exp_text" + experts = model.model.layers[0].mlp.experts + assert hasattr(experts, "0") + assert not hasattr(experts, "gate_up_proj") + + with torch.inference_mode(): + actual = model(input_ids=input_ids, use_cache=False).logits + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-7) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_qwen3_5_moe_runtime_defusion_preserves_cuda_device(): model = Qwen3_5MoeForConditionalGeneration(_tiny_qwen3_5_moe_config()).cuda().eval() diff --git a/tests/test_meta_model_defusion.py b/tests/test_meta_model_defusion.py index cc942f2..3dd96cf 100644 --- a/tests/test_meta_model_defusion.py +++ b/tests/test_meta_model_defusion.py @@ -750,6 +750,28 @@ def _validate_defused_module(case: dict, module) -> None: "validator": "experts", "min_targets": 2, }, + { + "model_type": "qwen4_exp", + "mode": "convert", + "model_module": "transformers.models.qwen4_exp.modeling_qwen4_exp", + "model_class": "Qwen4ExpForConditionalGeneration", + "config_module": "transformers.models.qwen4_exp.configuration_qwen4_exp", + "config_class": "Qwen4ExpConfig", + "target_class_paths": ("transformers.models.qwen4_exp.modeling_qwen4_exp.Qwen4ExpTextExperts",), + "validator": "experts", + "min_targets": 2, + }, + { + "model_type": "qwen4_exp_text", + "mode": "convert", + "model_module": "transformers.models.qwen4_exp.modeling_qwen4_exp", + "model_class": "Qwen4ExpForCausalLM", + "config_module": "transformers.models.qwen4_exp.configuration_qwen4_exp", + "config_class": "Qwen4ExpTextConfig", + "target_class_paths": ("transformers.models.qwen4_exp.modeling_qwen4_exp.Qwen4ExpTextExperts",), + "validator": "experts", + "min_targets": 2, + }, { "model_type": "qwen3_moe", "mode": "replace",