minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support user-supplied recv-side transform in dumper grafter (#24509)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-06 16:56:52 +08:00
committed by GitHub
parent 9a65f0ac26
commit ebd64f5d40
2 changed files with 283 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/dumper.py
+78 -8
View File
@@ -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,
test/registered/debug_utils/test_dumper.py
+205
View File
@@ -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__]))