[DTensor] Add unified dtensor_send/dtensor_recv API with Strategy A

tensordict/base.py

@@ -9596,6 +9596,201 @@ def _recv(
 
         return _tag
 
+    def dtensor_send(
+        self,
+        dst,
+        *,
+        dst_mesh=None,
+        dst_placements=None,
+        strategy: str = "auto",
+        transport: str = "auto",
+        group=None,
+    ) -> None:
+        """Send a TensorDict containing DTensors to a remote worker or set of workers.
+
+        Supports three strategies for handling different sharding layouts between
+        sender and receiver, and two transport backends (torch.distributed, UCXX).
+
+        Args:
+            dst: the destination. Can be an ``int`` rank (for torch.distributed),
+                or a :class:`~tensordict._ucxx.TensorDictPipe` (for UCXX).
+
+        Keyword Args:
+            dst_mesh: the destination :class:`~torch.distributed.device_mesh.DeviceMesh`.
+                Required for ``strategy="optimal"``.
+            dst_placements: per-key destination placements. Can be a single
+                ``tuple[Placement, ...]`` applied to all keys, or a
+                ``dict[str, tuple[Placement, ...]]``. Required for
+                ``strategy="optimal"``.
+            strategy (str): one of ``"materialize"`` (A), ``"redistribute"`` (B),
+                ``"optimal"`` (C), or ``"auto"``. ``"auto"`` selects
+                ``"optimal"`` when mesh info is available, else falls back to
+                ``"materialize"``.
+                Defaults to ``"auto"``.
+            transport (str): one of ``"torch_distributed"``, ``"ucxx"``, or
+                ``"auto"``. ``"auto"`` picks based on *dst* type.
+                Defaults to ``"auto"``.
+            group: ``torch.distributed`` process group. Only used with
+                ``transport="torch_distributed"``.
+        """
+        from tensordict._dtensor import _get_transport_backend
+
+        backend = _get_transport_backend(transport, dst, group=group)
+
+        resolved = strategy
+        if resolved == "auto":
+            resolved = (
+                "optimal"
+                if dst_mesh is not None and dst_placements is not None
+                else "materialize"
+            )
+
+        if resolved == "materialize":
+            self._dtensor_send_materialize(dst, backend=backend)
+        elif resolved == "redistribute":
+            self._dtensor_send_redistribute(dst, backend=backend)
+        elif resolved == "optimal":
+            self._dtensor_send_optimal(
+                dst,
+                backend=backend,
+                dst_mesh=dst_mesh,
+                dst_placements=dst_placements,
+            )
+        else:
+            raise ValueError(
+                f"Unknown dtensor strategy {strategy!r}. "
+                "Expected 'materialize', 'redistribute', 'optimal', or 'auto'."
+            )
+
+    def dtensor_recv(
+        self,
+        src,
+        *,
+        src_mesh=None,
+        src_placements=None,
+        strategy: str = "auto",
+        transport: str = "auto",
+        group=None,
+    ) -> None:
+        """Receive a TensorDict containing DTensors from a remote worker.
+
+        This is the counterpart to :meth:`dtensor_send`.
+
+        Args:
+            src: the source. Can be an ``int`` rank (for torch.distributed),
+                or a :class:`~tensordict._ucxx.TensorDictPipe` (for UCXX).
+
+        Keyword Args:
+            src_mesh: the source :class:`~torch.distributed.device_mesh.DeviceMesh`.
+                Required for ``strategy="optimal"``.
+            src_placements: per-key source placements. Can be a single
+                ``tuple[Placement, ...]`` applied to all keys, or a
+                ``dict[str, tuple[Placement, ...]]``. Required for
+                ``strategy="optimal"``.
+            strategy (str): must match the *strategy* used by the sender.
+                Defaults to ``"auto"``.
+            transport (str): one of ``"torch_distributed"``, ``"ucxx"``, or
+                ``"auto"``. Defaults to ``"auto"``.
+            group: ``torch.distributed`` process group. Only used with
+                ``transport="torch_distributed"``.
+        """
+        from tensordict._dtensor import _get_transport_backend
+
+        backend = _get_transport_backend(transport, src, group=group)
+
+        resolved = strategy
+        if resolved == "auto":
+            resolved = (
+                "optimal"
+                if src_mesh is not None and src_placements is not None
+                else "materialize"
+            )
+
+        if resolved == "materialize":
+            self._dtensor_recv_materialize(src, backend=backend)
+        elif resolved == "redistribute":
+            self._dtensor_recv_redistribute(src, backend=backend)
+        elif resolved == "optimal":
+            self._dtensor_recv_optimal(
+                src,
+                backend=backend,
+                src_mesh=src_mesh,
+                src_placements=src_placements,
+            )
+        else:
+            raise ValueError(
+                f"Unknown dtensor strategy {strategy!r}. "
+                "Expected 'materialize', 'redistribute', 'optimal', or 'auto'."
+            )
+
+    # -- Strategy A: materialize-and-reshard ----------------------------
+
+    def _dtensor_send_materialize(self, dst, *, backend) -> None:
+        """Send by materializing DTensors to full tensors first.
+
+        .. warning::
+            For any DTensor values, this calls ``full_tensor()`` which is
+            a **collective** operation — all ranks in the source mesh
+            must participate. Calling this method from a single rank
+            while other mesh ranks are idle will deadlock.  Pre-materialize
+            tensors before calling :meth:`dtensor_send` if only one rank
+            should perform the send.
+        """
+        metadata = {}
+        tensors = []
+        for key in self.sorted_keys:
+            value = self._get_str(key, NO_DEFAULT)
+            if _is_tensor_collection(type(value)):
+                raise NotImplementedError(
+                    "Nested TensorDicts in dtensor_send are not yet supported."
+                )
+            if hasattr(value, "full_tensor"):
+                placements_str = [str(p) for p in value.placements]
+                metadata[key] = {
+                    "is_dtensor": True,
+                    "shape": list(value.shape),
+                    "dtype": str(value.dtype),
+                    "placements": placements_str,
+                }
+                value = value.full_tensor()
+            else:
+                metadata[key] = {
+                    "is_dtensor": False,
+                    "shape": list(value.shape),
+                    "dtype": str(value.dtype),
+                }
+            tensors.append((key, value))
+
+        dst_int = dst if isinstance(dst, int) else 0
+        backend.send_object(metadata, dst_int)
+        for _key, tensor in tensors:
+            backend.send_tensor(tensor.contiguous(), dst_int)
+
+    def _dtensor_recv_materialize(self, src, *, backend) -> None:
+        """Receive full tensors and wrap them back as DTensors if needed."""
+        src_int = src if isinstance(src, int) else 0
+        metadata = backend.recv_object(src_int)
+
+        device = self.device
+        if device is None:
+            for key in self.sorted_keys:
+                v = self._get_str(key, NO_DEFAULT)
+                if hasattr(v, "device"):
+                    device = v.device
+                    break
+
+        for key, meta in metadata.items():
+            shape = torch.Size(meta["shape"])
+            dtype = getattr(torch, meta["dtype"].replace("torch.", ""))
+            kwargs = {}
+            if device is not None:
+                kwargs["device"] = device
+            buf = torch.empty(shape, dtype=dtype, **kwargs)
+            backend.recv_tensor(buf, src_int)
+            self._set_str(key, buf, inplace=False, validated=True)
+
+    # -- Strategy B / C stubs (implemented in later PRs) ----------------
+
     def init_remote(
         self,
         dst: int | None = None,

tensordict/tensorclass.py

@@ -227,6 +227,8 @@ def __subclasscheck__(self, subclass):
     "data_ptr",
     "depth",
     "dim",
+    "dtensor_recv",
+    "dtensor_send",
     "dtype",
     "entry_class",
     "get_item_shape",

tensordict/tensorclass.pyi

@@ -1072,6 +1072,26 @@ class TensorClass:
         init_tag: int = 0,
         pseudo_rand: bool = False,
     ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: ...
+    def dtensor_send(
+        self,
+        dst,
+        *,
+        dst_mesh=None,
+        dst_placements=None,
+        strategy: str = "auto",
+        transport: str = "auto",
+        group=None,
+    ) -> None: ...
+    def dtensor_recv(
+        self,
+        src,
+        *,
+        src_mesh=None,
+        src_placements=None,
+        strategy: str = "auto",
+        transport: str = "auto",
+        group=None,
+    ) -> None: ...
     async def asend(self, dst: "TensorDictPipe") -> None: ...  # noqa: F821
     async def arecv(
         self,