Upload modeling_hunyuan.py

modeling_hunyuan.py

@@ -1605,6 +1605,68 @@ 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
+
 
 @add_start_docstrings(
     """
@@ -1845,66 +1907,4 @@ 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)
-
-# -----------------------------------------------------------------------
-# 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
+        return out.view(bsz, seq_len, self.hidden_size)