Defuser 0.0.3__tar.gz → 0.0.4__tar.gz

{defuser-0.0.3 → defuser-0.0.4}/Defuser.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Defuser
-Version: 0.0.3
+Version: 0.0.4
 Summary: Model defuser helper for HF Transformers.
 Author-email: ModelCloud <qubitium@modelcloud.ai>
 License-Expression: Apache-2.0
@@ -22,7 +22,6 @@ License-File: LICENSE
 Requires-Dist: transformers
 Dynamic: license-file
 
-
 <div align=center>
 <img width="50%" alt="image" src="https://github.com/user-attachments/assets/f801617b-8959-474a-a565-6b8897e2fcbf" />
 <h1 align="center">Defuser</h1>
@@ -37,5 +36,8 @@ Dynamic: license-file
 </p>
 Model defuser helper for HF Transformers >= 5.0. In HF Transformers 5.x releases, many MoE modules became auto-stacked or auto-fused by new modeling code which has benefits but also downsides. 
 
-* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating in a simple naive state. 
-* There are cases, quantization libraries, where we need to run inference where module input/output needs to be individually captured and this pkg can help complete this task.  
+* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time
+  weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating
+  in a simple naive state.
+* There are cases, quantization libraries, where we need to run inference where module input/output needs to be
+  individually captured and this pkg can help complete this task.  

defuser-0.0.4/Defuser.egg-info/SOURCES.txt

@@ -0,0 +1,28 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+Defuser.egg-info/PKG-INFO
+Defuser.egg-info/SOURCES.txt
+Defuser.egg-info/dependency_links.txt
+Defuser.egg-info/requires.txt
+Defuser.egg-info/top_level.txt
+defuser/__init__.py
+defuser/defuser.py
+defuser/model_registry.py
+defuser/modeling/__init__.py
+defuser/modeling/fused_moe/__init__.py
+defuser/modeling/fused_moe/moe_experts_interface.py
+defuser/modeling/fused_moe/qwen3_5_moe.py
+defuser/modeling/fused_moe/replace_modules.py
+defuser/modeling/fused_moe/update_module.py
+defuser/modeling/unfused_moe/__init__.py
+defuser/modeling/unfused_moe/qwen3_moe.py
+defuser/utils/__init__.py
+defuser/utils/common.py
+defuser/utils/device.py
+defuser/utils/hf.py
+defuser/utils/model.py
+tests/test_convert_model.py
+tests/test_device_utils.py
+tests/test_replace_modules_tracker.py

