Defuser 0.0.4__tar.gz → 0.0.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Defuser
-Version: 0.0.4
+Version: 0.0.6
 Summary: Model defuser helper for HF Transformers.
 Author-email: ModelCloud <qubitium@modelcloud.ai>
 License-Expression: Apache-2.0

{defuser-0.0.4 → defuser-0.0.6}/Defuser.egg-info/SOURCES.txt

@@ -11,13 +11,9 @@ 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/modeling/moe_experts_interface.py
+defuser/modeling/replace_modules.py
+defuser/modeling/update_module.py
 defuser/utils/__init__.py
 defuser/utils/common.py
 defuser/utils/device.py

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Defuser
-Version: 0.0.4
+Version: 0.0.6
 Summary: Model defuser helper for HF Transformers.
 Author-email: ModelCloud <qubitium@modelcloud.ai>
 License-Expression: Apache-2.0

{defuser-0.0.4 → defuser-0.0.6}/defuser/__init__.py

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

defuser-0.0.6/defuser/defuser.py

@@ -0,0 +1,124 @@
+# 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.model_registry import MODEL_CONFIG
+from defuser.modeling.update_module import update_module
+from packaging import version
+import transformers
+from logbar import LogBar
+
+logger = LogBar(__name__)
+
+def check_model_compatibility(model: nn.Module) -> bool:
+    """Validate model type and transformers version compatibility."""
+    config = getattr(model, "config", None)
+    model_type = getattr(config, "model_type", None)
+    if model_type not in MODEL_CONFIG:
+        return False
+
+    min_ver = MODEL_CONFIG[model_type].get("min_transformers_version")
+    current_ver = version.parse(transformers.__version__)
+    if min_ver and current_ver < version.parse(min_ver):
+        logger.warn(
+            f"Skip conversion for model_type={model_type}: "
+            f"requires transformers>={min_ver}, current version is {transformers.__version__}."
+        )
+        return False
+
+    return True
+
+
+def convert_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.
+
+    # -----------------------------------------------------------------------
+    # 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.
+
+    if not check_model_compatibility(model):
+        return model
+
+    return update_module(
+        model,
+        cleanup_original=cleanup_original,
+        max_layers=max_layers,
+    )
+
+__all__ = ["convert_model"]

defuser-0.0.6/defuser/model_registry.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: 2026 ModelCloud.ai
+# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+MODEL_CONFIG = {
+    "qwen3_moe": {
+        "min_transformers_version": "5.0.0",
+    },
+    "qwen3_5_moe": {
+        "min_transformers_version": "5.2.0",
+    },
+    "qwen3_5_moe_text": {
+        "min_transformers_version": "5.2.0",
+    },
+}

{defuser-0.0.4/defuser/modeling/fused_moe → defuser-0.0.6/defuser/modeling}/moe_experts_interface.py

@@ -56,14 +56,6 @@ KNOWN_PROJECTION_PATTERNS = {
     "gate_proj": {"is_input_proj": True, "output_multiplier": 1},  # hidden -> intermediate (gate)
     "up_proj": {"is_input_proj": True, "output_multiplier": 1},  # hidden -> intermediate (up)
     "down_proj": {"is_input_proj": False, "output_multiplier": 1},  # intermediate -> hidden
-    # Mixtral-style
-    "w1": {"is_input_proj": True, "output_multiplier": 1},  # gate: hidden -> intermediate
-    "w2": {"is_input_proj": False, "output_multiplier": 1},  # down: intermediate -> hidden
-    "w3": {"is_input_proj": True, "output_multiplier": 1},  # up: hidden -> intermediate
-    # DBRX-style
-    "v1": {"is_input_proj": True, "output_multiplier": 1},
-    "w1_proj": {"is_input_proj": True, "output_multiplier": 1},
-    "w2_proj": {"is_input_proj": False, "output_multiplier": 1},
 }
 
 

{defuser-0.0.4/defuser/modeling/fused_moe → defuser-0.0.6/defuser/modeling}/replace_modules.py

@@ -16,34 +16,11 @@ import torch
 from logbar import LogBar
 from tqdm import tqdm
 
-from defuser.model_registry import MODEL_CONFIG, PATCH
-from defuser.utils.common import is_within_max_layers
+from defuser.utils.common import is_within_max_layers, is_transformers_version_greater_or_equal_5
 
 logger = LogBar(__name__)
 
 
