diff --git a/README.md b/README.md index 360eb58..2f8183f 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,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`. | +| Standard routed expert tensors 🧱 | `deepseek_v2`, `diffusion_gemma`, `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`, `glm5_next`, `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. | diff --git a/defuser/model_registry.py b/defuser/model_registry.py index 6dc674e..5058a33 100644 --- a/defuser/model_registry.py +++ b/defuser/model_registry.py @@ -46,6 +46,21 @@ class PATCH(str, Enum): "dia": { "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, }, + "diffusion_gemma": { + "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, + PATCH.EXPERTS_DEFUSE: [ + { + # DiffusionGemma's fused experts need an instance-level + # implementation after defusion so their per-expert Linear + # leaves remain compatible with the upstream routing call. + "module_class": ( + "transformers.models.diffusion_gemma.modeling_diffusion_gemma." + "DiffusionGemmaTextExperts" + ), + "forward_impl": "linear_loop", + } + ], + }, "dots1": { "min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION, }, diff --git a/pyproject.toml b/pyproject.toml index e65d0fc..e2c5655 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta" [project] name = "Defuser" -version = "0.0.27" +version = "0.0.28" description = "Model defuser helper for HF Transformers." readme = "README.md" requires-python = ">=3.9" diff --git a/tests/test_meta_model_defusion.py b/tests/test_meta_model_defusion.py index cab8a45..40a23d8 100644 --- a/tests/test_meta_model_defusion.py +++ b/tests/test_meta_model_defusion.py @@ -138,7 +138,49 @@ def _mutate_common_config_tree(config, visited: set[int] | None = None) -> None: def _build_model_config(case: dict): """Construct a small config tree for one registered public model type.""" - config = _load(case["config_module"], case["config_class"])() + config_cls = _load(case["config_module"], case["config_class"]) + if case["model_type"] == "diffusion_gemma": + # Unlike the generic cases, DiffusionGemma validates its nested text / + # vision configs during construction; provide a minimal valid pair so + # this registry test remains a meta-only structural check. + return config_cls( + text_config={ + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "moe_intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 1, + "num_global_key_value_heads": 1, + "head_dim": 16, + "global_head_dim": 16, + "num_experts": 4, + "top_k_experts": 1, + "max_position_embeddings": 64, + "sliding_window": 8, + "layer_types": ["sliding_attention", "full_attention"], + "use_bidirectional_attention": "vision", + }, + vision_config={ + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 16, + "max_position_embeddings": 64, + "patch_size": 4, + "pooling_kernel_size": 1, + "position_embedding_size": 16, + }, + canvas_length=4, + image_token_id=124, + boi_token_id=125, + eoi_token_id=126, + ) + + config = config_cls() _mutate_common_config_tree(config) model_type = case["model_type"] @@ -390,6 +432,21 @@ def _validate_defused_module(case: dict, module) -> None: "target_class_paths": ("transformers.models.dots1.modeling_dots1.Dots1NaiveMoe",), "validator": "experts", }, + { + "model_type": "diffusion_gemma", + "mode": "convert", + "model_module": "transformers.models.diffusion_gemma.modeling_diffusion_gemma", + "model_class": "DiffusionGemmaForBlockDiffusion", + "config_module": "transformers.models.diffusion_gemma.configuration_diffusion_gemma", + "config_class": "DiffusionGemmaConfig", + "target_class_paths": ( + "transformers.models.diffusion_gemma.modeling_diffusion_gemma.DiffusionGemmaTextExperts", + ), + "validator": "experts", + # There are separate experts blocks in the encoder and decoder; seeing + # both ensures the declarative class-path patch is applied everywhere. + "min_targets": 2, + }, { "model_type": "ernie4_5_moe", "mode": "convert",