From 0384c459a7b707a288751b77e5ecd4ad557f0585 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 22 Feb 2026 16:04:02 +0800 Subject: [PATCH] Support non-intrusive dumping in dumper (#19068) --- python/sglang/srt/debug_utils/dumper.py | 79 ++++++ test/registered/debug_utils/test_dumper.py | 309 +++++++++++++++++++++ 2 files changed, 388 insertions(+) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 7a830ce80..3edc59f91 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/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 ------------------------------------------ diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 53abdb5db..575d9f0eb 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/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__]))