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

Support per-call extras and dataclass transform input in dumper grafter (#24511)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-06 16:57:44 +08:00
committed by GitHub
parent 833279eb2e
commit 75943cfbcf
2 changed files with 161 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/dumper.py
+68 -19
View File
@@ -279,6 +279,7 @@ class _Dumper:
save: bool = True,
dims: Optional[str] = None,
dims_grad: Optional[str] = None,
grafter_extras: Optional[dict] = None,
**kwargs,
) -> None:
value_meta: dict = {}
@@ -302,6 +303,7 @@ class _Dumper:
grad_tag="Dumper.Grad",
value_meta_only_fields=value_meta,
grad_meta_only_fields=grad_meta,
grafter_extras=grafter_extras,
)
def dump_model(
@@ -458,6 +460,7 @@ class _Dumper:
grad_tag: str,
value_meta_only_fields: Optional[dict] = None,
grad_meta_only_fields: Optional[dict] = None,
grafter_extras: Optional[dict] = None,
) -> None:
self._http_manager # noqa: B018
@@ -480,7 +483,7 @@ class _Dumper:
recompute_meta = recompute_status.to_pseudo_parallel_meta()
value = _materialize_value(value)
self._grafter.maybe_intercept(value=value, tags=tags)
self._grafter.maybe_intercept(value=value, tags=tags, extras=grafter_extras)
if enable_value:
self._dump_single(
@@ -804,6 +807,29 @@ class _GraftDirection(enum.Enum):
T2B = "t2b" # name flows target -> baseline
@dataclass
class GraftTransformInput:
"""Single argument passed to a user-supplied transform function.
User transforms have signature::
def transform(graft_input: GraftTransformInput) -> torch.Tensor: ...
The dataclass shape lets us add fields (e.g., direction, sender ranks)
later without breaking existing transforms.
"""
# Full dumper.dump tags dict (name + recompute_status + extra_kwargs + ctx).
tags: "dict[str, Any]"
# One tensor per sender rank, in sender-rank order.
received_list: "list[torch.Tensor]"
# Parallel list of per-sender `grafter_extras` (the dict passed to
# dumper.dump on each sender; None if the sender omitted it).
received_extras_list: "list[Optional[dict]]"
# Recv side's local tensor that will be copy_'d into.
target: "torch.Tensor"
class _Grafter:
"""1+1 cross-system tensor grafter.
@@ -819,7 +845,13 @@ class _Grafter:
self._config = config
self._pg: Optional[dist.ProcessGroup] = None
def maybe_intercept(self, *, value, tags: dict) -> None:
def maybe_intercept(
self,
*,
value,
tags: dict,
extras: Optional[dict] = None,
) -> None:
cfg = self._config
if not cfg.grafter_enable:
return
@@ -842,35 +874,45 @@ class _Grafter:
role = _GraftRole(cfg.grafter_role)
is_send = self._is_sender(role=role, direction=direction)
# all-gather over the graft world; sender ranks contribute `value`,
# recv ranks contribute None (their local target is private and
# shouldn't leak). all_gather_object is pickle-routed, so tensor
# shapes may differ across sender ranks.
# all-gather over the graft world; sender ranks contribute (value,
# extras) tuples, recv ranks contribute None (their local target is
# private and shouldn't leak).
total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size
my_contribution = value if is_send else None
my_contribution = (value, extras) if is_send else None
gathered: list = [None] * total_world
dist.all_gather_object(gathered, my_contribution, group=self._pg)
if is_send:
_log(f"[Grafter] send role={role.value} dir={direction.value} tags={tags}")
_log(
f"[Grafter] send role={role.value} dir={direction.value} "
f"tags={tags} extras={extras}"
)
return
sender_contribs = self._sender_slice(direction=direction, gathered=gathered)
# Pickled CUDA tensors restore to their original-device name; that
# may not match this process's local device, so normalize.
sender_tensors = [
(t.to(value.device) if isinstance(t, torch.Tensor) else t)
for t in sender_contribs
(c[0].to(value.device) if isinstance(c[0], torch.Tensor) else c[0])
for c in sender_contribs
]
sender_extras = [c[1] for c in sender_contribs]
# 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(sender_tensors, target=value)
value_to_override = self._apply_transform(
tags=tags,
received_list=sender_tensors,
received_extras_list=sender_extras,
target=value,
)
_log(
f"[Grafter] recv role={role.value} dir={direction.value} "
f"tags={tags} n_senders={len(sender_tensors)}"
f"tags={tags} n_senders={len(sender_tensors)} "
f"sender_extras={sender_extras}"
)
value.copy_(value_to_override)
except Exception as e:
@@ -889,22 +931,29 @@ class _Grafter:
def _apply_transform(
self,
received_list: list,
*,
tags: dict,
received_list: list,
received_extras_list: list,
target: torch.Tensor,
) -> torch.Tensor:
graft_input = GraftTransformInput(
tags=tags,
received_list=received_list,
received_extras_list=received_extras_list,
target=target,
)
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)
fn = self._default_transform if path is None else _load_function(path)
return fn(graft_input)
@staticmethod
def _default_transform(
received_list: list, *, target: torch.Tensor
) -> torch.Tensor:
def _default_transform(graft_input: GraftTransformInput) -> torch.Tensor:
"""Identity-by-rank fallback. Requires #senders == #recvs and
shape(received_list[my_recv_rank]) == shape(target). Otherwise raises
and asks the user for a transform."""
received_list = graft_input.received_list
target = graft_input.target
my_recv_rank = dist.get_rank()
recv_world_size = dist.get_world_size()
if len(received_list) != recv_world_size:
test/registered/debug_utils/test_dumper.py
+93 -15
View File
@@ -3069,8 +3069,8 @@ class TestGrafterDistributed:
"""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"
"def transform(graft_input):\n"
" return graft_input.received_list[0] * 2\n"
)
graft_port = find_available_port(29610)
_run_graft_test(
@@ -3145,7 +3145,7 @@ class TestGrafterDistributed:
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"
"def transform(graft_input):\n"
" raise RuntimeError('intentional test error from user transform')\n"
)
graft_port = find_available_port(29635)
@@ -3200,8 +3200,8 @@ class TestGrafterDistributed:
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"
"def transform(graft_input):\n"
" return torch.zeros(99, device=graft_input.target.device)\n"
)
graft_port = find_available_port(29665)
_run_graft_test(
@@ -3245,6 +3245,84 @@ class TestGrafterDistributed:
if grafter._pg is not None:
dist.destroy_process_group(grafter._pg)
def test_extras_flow_to_recv_transform(self, tmp_path: Path):
"""Sender attaches per-call grafter_extras; recv transform reads them
and uses them to compute the override value."""
module_name = "_xform_uses_extras"
(tmp_path / f"{module_name}.py").write_text(
"import torch\n"
"def transform(graft_input):\n"
" fill = graft_input.received_extras_list[0]['fill_value']\n"
" return torch.full_like(graft_input.target, fill)\n"
)
graft_port = find_available_port(29645)
_run_graft_test(
self._test_extras_func,
graft_port=graft_port,
group_name="grafter_extras",
transform_dir=str(tmp_path),
transform_path=f"{module_name}.transform",
)
@staticmethod
def _test_extras_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"},
extras={"fill_value": 42.0},
)
else:
target = torch.zeros(3, device="cuda:1")
grafter.maybe_intercept(value=target, tags={"name": "x"})
assert target.tolist() == [42.0, 42.0, 42.0], target.tolist()
finally:
if grafter._pg is not None:
dist.destroy_process_group(grafter._pg)
def test_extras_default_none_flow(self):
"""When the sender omits `grafter_extras`, the recv transform sees a
list of Nones."""
graft_port = find_available_port(29650)
_run_graft_test(
self._test_extras_none_func,
graft_port=graft_port,
group_name="grafter_extras_none",
)
@staticmethod
def _test_extras_none_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.zeros(3, device="cuda:1")
with _capture_stdout() as captured:
grafter.maybe_intercept(value=target, tags={"name": "x"})
output = captured.getvalue()
assert "sender_extras=[None]" in output, output
assert target.tolist() == [1.0, 2.0, 3.0], target.tolist()
finally:
if grafter._pg is not None:
dist.destroy_process_group(grafter._pg)
def _run_graft_test_cpu_multi(
worker_func, *, baseline_world: int, target_world: int, **kwargs
@@ -3357,13 +3435,13 @@ class TestGrafterMultiRankCpu:
module_name = "_xform_assert_4_senders"
(tmp_path / f"{module_name}.py").write_text(
"import torch\n"
"def transform(received_list, target):\n"
" rl = received_list\n"
"def transform(graft_input):\n"
" rl = graft_input.received_list\n"
" assert len(rl) == 4, f'expected 4 senders, got {len(rl)}'\n"
" for i, t in enumerate(rl):\n"
" v = float(t.flatten()[0].item())\n"
" assert v == float(i), f'rl[{i}][0]={v}, want {float(i)}'\n"
" return torch.full_like(target, 999.0)\n"
" return torch.full_like(graft_input.target, 999.0)\n"
)
graft_port = find_available_port(29655)
_run_graft_test_cpu_multi(
@@ -3408,13 +3486,13 @@ class TestGrafterMultiRankCpu:
module_name = "_xform_assert_2_senders_t2b"
(tmp_path / f"{module_name}.py").write_text(
"import torch\n"
"def transform(received_list, target):\n"
" rl = received_list\n"
"def transform(graft_input):\n"
" rl = graft_input.received_list\n"
" assert len(rl) == 2, f'expected 2 senders, got {len(rl)}'\n"
" for i, t in enumerate(rl):\n"
" v = float(t.flatten()[0].item())\n"
" assert v == float(i + 100), f'rl[{i}][0]={v}'\n"
" return torch.full_like(target, 7.0)\n"
" return torch.full_like(graft_input.target, 7.0)\n"
)
graft_port = find_available_port(29670)
_run_graft_test_cpu_multi(
@@ -3503,11 +3581,11 @@ class TestGrafterMultiRankCpu:
module_name = "_xform_concat_mixed_shape"
(tmp_path / f"{module_name}.py").write_text(
"import torch\n"
"def transform(received_list, target):\n"
" expected_shapes = [(i + 1,) for i in range(len(received_list))]\n"
" actual_shapes = [tuple(t.shape) for t in received_list]\n"
"def transform(graft_input):\n"
" expected_shapes = [(i + 1,) for i in range(len(graft_input.received_list))]\n"
" actual_shapes = [tuple(t.shape) for t in graft_input.received_list]\n"
" assert actual_shapes == expected_shapes, actual_shapes\n"
" return torch.cat(received_list)\n"
" return torch.cat(graft_input.received_list)\n"
)
graft_port = find_available_port(29680)
_run_graft_test_cpu_multi(