add filter control

Author: Qubitium

README.md

@@ -33,9 +33,17 @@ from defuser import convert_model, replace_fused_blocks
 ```
 
 - `replace_fused_blocks(model_type)` patches supported HF model classes before `from_pretrained()` or direct model construction.
-- `convert_model(model, cleanup_original=True, max_layers=None)` converts an already loaded model in place. This is the runtime defusion path for supported post-load expert and MLP conversions, including `qwen3_5_moe` style checkpoints.
+- `convert_model(model, cleanup_original=True, max_layers=None, filter=None)` converts an already loaded model in place. This is the runtime defusion path for supported post-load expert and MLP conversions, including `qwen3_5_moe` style checkpoints.
 - Defuser is designed and CI-tested for `transformers>=5.3.0`, and support is only offered for that version range. Older versions log a warning on these public APIs and are skipped as unsupported.
 
+`filter` is an optional list of PCRE regex rules evaluated against full module paths such as `model.layers.0.mlp.experts`:
+
+- `+:regex` explicitly includes matching candidate module paths
+- `-:regex` explicitly excludes matching candidate module paths
+- `regex` is shorthand for `+:regex`
+- negative rules take priority over positive rules
+- when `filter` is provided, a candidate module is defused only if it matches at least one positive rule and no negative rules
+
 ## Supported Models
 
 Defuser currently supports the following `transformers==5.3.0` `model_type` values.
@@ -91,6 +99,20 @@ converted = convert_model(model)
 print(converted)  # True when runtime defusion happened
 ```
 
+Use `filter` when only specific blocks should be defused:
+
+```python
+from defuser import convert_model
+
+convert_model(
+    model,
+    filter=[
+        r"+:^model\.layers\.0\.mlp\.experts$",
+        r"-:^model\.layers\.0\.mlp\.experts\.shared_",
+    ],
+)
+```
+
 ## Real Qwen3.5 MoE Example
 
 The example below is written for the `transformers==5.3.0` public API surface and uses the real Hugging Face model `Qwen/Qwen3.5-35B-A3B-Instruct`. Defuser supports `transformers>=5.3.0`.

defuser/defuser.py

@@ -117,6 +117,7 @@ def convert_model(
         model: nn.Module,
         cleanup_original: bool = False,
         max_layers: int | None = None,
+        filter: list[str] | None = None,
 ) -> bool:
     """Convert one loaded model in place from fused experts to defused modules."""
     if warn_if_public_api_transformers_unsupported("convert_model()", logger):
@@ -200,7 +201,7 @@ def convert_model(
     if not check_model_compatibility(model):
         return False
 
-    apply_model_patches(model, max_layers=max_layers)
+    apply_model_patches(model, max_layers=max_layers, filter_rules=filter)
 
     # If fused blocks have already been structurally replaced at load model before,
     # there is no need to perform runtime defusing again
@@ -214,6 +215,7 @@ def convert_model(
         model,
         cleanup_original=cleanup_original,
         max_layers=max_layers,
+        filter_rules=filter,
     )
 
     return True

defuser/modeling/model_patches.py

@@ -14,7 +14,7 @@
     patch_parallel_experts,
     patch_split_gate_up_mlp,
 )
-from defuser.utils.common import is_within_max_layers
+from defuser.utils.common import compile_module_name_filter, is_within_max_layers, matches_module_name_filter
 import torch
 
 logger = LogBar(__name__)
@@ -87,7 +87,7 @@ def patched_init_weights(self, module):
 
 
 @register_model_patch("qwen3_omni_moe")
-def patch_qwen3_omni_text_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_qwen3_omni_text_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     """Restore text-only ``forward`` and ``generate`` behavior after class swapping."""
     model_cls = type(model)
     if not getattr(model_cls, "__module__", "").startswith("transformers.models.qwen3_omni_moe."):
@@ -122,11 +122,15 @@ def _patch_modules_by_class(
     patchers: dict[str, Callable],
     *,
     max_layers: int | None = None,
+    filter_rules=None,
 ) -> list[str]:
+    module_name_filter = compile_module_name_filter(filter_rules)
     applied = []
     for name, module in list(model.named_modules()):
         if not is_within_max_layers(name, max_layers):
             continue
+        if not matches_module_name_filter(name, module_name_filter):
+            continue
         class_path = f"{module.__class__.__module__}.{module.__class__.__name__}"
         patcher = patchers.get(class_path)
         if patcher is None:
