Support non-intrusive dumping in dumper (#19068)

python/sglang/srt/debug_utils/dumper.py

@@ -231,6 +231,12 @@ class _Dumper:
             k: v for k, v in (self._global_ctx | kwargs).items() if v is not None
         }
 
+    def register_non_intrusive_dumper(
+        self,
+        model: "torch.nn.Module",
+    ) -> "_NonIntrusiveDumper":
+        return _NonIntrusiveDumper(dumper=self, model=model)
+
     # ------------------------------- public :: secondary ---------------------------------
 
     def configure(self, **kwargs) -> None:
@@ -465,6 +471,79 @@ class _Dumper:
             print(f"[Dumper] Choose partial_name={name}")
 
 
+# -------------------------------------- hook dumper ------------------------------------------
+
+
+class _NonIntrusiveDumper:
+    """Registers forward hooks on model modules to non-invasively dump tensor outputs."""
+
+    _NAME_PREFIX = "non_intrusive__"
+
+    def __init__(
+        self,
+        dumper: _Dumper,
+        model: "torch.nn.Module",
+    ):
+        self._dumper = dumper
+
+        for module_name, module in model.named_modules():
+            module.register_forward_hook(
+                self._make_forward_hook(module_name=module_name)
+            )
+
+    def _make_forward_hook(self, module_name: str):
+        def _hook(_module, input, output):
+            for i, item in enumerate(input):
+                self._dump_value(module_name, item, role=f"inputs.{i}")
+
+            if output is not None:
+                self._dump_value(module_name, output, role="output")
+
+        return _hook
+
+    def _dump_value(self, module_name: str, value, role: str) -> None:
+        for key, tensor in self._convert_value(value).items():
+            parts = [p for p in (module_name, role, key) if p]
+            self._dumper.dump(self._NAME_PREFIX + ".".join(parts), tensor)
+
+    @staticmethod
+    def _convert_value(value) -> dict[str, torch.Tensor]:
+        if isinstance(value, torch.Tensor):
+            return {"": value}
+
+        if isinstance(value, (tuple, list)):
+            tensors = [t for t in value if isinstance(t, torch.Tensor)]
+            if len(tensors) == 1:
+                return {"": tensors[0]}
+            return {str(i): t for i, t in enumerate(tensors)}
+
+        # SGLang specific
+        try:
+            from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+            from sglang.srt.model_executor.forward_batch_info import (
+                ForwardBatch,
+                PPProxyTensors,
+            )
+
+            if isinstance(value, LogitsProcessorOutput):
+                return {"next_token_logits": value.next_token_logits}
+            if isinstance(value, ForwardBatch):
+                return {
+                    "input_ids": value.input_ids,
+                    "seq_lens": value.seq_lens,
+                    "positions": value.positions,
+                }
+            if isinstance(value, PPProxyTensors):
+                return {k: v for k, v in value.tensors.items()}
+        except ImportError:
+            pass
+
+        # Megatron specific
+        # TODO
+
+        return {}
+
+
 # -------------------------------------- util fn ------------------------------------------
 
 

test/registered/debug_utils/test_dumper.py

@@ -953,5 +953,314 @@ class TestDumperHttp:
         assert resp.status_code == 400
 
 