-def is_model_patchable(model: torch.nn.Module) -> bool:
-    """Check if the model has a custom replacement registered via MODEL_CONFIG.
-
-    Returns True if the model's model_type matches a key in MODEL_CONFIG.
-    """
-    if hasattr(model, "config") and hasattr(model.config, "model_type"):
-        return model.config.model_type in MODEL_CONFIG
-    return False
-
-
-def _import_required_replacements(model: torch.nn.Module) -> None:
-    """Import replacement modules required for the model's defuse workflow."""
-    if not is_model_patchable(model):
-        return
-    model_type = model.config.model_type
-    module_path = MODEL_CONFIG[model_type].get(PATCH.DEFUSE)
-    if not module_path:
-        return
-    importlib.import_module(module_path)
-    logger.debug(f"Loaded replacement module for {model_type}: {module_path}")
-
-
 def materialize_model(model: torch.nn.Module) -> None:
     def _materialize_module(module: torch.nn.Module) -> None:
         if isinstance(module, ReplacementModuleBase):
@@ -216,7 +193,7 @@ def _handle_moe_modules(model: torch.nn.Module) -> list[str]:
     Returns:
         List of module names that were processed
     """
-    from defuser.modeling.fused_moe.moe_experts_interface import (
+    from defuser.modeling.moe_experts_interface import (
         is_linear_loop_available,
         prepare_model_for_moe_quantization,
     )
@@ -259,15 +236,10 @@ def apply_replacements(
     Returns:
         The model with modules replaced.
     """
-    _import_required_replacements(model)
-
     _log_first_moe_block(model, "before replacement")
 
-    # Custom replacements first
-    if is_model_patchable(model):
-        _apply_custom_replacements(model, max_layers=max_layers)
-    # if auto_detect_moe and is_transformers_version_greater_or_equal_5():
-    #     _handle_moe_modules(model)
+    if auto_detect_moe and is_transformers_version_greater_or_equal_5():
+        _handle_moe_modules(model)
 
     _log_first_moe_block(model, "after replacement")
 

{defuser-0.0.4/defuser/modeling/fused_moe → defuser-0.0.6/defuser/modeling}/update_module.py

@@ -6,7 +6,7 @@
 # Adapted from intel/auto-round
 # at https://github.com/intel/auto-round/blob/main/auto_round/special_model_handler.py
 
