added LinearQwen3_5MoeSparseMoeBlock

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

Author: ZX-ModelCloud

defuser/modeling/fused_moe/__init__.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: 2026 ModelCloud.ai
+# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+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, 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
+        with torch.device("meta"):
+            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
+
+        with torch.device("meta"):
+            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
+            original.to_empty(device="meta")  # release original experts parameters
+            clear_memory()

defuser/modeling/fused_moe/qwen3_moe.py

@@ -0,0 +1,209 @@
+# 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 abc import ABC, abstractmethod
+
+import torch
+from typing import Dict, Type
+from defuser.logger import logger
+from dataclasses import dataclass
+
+
+class ReplacementModuleBase(ABC, torch.nn.Module):
+    """
+    Abstract base class for module replacement during calibration phase.
+
+    Replacement modules replace original modules to ensure all components
+    receive data for proper quantization statistics.
+
+    Subclasses must:
+    1. Implement `original_module_class()` to return the target module class name
+    2. Implement `__init__()` with signature:
+       (self, original, config)
+    """
+
+    # Registry: module class name -> replacement module class
+    _replacement_registry: Dict[str, Type["ReplacementModuleBase"]] = {}
+
+    def __init_subclass__(cls, **kwargs):
+        """Automatically register subclasses in the replacement registry."""
+        super().__init_subclass__(**kwargs)
+
+        # Only register if it's a concrete implementation (not ABC)
+        if not getattr(cls, "__abstractmethods__", None):
+            if cls.original_module_class() is None:
+                raise TypeError(
+                    f"{cls.__name__} must implement 'original_module_class()' class method "
+                    "to return the name of the module class it replaces"
+                )
+
+            if cls.original_module_class() in cls._replacement_registry:
+                existing = cls._replacement_registry[cls.original_module_class()]
+                raise ValueError(
+                    f"Module '{cls.original_module_class()}' already registered to "
+                    f"{existing.__name__}. Cannot register {cls.__name__}."
+                )
+
+            cls._replacement_registry[cls.original_module_class()] = cls
+            logger.trace(f"Registered {cls.__name__} for replacing {cls.original_module_class()}")
+
+    def __init__(self, original: torch.nn.Module):
+        super().__init__()
+        _global_tracker.register_replacement(
+            name=str(id(self)),
+            original=original,
+            replacement=self,
+        )
+        self._materialized = False
+
+    @classmethod
+    def get_replacement_class(cls, module_class_name: str) -> Type["ReplacementModuleBase"]:
+        """Get replacement class for a given module class name."""
+        return cls._replacement_registry.get(module_class_name)
+
+    @classmethod
+    def is_registered(cls, module_class_name: str) -> bool:
+        """Check if a module class has a replacement implementation."""
+        return module_class_name in cls._replacement_registry
+
+    @classmethod
+    def is_to_be_replaced(
+            cls,
+            original: torch.nn.Module,
+    ) -> bool:
+        """Determine if the given module should be replaced.
+
+        Users can extend this method to add custom logic for replacement.
+        """
+        return cls.is_registered(original.__class__.__name__)
+
+    @classmethod
+    def get_registered_modules(cls) -> list:
+        """Get list of all registered module class names."""
+        return list(cls._replacement_registry.keys())
+
+    @classmethod
+    @abstractmethod
+    def original_module_class(cls) -> str:
+        """Return the class name of the module this replaces."""
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_original(
+            cls,
+            original: torch.nn.Module,
+            config,
+    ) -> "ReplacementModuleBase":
+        """Create replacement module from original module."""
+        pass
+
+    def materialize_weights(self):
+        """Materialize weights if needed."""
+        if not self._materialized:
+            self._materialize_weights()
+            self.post_process_materialization()
+
+    def _materialize_weights(self) -> None:
+        """Materialize weights from the original module.
+
+        Subclasses should override this method to implement
+        weight materialization logic.
+        """
+        pass
+
+    def release_original_module(self) -> None:
+        """Release reference to the original module to free memory."""
+        # Release from global tracker
+        _global_tracker.release_original(self)
+
+    def _get_original_module(self) -> torch.nn.Module:
+        """Get the original module associated with this replacement."""
+        return _global_tracker.get_original(self)
+
+    def post_process_materialization(self) -> None:
+        """Mark the replacement module as materialized."""
+        self._materialized = True
+        self.release_original_module()
+
+
+@dataclass
+class ReplacedModuleInfo:
+    original_module: torch.nn.Module
+    replacement_module: ReplacementModuleBase
+
+
+class ModuleReplacementTracker:
+    """Tracker to maintain mapping between replacement modules and their original modules.
+
+    This is a singleton class - only one instance can exist.
+    """
+
+    _instance = None
+    _initialized = False
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(ModuleReplacementTracker, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        # Only initialize once
+        if ModuleReplacementTracker._initialized:
+            return
+
+        # Map from replacement module id to original module
+        self._replacement_to_original: Dict[int, torch.nn.Module] = {}
+        # Map from module name to ReplacedModuleInfo
+        self._name_to_info: Dict[str, ReplacedModuleInfo] = {}
+
+        ModuleReplacementTracker._initialized = True
+
+    @classmethod
+    def get_instance(cls) -> "ModuleReplacementTracker":
+        """Get the singleton instance of the tracker."""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def register_replacement(self, name: str, original: torch.nn.Module, replacement: ReplacementModuleBase) -> None:
+        """Register a module replacement."""
+        self._replacement_to_original[id(replacement)] = original
+        self._name_to_info[name] = ReplacedModuleInfo(original_module=original, replacement_module=replacement)
+        logger.trace(f"Registered replacement for module: {name}")
+
+    def get_original(self, replacement: ReplacementModuleBase) -> torch.nn.Module:
+        """Get the original module for a given replacement module."""
+        return self._replacement_to_original.get(id(replacement))
+
+    def get_info_by_name(self, name: str) -> ReplacedModuleInfo:
+        """Get replacement info by module name."""
+        return self._name_to_info.get(name)
+
+    def release_original(self, replacement: ReplacementModuleBase) -> None:
+        """Release the original module associated with a replacement module."""
+        replacement_id = id(replacement)
+        if replacement_id in self._replacement_to_original:
+            original = self._replacement_to_original[replacement_id]
+            # Delete the original module to free memory
+            del original
+            del self._replacement_to_original[replacement_id]
+            logger.trace(f"Released original module for replacement {replacement_id}")
+
+    def release_all_originals(self) -> None:
+        """Release all tracked original modules."""
+        count = len(self._replacement_to_original)
+        if count > 0:
+            self._replacement_to_original.clear()
+            logger.debug(f"Released {count} original modules from tracker")
+
+    def clear(self) -> None:
+        """Clear all tracked information."""
+        self._replacement_to_original.clear()
+        self._name_to_info.clear()
+        logger.debug("Cleared module replacement tracker")
+
+
+_global_tracker = ModuleReplacementTracker()