{defuser-0.0.3 → defuser-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Defuser
-Version: 0.0.3
+Version: 0.0.4
 Summary: Model defuser helper for HF Transformers.
 Author-email: ModelCloud <qubitium@modelcloud.ai>
 License-Expression: Apache-2.0
@@ -22,7 +22,6 @@ License-File: LICENSE
 Requires-Dist: transformers
 Dynamic: license-file
 
-
 <div align=center>
 <img width="50%" alt="image" src="https://github.com/user-attachments/assets/f801617b-8959-474a-a565-6b8897e2fcbf" />
 <h1 align="center">Defuser</h1>
@@ -37,5 +36,8 @@ Dynamic: license-file
 </p>
 Model defuser helper for HF Transformers >= 5.0. In HF Transformers 5.x releases, many MoE modules became auto-stacked or auto-fused by new modeling code which has benefits but also downsides. 
 
-* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating in a simple naive state. 
-* There are cases, quantization libraries, where we need to run inference where module input/output needs to be individually captured and this pkg can help complete this task.  
+* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time
+  weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating
+  in a simple naive state.
+* There are cases, quantization libraries, where we need to run inference where module input/output needs to be
+  individually captured and this pkg can help complete this task.  

{defuser-0.0.3 → defuser-0.0.4}/README.md

@@ -1,4 +1,3 @@
-
 <div align=center>
 <img width="50%" alt="image" src="https://github.com/user-attachments/assets/f801617b-8959-474a-a565-6b8897e2fcbf" />
 <h1 align="center">Defuser</h1>
@@ -13,5 +12,8 @@
 </p>
 Model defuser helper for HF Transformers >= 5.0. In HF Transformers 5.x releases, many MoE modules became auto-stacked or auto-fused by new modeling code which has benefits but also downsides. 
 
-* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating in a simple naive state. 
-* There are cases, quantization libraries, where we need to run inference where module input/output needs to be individually captured and this pkg can help complete this task.  
+* Goal is to provide naive module/layer forwarding code for all models supported by HF transformers where run-time
+  weight and structure level optimizations such weight merging, stacking, fusing are reversed so the model is operating
+  in a simple naive state.
+* There are cases, quantization libraries, where we need to run inference where module input/output needs to be
+  individually captured and this pkg can help complete this task.  

defuser-0.0.4/defuser/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: 2026 ModelCloud.ai
+# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+def convert_hf_model(*args, **kwargs):
+    """Lazily import conversion entrypoint to avoid import-time cycles."""
+    from .defuser import convert_hf_model as _convert_hf_model
+
+    return _convert_hf_model(*args, **kwargs)
+
+
+__all__ = ["convert_hf_model"]

defuser-0.0.4/defuser/defuser.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: 2026 ModelCloud.ai
+# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+from torch import nn
+
+from defuser.modeling.fused_moe.update_module import update_module
+from defuser.utils.hf import patch
+
+
+def convert_hf_model(
+    model: nn.Module,
+    cleanup_original: bool = False,
+    max_layers: int | None = None,
+) -> nn.Module:
+    if max_layers is not None and max_layers < 1:
+        raise ValueError("max_layers must be >= 1 when provided")
+
+    # Patch modeling structure for legacy Qwen3 MoE
+    #
+    # There are two slightlyfis_within_max_layers different checkpoint formats we need to support:
+    #   1) Qwen3 MoE
+    #   2) Qwen3.5 MoE
+    #
+    # The key difference is how the expert MLP weights are stored in the original
+    # checkpoint (fused vs. defused). Because of that, the amount of work needed
+    # after replacing the modeling structure is different.
+    #
+    # ---------------------------------------------------------------------------
+    # Step 1: Try applying a lightweight modeling patch
+    # ---------------------------------------------------------------------------
+    # `apply_modeling_patch(model)` only replaces the *modeling structure*
+    # (module definitions / forward logic) to match our runtime implementation.
+    #
+    # For **Qwen3 MoE**, this is sufficient because:
+    #   - The original checkpoint already stores `mlp.experts` weights in a
+    #     **defused format**.
+    #   - In other words, the tensors are already separated as:
+    #
+    #       gate_proj
+    #       up_proj
+    #       down_proj
+    #
+    #   - Therefore we only need to swap the modeling implementation so that the
+    #     module structure matches the expected layout, without touching the
+    #     underlying tensors.
+    #
+    # If this patch succeeds, it means the model is in the Qwen3 MoE format and
+    # no further tensor transformation is required.
+    is_applied = patch(model, max_layers=max_layers)
+    if not is_applied:
+        # -----------------------------------------------------------------------
+        # Step 2: Handle Qwen3.5 MoE checkpoints
+        # -----------------------------------------------------------------------
+        #
+        # If `apply_modeling_patch` fails, we assume the checkpoint corresponds to
+        # **Qwen3.5 MoE**.
+        #
+        # In Qwen3.5 MoE, the expert MLP weights are stored in a **fused format**.
+        # Specifically, the checkpoint keeps tensors such as:
+        #
+        #     gate_up_proj
+        #     down_proj
+        #
+        # where `gate_proj` and `up_proj` are fused together.
+        #
+        # Because our runtime modeling expects **defused tensors**, simply replacing
+        # the module structure is not enough. We must also convert the stored
+        # parameters.
+        #
+        # `update_module()` performs two tasks:
+        #
+        #   1) Replace the modeling structure so that it matches the expected
+        #      defused MoE implementation.
+        #
+        #   2) Prepare the module for **tensor defusion** of the expert weights.
+        #
+        # After the structure update, `materialize_model_()` will be invoked to
+        # actually split the fused tensors:
+        #
+        #     gate_up_proj  -->  gate_proj + up_proj
+        #
+        # and ensure the module finally contains the expected parameters:
+        #
+        #     gate_proj
+        #     up_proj
+        #     down_proj
+        #
+        # This ensures compatibility between the Qwen3.5 fused checkpoint format
+        # and the runtime model implementation that operates on defused weights.
+        model = update_module(
+            model,
+            cleanup_original=cleanup_original,
+            max_layers=max_layers,
+        )
+    return model
+
+
+__all__ = ["convert_hf_model"]

defuser-0.0.4/defuser/model_registry.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2026 ModelCloud.ai
+# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+from enum import Enum
+
+
+class PATCH(str, Enum):
+    DEFUSE = "defuse"
+    REPLACE_MODULE = "replace_module"
+
+
+MODEL_CONFIG = {
+    "qwen3_moe": {
+        "min_transformers_version": "5.0.0",
+        # structure path only replaces modeling structure
+        PATCH.REPLACE_MODULE: [
+            (
+                "transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock",
+                "defuser.modeling.unfused_moe.qwen3_moe.LinearQwen3MoeSparseMoeBlock",
+            )
+        ],
+    },
+    "qwen3_5_moe": {
+        "min_transformers_version": "5.2.0",
+        # Replacement module path imported only when the defuse workflow runs
+        PATCH.DEFUSE: "defuser.modeling.fused_moe.qwen3_5_moe",
+    },
+    "qwen3_5_moe_text": {
+        "min_transformers_version": "5.2.0",
+        # Replacement module path imported only when the defuse workflow runs
+        PATCH.DEFUSE: "defuser.modeling.fused_moe.qwen3_5_moe",
+    },
+}