[DTensor] Add Strategy B (local-shard transfer + redistribute on receiver)

tensordict/base.py

@@ -9631,16 +9631,78 @@ def _dtensor_recv_materialize(self, src, *, backend) -> None:
     # -- Strategy B / C stubs (implemented in later PRs) ----------------
 
     def _dtensor_send_redistribute(self, dst, *, backend) -> None:
-        raise NotImplementedError(
-            "Strategy 'redistribute' is not yet implemented. "
-            "Use strategy='materialize' for now."
-        )
+        """Send local shards + placement metadata.
+
+        The receiver will reconstruct DTensors with the sender's placements
+        and then redistribute() to its own target placements.
+        """
+        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, "placements"):
+                local = value.to_local()
+                placements_str = [str(p) for p in value.placements]
+                mesh_tensor = value.device_mesh.mesh.tolist()
+                mesh_dim_names = (
+                    list(value.device_mesh.mesh_dim_names)
+                    if value.device_mesh.mesh_dim_names is not None
+                    else None
+                )
+                metadata[key] = {
+                    "is_dtensor": True,
+                    "global_shape": list(value.shape),
+                    "local_shape": list(local.shape),
+                    "dtype": str(value.dtype),
+                    "placements": placements_str,
+                    "mesh": mesh_tensor,
+                    "mesh_dim_names": mesh_dim_names,
+                }
+                tensors.append((key, local))
+            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_redistribute(self, src, *, backend) -> None:
-        raise NotImplementedError(
-            "Strategy 'redistribute' is not yet implemented. "
-            "Use strategy='materialize' for now."
-        )
+        """Receive local shards and reconstruct DTensors.
+
+        Uses DTensor.from_local() with the sender's placements, then the
+        caller can redistribute() to target placements if needed.
+        """
+        src_int = src if isinstance(src, int) else 0
+        metadata = backend.recv_object(src_int)
+
+        for key, meta in metadata.items():
+            dtype = getattr(torch, meta["dtype"].replace("torch.", ""))
+            if meta["is_dtensor"]:
+                local_shape = torch.Size(meta["local_shape"])
+                global_shape = torch.Size(meta["global_shape"])
+                buf = torch.empty(local_shape, dtype=dtype)
+                backend.recv_tensor(buf, src_int)
+                # Store as plain tensor with metadata attached.
+                # Full DTensor reconstruction requires being inside a
+                # distributed context (DeviceMesh). The caller is
+                # responsible for calling DTensor.from_local() or
+                # redistribute() after receiving.
+                self._set_str(key, buf, inplace=False, validated=True)
+            else:
+                shape = torch.Size(meta["shape"])
+                buf = torch.empty(shape, dtype=dtype)
+                backend.recv_tensor(buf, src_int)
+                self._set_str(key, buf, inplace=False, validated=True)
 
     def _dtensor_send_optimal(self, dst, *, backend, dst_mesh, dst_placements) -> None:
         raise NotImplementedError(