-from defuser.modeling.fused_moe.replace_modules import apply_replacements, release_original_module_
+from defuser.modeling.replace_modules import apply_replacements, release_original_module_
 
 
 def update_module(

{defuser-0.0.4 → defuser-0.0.6}/defuser/utils/hf.py

@@ -16,8 +16,7 @@ from logbar import LogBar
 from packaging import version
 from transformers import AutoConfig
 
-from defuser.model_registry import MODEL_CONFIG, PATCH
-from defuser.utils.common import is_within_max_layers
+from defuser.model_registry import MODEL_CONFIG
 
 logger = LogBar(__name__)
 
@@ -105,34 +104,3 @@ def pre_check_config(model_name: str | torch.nn.Module):
     except:
         return True
     return True
-
-
-def patch(model: torch.nn.Module, max_layers: int | None = None) -> bool:
-    res = pre_check_config(model)
-    if not res:
-        return False
-    model_type = getattr(model.config, "model_type")
-    cfg = MODEL_CONFIG[model_type]
-    # patch blocks
-    for orig_path, custom_path in cfg.get(PATCH.REPLACE_MODULE, []):
-        orig_module_path, orig_class_name = orig_path.rsplit(".", 1)
-        custom_module_path, custom_class_name = custom_path.rsplit(".", 1)
-        try:
-            orig_module = importlib.import_module(orig_module_path)
-            custom_module = importlib.import_module(custom_module_path)
-            custom_class = getattr(custom_module, custom_class_name)
-            orig_class = getattr(orig_module, orig_class_name)
-            names = []
-            for n, m in model.named_modules():
-                if isinstance(m, orig_class):
-                    if not is_within_max_layers(n, max_layers):
-                        continue
-                    names.append((n, next(m.parameters()).dtype))
-            for (n, orig_dtype) in names:
-                model.set_submodule(n, custom_class(model.config).to(orig_dtype), True)
-            logger.info(f"Patched module: {orig_path} -> {custom_path}")
-            return True
-        except Exception as e:
-            logger.warn(f"Failed to patch {orig_path}: {e}")
-            return False
-    return False

{defuser-0.0.4 → defuser-0.0.6}/pyproject.toml

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

{defuser-0.0.4 → defuser-0.0.6}/tests/test_convert_model.py

@@ -5,30 +5,34 @@
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText
 
-from defuser import convert_hf_model
-from defuser.modeling.fused_moe.replace_modules import materialize_model
+from defuser import convert_model
+from defuser.modeling.replace_modules import materialize_model
 
 
 def test_qwen3_moe():
-    from defuser.modeling.unfused_moe.qwen3_moe import LinearQwen3MoeSparseMoeBlock
-
-    config = AutoConfig.from_pretrained("/monster/data/model/Qwen3-30B-A3B")
+    model_id = "Qwen/Qwen3-30B-A3B"
+    config = AutoConfig.from_pretrained(model_id)
     config.num_hidden_layers = 1
     model = AutoModelForCausalLM.from_pretrained(
-        "/monster/data/model/Qwen3-30B-A3B",
+        model_id,
         config=config,
         ignore_mismatched_sizes=True,
     )
 
     assert model.config.model_type == "qwen3_moe"
 
-    converted = convert_hf_model(model, max_layers=1)
+    converted = convert_model(model, max_layers=1)
     assert converted
-    assert isinstance(model.model.layers[0].mlp, LinearQwen3MoeSparseMoeBlock)
+
+    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")
 
 
 def test_qwen3_5_moe():
-    from defuser.modeling.fused_moe.qwen3_5_moe import LinearQwen3_5MoeSparseMoeBlock
     from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeSparseMoeBlock
 
     config = AutoConfig.from_pretrained("/monster/data/model/Qwen3.5-35B-A3B")
@@ -50,14 +54,20 @@ def test_qwen3_5_moe():
     expected_up = original_moe_block.experts.gate_up_proj[0, intermediate_dim:, :hidden_dim].contiguous().clone()
     expected_down = original_moe_block.experts.down_proj[0, :hidden_dim, :intermediate_dim].contiguous().clone()
 
-    converted = convert_hf_model(model, cleanup_original=False, max_layers=1)
+    converted = convert_model(model, cleanup_original=False, max_layers=1)
     assert converted
 
     moe_block = model.model.language_model.layers[0].mlp
-    assert isinstance(moe_block, LinearQwen3_5MoeSparseMoeBlock)
+    experts = moe_block.experts
+
+    assert hasattr(experts, "0")
+    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])
 
-    torch.testing.assert_close(moe_block.experts[0].gate_proj.weight, expected_gate)
-    torch.testing.assert_close(moe_block.experts[0].up_proj.weight, expected_up)
-    torch.testing.assert_close(moe_block.experts[0].down_proj.weight, expected_down)
+    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)

{defuser-0.0.4 → defuser-0.0.6}/tests/test_replace_modules_tracker.py

@@ -9,7 +9,7 @@ import weakref
 import pytest
 import torch
 
-from defuser.modeling.fused_moe.replace_modules import (
+from defuser.modeling.replace_modules import (
     ModuleReplacementTracker,
     ReplacementModuleBase,
     release_original_module_,

defuser-0.0.4/defuser/defuser.py

@@ -1,100 +0,0 @@
-# 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

@@ -1,35 +0,0 @@
-# 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",
-    },
-}

defuser-0.0.4/defuser/modeling/fused_moe/qwen3_5_moe.py

