Upload modeling_hunyuan.py

modeling_hunyuan.py

@@ -1605,68 +1605,15 @@ class HunYuanMoEV1ForCausalLM(HunYuanPreTrainedModel):
             )
         return reordered_past
 
-    # -----------------------------------------------------------------------
-    # Helper for module replacement
-    # -----------------------------------------------------------------------
-    def _replace_submodule(root: nn.Module, target: str, new_module: nn.Module):
-        """Replace a (possibly nested) sub‑module.
-
-        ``target`` is the dotted path returned by ``model.named_modules()``.
-        """
-        parts = target.split('.')
-        parent = root
-        for p in parts[:-1]:
-            parent = getattr(parent, p)
-        setattr(parent, parts[-1], new_module)
-
-    # -----------------------------------------------------------------------
-    # Public APIs
-    # -----------------------------------------------------------------------
-    def densify(model: nn.Module):
-        """Convert all :class:`HunYuanMoE` modules under *model* to
-        :class:`HunYuanDenseMoE`. Operates **in‑place**."""
-        replacements = []
-        for name, module in model.named_modules():
-            if isinstance(module, HunYuanMoE):
-                replacements.append((name, module))
-        for name, sparse_moe in replacements:
-            dense_moe = HunYuanDenseMoE(sparse_moe).to(next(sparse_moe.parameters()).device)
-            _replace_submodule(model, name, dense_moe)
-        return model
-
-
-    def sparsify(model: nn.Module):
-        """Rebuild standard sparse :class:`HunYuanMoE` modules from their
-        fused :class:`HunYuanDenseMoE` form. Operates **in‑place**."""
-        replacements = []
-        for name, module in model.named_modules():
-            if isinstance(module, HunYuanDenseMoE):
-                replacements.append((name, module))
-        for name, dense_moe in replacements:
-            cfg = dense_moe.config
-            sparse_moe = HunYuanMoE(cfg, layer_idx=dense_moe.layer_idx).to(next(dense_moe.parameters()).device)
-
-            # Copy router
-            sparse_moe.gate.load_state_dict(dense_moe.gate.state_dict())
-
-            # Slice fused weights back to per‑expert
-            for idx, expert in enumerate(sparse_moe.experts):
-                start = idx * dense_moe.intermediate_size
-                end   = (idx + 1) * dense_moe.intermediate_size
-
-                expert.gate_proj.weight.data.copy_(
-                    dense_moe.fused_gate_proj.weight.data[start:end]
-                )
-                expert.up_proj.weight.data.copy_(
-                    dense_moe.fused_up_proj.weight.data[start:end]
-                )
-                expert.down_proj.weight.data.copy_(
-                    dense_moe.fused_down_proj.weight.data[:, start:end]
-                )
-
-            _replace_submodule(model, name, sparse_moe)
-        return model
+    def densify(self):
+        """In-place fusion of every HunYuanMoE block under this model."""
+        densify(self)          # just call the standalone helper
+        return self            # for chaining
 
+    def sparsify(self):
+        """Restore the original sparse experts (inverse of densify)."""
+        sparsify(self)
+        return self
 
 @add_start_docstrings(
     """
@@ -1907,4 +1854,66 @@ class HunYuanDenseMoE(nn.Module):
         sparse_out = self._sparse_path(x, probs)
 
         out = dense_out + (sparse_out - dense_out).detach()  # STE
-        return out.view(bsz, seq_len, self.hidden_size)
+        return out.view(bsz, seq_len, self.hidden_size)
+
+# -----------------------------------------------------------------------
+# Helper for module replacement
+# -----------------------------------------------------------------------
+def _replace_submodule(root: nn.Module, target: str, new_module: nn.Module):
+    """Replace a (possibly nested) sub‑module.
+
+    ``target`` is the dotted path returned by ``model.named_modules()``.
+    """
+    parts = target.split('.')
+    parent = root
+    for p in parts[:-1]:
+        parent = getattr(parent, p)
+    setattr(parent, parts[-1], new_module)
+
+# -----------------------------------------------------------------------
+# Public APIs
+# -----------------------------------------------------------------------
+def densify(model: nn.Module):
+    """Convert all :class:`HunYuanMoE` modules under *model* to
+    :class:`HunYuanDenseMoE`. Operates **in‑place**."""
+    replacements = []
+    for name, module in model.named_modules():
+        if isinstance(module, HunYuanMoE):
+            replacements.append((name, module))
+    for name, sparse_moe in replacements:
+        dense_moe = HunYuanDenseMoE(sparse_moe).to(next(sparse_moe.parameters()).device)
+        _replace_submodule(model, name, dense_moe)
+    return model
+
+
+def sparsify(model: nn.Module):
+    """Rebuild standard sparse :class:`HunYuanMoE` modules from their
+    fused :class:`HunYuanDenseMoE` form. Operates **in‑place**."""
+    replacements = []
+    for name, module in model.named_modules():
+        if isinstance(module, HunYuanDenseMoE):
+            replacements.append((name, module))
+    for name, dense_moe in replacements:
+        cfg = dense_moe.config
+        sparse_moe = HunYuanMoE(cfg, layer_idx=dense_moe.layer_idx).to(next(dense_moe.parameters()).device)
+
+        # Copy router
+        sparse_moe.gate.load_state_dict(dense_moe.gate.state_dict())
+
+        # Slice fused weights back to per‑expert
+        for idx, expert in enumerate(sparse_moe.experts):
+            start = idx * dense_moe.intermediate_size
+            end   = (idx + 1) * dense_moe.intermediate_size
+
+            expert.gate_proj.weight.data.copy_(
+                dense_moe.fused_gate_proj.weight.data[start:end]
+            )
+            expert.up_proj.weight.data.copy_(
+                dense_moe.fused_up_proj.weight.data[start:end]
+            )
+            expert.down_proj.weight.data.copy_(
+                dense_moe.fused_down_proj.weight.data[:, start:end]
+            )
+
+        _replace_submodule(model, name, sparse_moe)
+    return model