From ebd64f5d40e20a77b35ca6fae9feaf1096737d91 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 6 May 2026 16:56:52 +0800 Subject: [PATCH] Support user-supplied recv-side transform in dumper grafter (#24509) --- python/sglang/srt/debug_utils/dumper.py | 86 ++++++++- test/registered/debug_utils/test_dumper.py | 205 +++++++++++++++++++++ 2 files changed, 283 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index f94de7c01..249106a94 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -7,6 +7,7 @@ import re import socket import threading import time +import traceback from abc import ABC, abstractmethod from collections.abc import Callable from contextlib import contextmanager @@ -151,6 +152,11 @@ class DumperConfig(_BaseConfig): grafter_backend: str = "nccl" grafter_group_name: str = "graft" grafter_timeout: int = 300 + # Fully-qualified Python path "pkg.subpkg.module.fn_name". When set, the + # recv side calls this function with (received_list, target) and copies + # the result into target. None -> use the default identity-by-rank + # fallback in `_Grafter._default_transform`. + grafter_transform_path: Optional[str] = None @classmethod def _env_prefix(cls) -> str: @@ -837,21 +843,65 @@ class _Grafter: is_send = self._is_sender(role=role, direction=direction) # 1+1 broadcast: sender side ships the tensor as a pickled object; - # recv side calls `value.copy_()` with the received tensor. + # recv side feeds it through the user transform (default: identity) + # and `value.copy_()` the result. sender_rank = 0 if direction == _GraftDirection.B2T else 1 obj_list: list = [None] if is_send: obj_list = [value] _log(f"[Grafter] send role={role.value} dir={direction.value} tags={tags}") dist.broadcast_object_list(obj_list, src=sender_rank, group=self._pg) - if not is_send: - received = obj_list[0] - if isinstance(received, torch.Tensor): - # Pickled CUDA tensors restore to their original-device name; - # that may not match this process's local device, so normalize. - received = received.to(value.device) + if is_send: + return + + received = obj_list[0] + if isinstance(received, torch.Tensor): + # Pickled CUDA tensors restore to their original-device name; + # that may not match this process's local device, so normalize. + received = received.to(value.device) + # Transform + copy_ are wrapped: a buggy user transform must NOT + # crash the whole training/inference run. On error we log the full + # traceback and skip this graft point; downstream sees the recv + # side's original tensor unchanged. + try: + value_to_override = self._apply_transform([received], target=value) _log(f"[Grafter] recv role={role.value} dir={direction.value} tags={tags}") - value.copy_(received) + value.copy_(value_to_override) + except Exception as e: + _log( + f"[Grafter] recv role={role.value} dir={direction.value} " + f"tags={tags} transform/copy_ raised {type(e).__name__}: {e}; " + f"skipping graft for this call (target tensor unchanged)\n" + f"{traceback.format_exc()}" + ) + + def _apply_transform( + self, + received_list: list, + *, + target: torch.Tensor, + ) -> torch.Tensor: + path = self._config.grafter_transform_path + if path is None: + return self._default_transform(received_list, target=target) + return _load_function(path)(received_list, target) + + @staticmethod + def _default_transform( + received_list: list, *, target: torch.Tensor + ) -> torch.Tensor: + """Identity-by-rank fallback. For the 1+1 setup currently supported, + just returns the single received tensor; requires shape match.""" + candidate = received_list[0] + if candidate.shape != target.shape: + raise RuntimeError( + f"[Grafter] no grafter_transform_path set; default " + f"identity-by-rank requires matching shapes but " + f"received_list[0].shape={tuple(candidate.shape)} != " + f"target.shape={tuple(target.shape)}. Provide a transform via " + f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." + ) + return candidate def _classify_direction(self, tags: dict) -> Optional["_GraftDirection"]: cfg = self._config @@ -1339,6 +1389,26 @@ def _get_local_ip_by_remote() -> Optional[str]: return None +@functools.lru_cache(maxsize=None) +def _load_function(path: str) -> Callable: + """Resolve a fully-qualified Python path 'pkg.module.symbol' to its object. + + Copied (verbatim, minus the function-registry branch) from + miles.utils.misc.load_function -- kept inline so dumper.py has no + cross-package dependency. + """ + import importlib + + module_path, _, attr = path.rpartition(".") + if not module_path: + raise ValueError( + f"_load_function expects 'pkg.module.symbol', got {path!r} " + f"(missing dotted prefix)" + ) + module = importlib.import_module(module_path) + return getattr(module, attr) + + def _init_custom_process_group( *, backend: str, diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 8a26ff625..a7f5cec63 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -23,6 +23,7 @@ from sglang.srt.debug_utils.dumper import ( _format_tags, _get_default_exp_name, _Grafter, + _load_function, _log, _map_tensor, _materialize_value, @@ -2809,6 +2810,28 @@ class TestGrafterFilterMatching: grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "other"}) assert grafter._pg is None + def test_load_function_bad_module(self): + with pytest.raises(ModuleNotFoundError): + _load_function("no_such_pkg.no_such_module.transform") + + def test_load_function_missing_attr(self): + # `os.path` exists but has no `definitely_no_such_attr`. + with pytest.raises(AttributeError): + _load_function("os.path.definitely_no_such_attr") + + def test_load_function_no_dotted_prefix(self): + with pytest.raises(ValueError, match=r"missing dotted prefix"): + _load_function("only_one_segment") + + def test_load_function_non_callable_resolves_but_call_fails(self): + """`_load_function` itself only does attribute lookup -- it doesn't + verify the result is callable. A non-callable target manifests at + call time as TypeError; we still want the failure to be debuggable.""" + sep = _load_function("os.path.sep") # str, not a callable + assert isinstance(sep, str) + with pytest.raises(TypeError): + sep() + def _run_graft_test(worker_func, **kwargs): """Spawn one GPU-using process per role (rank 0 = baseline, rank 1 = target). @@ -2866,6 +2889,7 @@ def _make_grafter_test_config( timeout: int = 30, b2t_filter: Optional[str] = "name == 'x'", t2b_filter: Optional[str] = None, + transform_path: Optional[str] = None, ) -> DumperConfig: """Build a DumperConfig for distributed grafter tests. rank 0 -> baseline, rank 1 -> target. Both sides are world_size=1 within their own role's @@ -2883,6 +2907,7 @@ def _make_grafter_test_config( grafter_target_world_size=1, grafter_group_name=group_name, grafter_timeout=timeout, + grafter_transform_path=transform_path, ) @@ -3040,6 +3065,186 @@ class TestGrafterDistributed: if grafter._pg is not None: dist.destroy_process_group(grafter._pg) + def test_recv_with_user_transform(self, tmp_path: Path): + """User transform doubles the received tensor before copy_.""" + module_name = "_xform_user_basic" + (tmp_path / f"{module_name}.py").write_text( + "def transform(received_list, target):\n" + " return received_list[0] * 2\n" + ) + graft_port = find_available_port(29610) + _run_graft_test( + self._test_user_transform_func, + graft_port=graft_port, + group_name="grafter_transform", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + ) + + @staticmethod + def _test_user_transform_func( + rank, graft_port, group_name, transform_dir, transform_path + ): + sys.path.insert(0, transform_dir) + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, + graft_port=graft_port, + group_name=group_name, + transform_path=transform_path, + ) + ) + try: + if rank == 0: + tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.zeros(3, device="cuda:1") + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [2.0, 4.0, 6.0], f"got {target.tolist()}" + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_default_fallback_shape_mismatch_does_not_crash(self): + """When sender shape != target shape, default identity fallback raises; + the grafter must catch it, log, and leave target unchanged.""" + graft_port = find_available_port(29615) + _run_graft_test( + self._test_shape_mismatch_func, + graft_port=graft_port, + group_name="grafter_shape_mismatch", + ) + + @staticmethod + def _test_shape_mismatch_func(rank, graft_port, group_name): + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, graft_port=graft_port, group_name=group_name + ) + ) + try: + if rank == 0: + tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.tensor([7.0, 7.0, 7.0, 7.0], device="cuda:1") + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [ + 7.0, + 7.0, + 7.0, + 7.0, + ], f"target should be unchanged after shape-mismatch graft, got {target.tolist()}" + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_user_transform_exception_does_not_crash(self, tmp_path: Path): + """A user transform that raises must NOT bring down the system; the + grafter logs and skips the copy_, leaving target unchanged.""" + module_name = "_xform_throws" + (tmp_path / f"{module_name}.py").write_text( + "def transform(received_list, target):\n" + " raise RuntimeError('intentional test error from user transform')\n" + ) + graft_port = find_available_port(29635) + _run_graft_test( + self._test_transform_throws_func, + graft_port=graft_port, + group_name="grafter_throws", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + module_name=module_name, + ) + + @staticmethod + def _test_transform_throws_func( + rank, graft_port, group_name, transform_dir, transform_path, module_name + ): + sys.path.insert(0, transform_dir) + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, + graft_port=graft_port, + group_name=group_name, + transform_path=transform_path, + ) + ) + try: + if rank == 0: + tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.tensor([9.0, 9.0, 9.0], device="cuda:1") + with _capture_stdout() as captured: + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [ + 9.0, + 9.0, + 9.0, + ], f"target must be unchanged when transform throws, got {target.tolist()}" + output = captured.getvalue() + assert "transform/copy_ raised RuntimeError" in output, output + assert "intentional test error" in output, output + assert "Traceback (most recent call last)" in output, output + assert f"{module_name}.py" in output, output + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_copy_failure_does_not_crash(self, tmp_path: Path): + """If the user transform returns a tensor whose shape doesn't match + target, `value.copy_(value_to_override)` raises -- and that error + must be caught, logged with traceback, and target left unchanged.""" + module_name = "_xform_returns_wrong_shape" + (tmp_path / f"{module_name}.py").write_text( + "import torch\n" + "def transform(received_list, target):\n" + " return torch.zeros(99, device=target.device)\n" + ) + graft_port = find_available_port(29665) + _run_graft_test( + self._test_copy_failure_func, + graft_port=graft_port, + group_name="grafter_copy_fail", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + ) + + @staticmethod + def _test_copy_failure_func( + rank, graft_port, group_name, transform_dir, transform_path + ): + sys.path.insert(0, transform_dir) + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, + graft_port=graft_port, + group_name=group_name, + transform_path=transform_path, + ) + ) + try: + if rank == 0: + tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.tensor([7.0, 7.0, 7.0], device="cuda:1") + with _capture_stdout() as captured: + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [ + 7.0, + 7.0, + 7.0, + ], f"target must be unchanged on copy_ failure, got {target.tolist()}" + output = captured.getvalue() + assert "transform/copy_ raised" in output, output + assert "Traceback (most recent call last)" in output, output + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + if __name__ == "__main__": sys.exit(pytest.main([__file__]))