@@ -1,116 +0,0 @@
-# SPDX-FileCopyrightText: 2026 ModelCloud.ai
-# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
-# SPDX-License-Identifier: Apache-2.0
-# Contact: qubitium@modelcloud.ai, x.com/qubitium
-
-# Adapted from intel/auto-round
-# at https://github.com/intel/auto-round/blob/main/auto_round/modeling/fused_moe/qwen3_5_moe.py
-
-import torch
-from torch.nn import functional as F
-from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeMLP
-from transformers.utils.versions import require_version
-
-from defuser.modeling.fused_moe.replace_modules import ReplacementModuleBase
-from defuser.utils.device import clear_memory, to_meta, unsupported_meta_device
-
-require_version("transformers>=5.2.0")
-
-from defuser.utils.model import _update_parameter
-
-
-class LinearQwen3_5MoeSparseMoeBlock(ReplacementModuleBase):
-    def __init__(self, original, config):
-        super().__init__(original)
-        self.gate = original.gate
-        text_config = config.get_text_config()
-        self.shared_expert = original.shared_expert
-        self.experts = SequentialQwen3_5MoeExperts(text_config, original.experts)
-        self.shared_expert_gate = original.shared_expert_gate
-        self.num_experts = text_config.num_experts
-
-    @classmethod
-    def original_module_class(cls) -> str:
-        """Return the class name of the module this replaces."""
-        return "Qwen3_5MoeSparseMoeBlock"
-
-    def _materialize_weights(self) -> None:
-        original = self._get_original_module()
-        self.experts._materialize_weights(original.experts)
-        clear_memory()
-
-    def experts_forward(
-            self,
-            hidden_states: torch.Tensor,
-            top_k_index: torch.Tensor,
-            top_k_weights: torch.Tensor,
-    ) -> torch.Tensor:
-        final_hidden_states = torch.zeros_like(hidden_states)
-        with torch.no_grad():
-            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
-            expert_mask = expert_mask.permute(2, 1, 0)
-            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
-
-        for expert_idx in expert_hit:
-            expert_idx = expert_idx[0]
-            if expert_idx == self.num_experts:
-                continue
-            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
-            current_state = hidden_states[token_idx]
-            # gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
-            # current_hidden_states = self.act_fn(gate) * up
-            # current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
-            current_hidden_states = self.experts[expert_idx](current_state)
-            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
-            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
-
-        return final_hidden_states
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        batch_size, sequence_length, hidden_dim = hidden_states.shape
-        hidden_states_reshaped = hidden_states.view(-1, hidden_dim)
-        shared_expert_output = self.shared_expert(hidden_states_reshaped)
-        _, routing_weights, selected_experts = self.gate(hidden_states_reshaped)
-        expert_output = self.experts_forward(hidden_states_reshaped, selected_experts, routing_weights)
-
-        shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output
-
-        expert_output += shared_expert_output
-        expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim)
-        return expert_output
-
-    @classmethod
-    def from_original(
-            cls,
-            original,
-            config,
-            **kwargs,
-    ):
-        """Create an instance from the original module."""
-        return cls(original, config)
-
-
-class SequentialQwen3_5MoeExperts(torch.nn.ModuleList):
-    def __init__(self, config, original):
-        super().__init__()
-        self.num_experts = original.gate_up_proj.shape[0]
-        intermediate_size = config.moe_intermediate_size
-
-        super().__init__([Qwen3_5MoeMLP(config, intermediate_size) for _ in range(self.num_experts)])
-
-    def _materialize_weights(self, original) -> None:
-        intermediate_size = original.down_proj.shape[-1]
-        if not unsupported_meta_device(original):
-            for i in range(self.num_experts):
-                gate_up = original.gate_up_proj[i]
-                down = original.down_proj[i]
-
-                gate_proj = gate_up[:intermediate_size, :]
-                up_proj = gate_up[intermediate_size:, :]
-
-                _update_parameter(self[i].gate_proj, "weight", gate_proj.contiguous())
-                _update_parameter(self[i].up_proj, "weight", up_proj.contiguous())
-                _update_parameter(self[i].down_proj, "weight", down.contiguous())
-            del gate_up, down, gate_proj, up_proj
-            to_meta(original)  # release original experts parameters
-            clear_memory()

defuser-0.0.4/defuser/modeling/unfused_moe/qwen3_moe.py

@@ -1,70 +0,0 @@
-# SPDX-FileCopyrightText: 2026 ModelCloud.ai
-# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
-# SPDX-License-Identifier: Apache-2.0
-# Contact: qubitium@modelcloud.ai, x.com/qubitium
-
-# Adapted from intel/auto-round
-# at https://github.com/intel/auto-round/blob/main/auto_round/modeling/unfused_moe/qwen3_moe.py
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-
-
-class LinearQwen3MoeSparseMoeBlock(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeMLP
-
-        self.num_experts = config.num_experts
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-
-        # This must be linear for vllm alignment
-        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
-        self.experts = nn.ModuleList(
-            [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
-        )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """ """
-        batch_size, sequence_length, hidden_dim = hidden_states.shape
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        # router_logits: (batch * sequence_length, n_experts)
-        router_logits = self.gate(hidden_states)
-
-        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
-        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
-        if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
-            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
-        # we cast back to the input dtype
-        routing_weights = routing_weights.to(hidden_states.dtype)
-
-        final_hidden_states = torch.zeros(
-            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
-        )
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be solicited
-        with torch.no_grad():
-            expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
-
-            # Loop over all available experts in the model and perform the computation on each expert
-            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
-        for expert_idx in expert_hit:
-            if expert_idx == self.num_experts:
-                continue
-            expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
-            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
-        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
-        return final_hidden_states