diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 364ebbf63..07ca7d74d 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -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: diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 6da7ef7c6..791031b75 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -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(