@@ -141,6 +145,7 @@ def _patch_split_gate_up_mlps(
     patchers: dict[str, str],
     *,
     max_layers: int | None = None,
+    filter_rules=None,
 ) -> list[str]:
     return _patch_modules_by_class(
         model,
@@ -149,6 +154,7 @@ def _patch_split_gate_up_mlps(
             for class_path, variant in patchers.items()
         },
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
@@ -166,61 +172,67 @@ def _patch_split_gate_up_mlps(
 
 
 @register_model_patch("dia")
-def patch_dia_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_dia_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.dia.modeling_dia.DiaMLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.dia.modeling_dia.DiaMLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("glm")
-def patch_glm_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_glm_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.glm.modeling_glm.GlmMLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.glm.modeling_glm.GlmMLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("glm4")
-def patch_glm4_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_glm4_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.glm4.modeling_glm4.Glm4MLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.glm4.modeling_glm4.Glm4MLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("glm_image")
-def patch_glm_image_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_glm_image_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.glm_image.modeling_glm_image.GlmImageTextMLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.glm_image.modeling_glm_image.GlmImageTextMLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("glm_ocr")
-def patch_glm_ocr_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_glm_ocr_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.glm_ocr.modeling_glm_ocr.GlmOcrTextMLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.glm_ocr.modeling_glm_ocr.GlmOcrTextMLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("phi3")
-def patch_phi3_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_phi3_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.phi3.modeling_phi3.Phi3MLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.phi3.modeling_phi3.Phi3MLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("phi4_multimodal")
-def patch_phi4_multimodal_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_phi4_multimodal_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {
@@ -234,73 +246,86 @@ def patch_phi4_multimodal_runtime(model, max_layers: int | None = None) -> list[
                 ],
         },
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("zamba2")
-def patch_zamba2_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_zamba2_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_split_gate_up_mlps(
         model,
         {"transformers.models.zamba2.modeling_zamba2.Zamba2MLP": _STANDARD_SPLIT_GATE_UP_CLASSES["transformers.models.zamba2.modeling_zamba2.Zamba2MLP"]},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("dbrx")
-def patch_dbrx_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_dbrx_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_modules_by_class(
         model,
         {"transformers.models.dbrx.modeling_dbrx.DbrxExperts": patch_dbrx_experts},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
-def _patch_parallel_runtime(model, class_path: str, *, max_layers: int | None = None) -> list[str]:
-    return _patch_modules_by_class(model, {class_path: patch_parallel_experts}, max_layers=max_layers)
+def _patch_parallel_runtime(model, class_path: str, *, max_layers: int | None = None, filter_rules=None) -> list[str]:
+    return _patch_modules_by_class(
+        model,
+        {class_path: patch_parallel_experts},
+        max_layers=max_layers,
+        filter_rules=filter_rules,
+    )
 
 
 @register_model_patch("granitemoe")
-def patch_granitemoe_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_granitemoe_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_parallel_runtime(
         model,
         "transformers.models.granitemoe.modeling_granitemoe.GraniteMoeParallelExperts",
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("granitemoehybrid")
-def patch_granitemoehybrid_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_granitemoehybrid_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_parallel_runtime(
         model,
         "transformers.models.granitemoehybrid.modeling_granitemoehybrid.GraniteMoeHybridParallelExperts",
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("granitemoeshared")
-def patch_granitemoeshared_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_granitemoeshared_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_parallel_runtime(
         model,
         "transformers.models.granitemoeshared.modeling_granitemoeshared.GraniteMoeSharedParallelExperts",
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("jetmoe")
-def patch_jetmoe_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_jetmoe_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_parallel_runtime(
         model,
         "transformers.models.jetmoe.modeling_jetmoe.JetMoeParallelExperts",
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
 @register_model_patch("longcat_flash")
-def patch_longcat_flash_runtime(model, max_layers: int | None = None) -> list[str]:
+def patch_longcat_flash_runtime(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     return _patch_modules_by_class(
         model,
         {"transformers.models.longcat_flash.modeling_longcat_flash.LongcatFlashExperts": patch_longcat_flash_experts},
         max_layers=max_layers,
+        filter_rules=filter_rules,
     )
 
 
@@ -316,15 +341,15 @@ def apply_model_class_patches(model_type) -> list[str]:
     return applied
 
 
-def apply_model_patches(model, max_layers: int | None = None) -> list[str]:
+def apply_model_patches(model, max_layers: int | None = None, filter_rules=None) -> list[str]:
     """Run any registered runtime patch for the instantiated ``model``."""
     config = getattr(model, "config", None)
     model_type = getattr(config, "model_type", None)
     patch = _MODEL_PATCH_REGISTRY.get(model_type)
     if patch is None:
         return []
 
-    applied = patch(model, max_layers=max_layers)
+    applied = patch(model, max_layers=max_layers, filter_rules=filter_rules)
     if applied and DEBUG_ON:
         logger.debug(f"Applied model patches for model_type={model_type}: {', '.join(applied)}")
     return applied

defuser/modeling/moe_experts_interface.py

@@ -32,6 +32,7 @@
 from torch import nn
 
 from defuser.model_registry import MODEL_CONFIG, PATCH
+from defuser.utils.common import compile_module_name_filter, matches_module_name_filter
 from defuser.utils.device import clear_memory, to_meta
 
 from defuser import DEBUG_ON
@@ -693,7 +694,11 @@ def _unfuse_experts_weights_inplace(
     return True
 
 
-def prepare_model_for_moe_quantization(model: nn.Module, implementation: str = LINEAR_LOOP_IMPL) -> list[str]:
+def prepare_model_for_moe_quantization(
+    model: nn.Module,
+    implementation: str = LINEAR_LOOP_IMPL,
+    filter_rules=None,
+) -> list[str]:
     """Prepare a model for MOE quantization using transformers' experts interface.
 
     This function:
@@ -722,7 +727,10 @@ def prepare_model_for_moe_quantization(model: nn.Module, implementation: str = L
     unfused_modules = []
     decorated_unfused_modules = []
     experts_defuse_specs = _model_experts_defuse_specs(model)
+    module_name_filter = compile_module_name_filter(filter_rules)
     for name, module in model.named_modules():
+        if not matches_module_name_filter(name, module_name_filter):
+            continue
         spec = _matching_experts_defuse_spec(module, experts_defuse_specs)
         if spec is not None and _unfuse_experts_weights_inplace(
             module,