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

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`,
+        # all-gather over the graft world; sender ranks contribute (value,
-        # recv ranks contribute None (their local target is private and
+        # extras) tuples, recv ranks contribute None (their local target is
-        # shouldn't leak). all_gather_object is pickle-routed, so tensor
+        # private and shouldn't leak).
-        # shapes may differ across sender ranks.
         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)
+            (c[0].to(value.device) if isinstance(c[0], torch.Tensor) else c[0])
-            for t in sender_contribs
+            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:
+        fn = self._default_transform if path is None else _load_function(path)
-            return self._default_transform(received_list, target=target)
+        return fn(graft_input)
-        return _load_function(path)(received_list, target)
 
     @staticmethod
-    def _default_transform(
+    def _default_transform(graft_input: GraftTransformInput) -> torch.Tensor:
-        received_list: list, *, target: torch.Tensor
-    ) -> 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

@@ -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"
+            "def transform(graft_input):\n"
-            "    return received_list[0] * 2\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"
+            "def transform(graft_input):\n"
-            "    return torch.zeros(99, device=target.device)\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"
+            "def transform(graft_input):\n"
-            "    rl = received_list\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"
+            "def transform(graft_input):\n"
-            "    rl = received_list\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"
+            "def transform(graft_input):\n"
-            "    expected_shapes = [(i + 1,) for i in range(len(received_list))]\n"
+            "    expected_shapes = [(i + 1,) for i in range(len(graft_input.received_list))]\n"
-            "    actual_shapes = [tuple(t.shape) for t in 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(