Fix wrongly large dumped file and handle non intrusive hook reset in dumper (#19124)

python/sglang/srt/debug_utils/dumper.py

@@ -7,6 +7,7 @@ import socket
 import threading
 import time
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import asdict, dataclass, field, fields, replace
@@ -130,8 +131,8 @@ class DumperConfig(_BaseConfig):
     enable_output_console: bool = True
     enable_value: bool = True
     enable_grad: bool = False
-    enable_model_value: bool = True
-    enable_model_grad: bool = True
+    enable_model_value: bool = False
+    enable_model_grad: bool = False
     exp_name: Optional[str] = None
     enable_http_server: bool = True
     cleanup_previous: bool = False
@@ -200,6 +201,7 @@ class _Dumper:
     def __init__(self, *, config: DumperConfig):
         self._config = config
         self._state = _DumperState()
+        self._non_intrusives: list["_NonIntrusiveDumper"] = []
 
     # ------------------------------- public :: core ---------------------------------
 
@@ -280,7 +282,9 @@ class _Dumper:
         mode = self._config.non_intrusive_mode
         if mode == "off":
             return None
-        return _NonIntrusiveDumper(dumper=self, model=model, mode=mode)
+        non_intrusive = _NonIntrusiveDumper(dumper=self, model=model, mode=mode)
+        self._non_intrusives.append(non_intrusive)
+        return non_intrusive
 
     # ------------------------------- public :: secondary ---------------------------------
 
@@ -291,6 +295,9 @@ class _Dumper:
         self._config = self._config.with_defaults(**kwargs)
 
     def reset(self) -> None:
+        for non_intrusive in self._non_intrusives:
+            non_intrusive.remove()
+        self._non_intrusives.clear()
         self._state = _DumperState()
 
     @contextmanager
@@ -488,6 +495,7 @@ class _NonIntrusiveDumper:
     ):
         self._dumper = dumper
         self._mode = mode
+        self._handles: list = []
 
         for module_name, module in model.named_modules():
             if ctx := self._detect_module_ctx(module_name, module):
@@ -498,13 +506,18 @@ class _NonIntrusiveDumper:
                 module_name=module_name, is_root=is_root
             )
             hook = self._make_forward_hook(module_name=module_name, is_root=is_root)
-            _register_forward_hook_or_replace_fn(
+            self._handles += _register_forward_hook_or_replace_fn(
                 module,
                 pre_hook=pre_hook,
                 hook=hook,
                 mode="replace_fn" if is_root else "hook",
             )
 
+    def remove(self) -> None:
+        for handle in self._handles:
+            handle.remove()
+        self._handles.clear()
+
     @classmethod
     def _detect_module_ctx(
         cls, module_name: str, module: "torch.nn.Module"
@@ -518,12 +531,16 @@ class _NonIntrusiveDumper:
 
     def _register_ctx_hooks(self, module: "torch.nn.Module", *, ctx: dict) -> None:
         clear_ctx = {k: None for k in ctx}
-        module.register_forward_pre_hook(
-            lambda _mod, _input, _ctx=ctx: self._dumper.set_ctx(**_ctx)
+        self._handles.append(
+            module.register_forward_pre_hook(
+                lambda _mod, _input, _ctx=ctx: self._dumper.set_ctx(**_ctx)
+            )
         )
-        module.register_forward_hook(
-            lambda _mod, _input, _output, _clear=clear_ctx: self._dumper.set_ctx(
-                **_clear
+        self._handles.append(
+            module.register_forward_hook(
+                lambda _mod, _input, _output, _clear=clear_ctx: self._dumper.set_ctx(
+                    **_clear
+                )
             )
         )
 
@@ -578,7 +595,7 @@ def _register_forward_hook_or_replace_fn(
     pre_hook,
     hook,
     mode: str,
-) -> None:
+) -> list:
     """Attach pre/post forward hooks to *module*.
 
     mode="hook"       — standard ``register_forward_pre_hook`` / ``register_forward_hook``
@@ -586,10 +603,15 @@ def _register_forward_hook_or_replace_fn(
     mode="replace_fn" — monkey-patch ``module.forward`` so hooks fire even when
                         callers invoke ``.forward()`` directly (as sglang does for the
                         root model).
+
+    Returns a list of handle objects with a ``.remove()`` method that undoes
+    the registration.
     """
     if mode == "hook":
-        module.register_forward_pre_hook(pre_hook)
-        module.register_forward_hook(hook)
+        return [
+            module.register_forward_pre_hook(pre_hook),
+            module.register_forward_hook(hook),
+        ]
     elif mode == "replace_fn":
         original_forward = module.forward
 
@@ -601,6 +623,13 @@ def _register_forward_hook_or_replace_fn(
             return output
 
         module.forward = _wrapped
+
+        class _Handle:
+            def remove(self) -> None:
+                assert module.forward is _wrapped
+                module.forward = original_forward
+
+        return [_Handle()]
     else:
         raise ValueError(f"Unknown mode {mode!r}")
 
@@ -609,6 +638,7 @@ def _register_forward_hook_or_replace_fn(
 
 
 def _torch_save(value, path: str):
+    value = _clone_if_view(value)
     try:
         try:
             return torch.save(value, path)
@@ -623,17 +653,32 @@ def _torch_save(value, path: str):
         print(f"[Dumper] Observe error={e} when saving data, skip the tensor")
 
 
-def _strip_parameter(value):
-    """Strip nn.Parameter to plain Tensor so it can be pickled."""
-    if isinstance(value, torch.nn.Parameter):
-        return value.data
-    if isinstance(value, dict) and isinstance(
-        value.get("value"), torch.nn.Parameter
-    ):
-        return {**value, "value": value["value"].data}
+def _map_tensor(value, fn: Callable[[torch.Tensor], torch.Tensor]):
+    if isinstance(value, dict):
+        return {k: _map_tensor(v, fn) for k, v in value.items()}
+    if isinstance(value, torch.Tensor):
+        return fn(value)
     return value
 
 
+def _clone_if_view(value):
+    def _fn(t: torch.Tensor) -> torch.Tensor:
+        if t.untyped_storage().nbytes() > t.nelement() * t.element_size():
+            return t.clone()
+        return t
+
+    return _map_tensor(value, _fn)
+
+
+def _strip_parameter(value):
+    def _fn(t: torch.Tensor) -> torch.Tensor:
+        if isinstance(t, torch.nn.Parameter):
+            return t.data
+        return t
+
+    return _map_tensor(value, _fn)
+
+
 def _collective_with_timeout(fn, operation_name: str, timeout_seconds: int = 60):
     completed = threading.Event()