+class TestNonIntrusiveDumper:
+    _PREFIX = "non_intrusive__"
+
+    @staticmethod
+    def _assert_captured_contains(
+        captured: dict, expected: list[str], prefix: str = "non_intrusive__"
+    ) -> None:
+        for suffix in expected:
+            key = f"{prefix}{suffix}"
+            assert key in captured, f"missing {key}"
+
+    @staticmethod
+    def _wrap_as_outer(inner_cls: type) -> torch.nn.Module:
+        """Wrap an inner module class as OuterModel.model, mimicking typical model nesting."""
+
+        class OuterModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.model = inner_cls()
+
+            def forward(self, *args, **kwargs):
+                return self.model(*args, **kwargs)
+
+        return OuterModel()
+
+    def test_basic_inputs_and_outputs(self, tmp_path):
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                return self.relu(self.linear(x))
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            output = model(x)
+
+        self._assert_captured_contains(
+            captured,
+            [
+                "output",
+                "inputs.0",
+                "model.output",
+                "model.inputs.0",
+                "model.linear.output",
+                "model.linear.inputs.0",
+                "model.relu.output",
+                "model.relu.inputs.0",
+            ],
+        )
+        P = self._PREFIX
+        assert torch.allclose(captured[f"{P}output"]["value"], output)
+
+    def test_hooks_all_module_levels(self, tmp_path):
+        class Attention(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.qkv_proj = torch.nn.Linear(4, 12)
+                self.o_proj = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                _qkv = self.qkv_proj(x)
+                return self.o_proj(x)
+
+        class Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.self_attn = Attention()
+                self.mlp = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                x = self.self_attn(x)
+                return self.mlp(x)
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer()])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = layer(x)
+                return x
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        self._assert_captured_contains(
+            captured,
+            [
+                "output",
+                "model.output",
+                "model.layers.0.output",
+                "model.layers.0.self_attn.output",
+                "model.layers.0.self_attn.qkv_proj.output",
+                "model.layers.0.self_attn.o_proj.output",
+                "model.layers.0.mlp.output",
+                "model.layers.0.self_attn.qkv_proj.inputs.0",
+                "model.layers.0.self_attn.o_proj.inputs.0",
+                "model.layers.0.mlp.inputs.0",
+            ],
+        )
+        P = self._PREFIX
+        assert f"{P}model.layers.output" not in captured
+
+    def test_multi_tensor_tuple_output(self, tmp_path):
+        class TupleModule(torch.nn.Module):
+            def forward(self, x):
+                return x, x * 2
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.split = TupleModule()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                a, b = self.split(x)
+                return self.linear(a + b)
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.split.output.0" in captured
+        assert "non_intrusive__model.split.output.1" in captured
+        assert torch.allclose(
+            captured["non_intrusive__model.split.output.0"]["value"], x
+        )
+
+    def test_single_tensor_tuple_collapses(self, tmp_path):
+        class SingleTupleModule(torch.nn.Module):
+            def forward(self, x):
+                return (x * 3,)
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.wrap = SingleTupleModule()
+
+            def forward(self, x):
+                return self.wrap(x)[0]
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.wrap.output" in captured
+        assert "non_intrusive__model.wrap.output.0" not in captured
+
+    def test_multiple_forward_inputs(self, tmp_path):
+        class TwoInputModule(torch.nn.Module):
+            def forward(self, x, mask):
+                return x * mask
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.mul = TwoInputModule()
+
+            def forward(self, x):
+                mask = torch.ones_like(x)
+                return self.mul(x, mask)
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.mul.inputs.0" in captured
+        assert "non_intrusive__model.mul.inputs.1" in captured
+
+    def test_none_output_only_dumps_inputs(self, tmp_path):
+        class NoneModule(torch.nn.Module):
+            def forward(self, x):
+                return None
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.sink = NoneModule()
+
+            def forward(self, x):
+                self.sink(x)
+                return x
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.sink.inputs.0" in captured
+        assert not any(
+            k.startswith("non_intrusive__model.sink.output") for k in captured
+        )
+
+    def test_non_tensor_value_silently_skipped(self, tmp_path):
+        class IntModule(torch.nn.Module):
+            def forward(self, x):
+                return 42
+
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.const = IntModule()
+
+            def forward(self, x):
+                self.const(x)
+                return x
+
+        d = _make_test_dumper(tmp_path)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.const.inputs.0" in captured
+        assert not any(
+            k.startswith("non_intrusive__model.const.output") for k in captured
+        )
+
+    def test_root_module_name_no_malformed_dots(self, tmp_path):
+        d = _make_test_dumper(tmp_path)
+        model = torch.nn.Linear(4, 4)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        for key in captured:
+            assert not key.startswith("non_intrusive__."), f"malformed key: {key}"
+            assert ".." not in key, f"double dot in key: {key}"
+
+        assert "non_intrusive__output" in captured
+        assert "non_intrusive__inputs.0" in captured
+
+    def test_respects_dumper_filter(self, tmp_path):
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                return self.relu(self.linear(x))
+
+        d = _make_test_dumper(
+            tmp_path, filter="name=non_intrusive__model.linear.output"
+        )
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert "non_intrusive__model.linear.output" in captured
+        assert "non_intrusive__model.relu.output" not in captured
+        assert "non_intrusive__model.linear.inputs.0" not in captured
+
+    def test_disabled_dumper_no_output(self, tmp_path):
+        class Inner(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        d = _make_test_dumper(tmp_path)
+        d.configure(enable=False)
+        model = self._wrap_as_outer(Inner)
+        d.register_non_intrusive_dumper(model)
+
+        x = torch.randn(2, 4)
+        with d.capture_output() as captured:
+            model(x)
+
+        assert len(captured) == 0
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))