[BugFix] Stage collector weight sync through CPU when device sharing …

test/test_collectors.py

@@ -62,8 +62,10 @@
 from torchrl.collectors.distributed.ray import _has_ray, RayCollector
 
 from torchrl.collectors.utils import (
+    _device_shareable_across_processes,
     _make_policy_factory,
     _maybe_normalize_replay_buffer_tensordict_device,
+    _stage_unshareable_weights_on_cpu,
     _traj_chunk_ends_done,
     _traj_emit,
     _traj_ingest,
@@ -188,6 +190,130 @@ def test_weight_sync_scheme_from_backend_rejects_unknown():
         WeightSyncScheme.from_backend("unknown")
 
 
+class TestWeightSyncCPUStaging:
+    """Weights on devices without cross-process sharing support (CUDA on
+    Windows/WSL2, MPS anywhere) must be staged through CPU shared memory,
+    otherwise they are silently received as zeros and can corrupt the
+    sender's memory (https://github.com/pytorch/rl/issues/3985)."""
+
+    def test_cpu_and_meta_always_shareable(self):
+        assert _device_shareable_across_processes(torch.device("cpu"))
+        assert _device_shareable_across_processes(torch.device("meta"))
+
+    def test_mps_never_shareable(self):
+        assert not _device_shareable_across_processes(torch.device("mps"))
+
+    @pytest.mark.parametrize("ipc_supported", [True, False])
+    def test_cuda_shareable_iff_cuda_ipc(self, monkeypatch, ipc_supported):
+        import torchrl.collectors.utils as collector_utils
+
+        monkeypatch.setattr(
+            collector_utils, "_platform_supports_cuda_ipc", lambda: ipc_supported
+        )
+        assert (
+            _device_shareable_across_processes(torch.device("cuda:0")) is ipc_supported
+        )
+
+    @pytest.mark.parametrize(
+        "scheme_cls", [SharedMemWeightSyncScheme, MultiProcessWeightSyncScheme]
+    )
+    def test_params_map_stages_cuda_weights_without_ipc(self, monkeypatch, scheme_cls):
+        import torchrl.collectors.utils as collector_utils
+
+        monkeypatch.setattr(
+            collector_utils, "_platform_supports_cuda_ipc", lambda: False
+        )
+        model = nn.Linear(2, 1)
+        scheme = scheme_cls()
+        with pytest.warns(UserWarning, match="staged through CPU shared memory"):
+            params_map = scheme._get_params_map(
+                model=model,
+                devices=[torch.device("cuda:0"), torch.device("cuda:0")],
+            )
+        assert params_map[0] is params_map[1]
+        for weights in params_map.values():
+            assert weights.is_shared()
+            for tensor in weights.values(True, True):
+                assert tensor.device.type == "cpu"
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_params_map_keeps_cuda_weights_with_ipc(self, monkeypatch):
+        import torchrl.collectors.utils as collector_utils
+
+        monkeypatch.setattr(
+            collector_utils, "_platform_supports_cuda_ipc", lambda: True
+        )
+        model = nn.Linear(2, 1)
+        scheme = SharedMemWeightSyncScheme()
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            params_map = scheme._get_params_map(
+                model=model, devices=[torch.device("cuda:0")]
+            )
+        for tensor in params_map[0].values(True, True):
+            assert tensor.device.type == "cuda"
+
+    def test_params_map_cpu_weights_unchanged(self):
+        model = nn.Linear(2, 1)
+        scheme = SharedMemWeightSyncScheme()
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            params_map = scheme._get_params_map(
+                model=model, devices=[torch.device("cpu")]
+            )
+        for tensor in params_map[0].values(True, True):
+            assert tensor.device.type == "cpu"
+
+    @pytest.mark.parametrize("weight_format", ["tensordict", "state_dict"])
+    def test_stage_unshareable_weights_on_cpu(self, monkeypatch, weight_format):
+        import torchrl.collectors.utils as collector_utils
+
+        monkeypatch.setattr(
+            collector_utils, "_platform_supports_cuda_ipc", lambda: False
+        )
+        model = nn.Linear(2, 1)
+        if weight_format == "tensordict":
+            weights = TensorDict.from_module(model).data
+        else:
+            weights = {key: value.detach() for key, value in model.state_dict().items()}
+        staged = _stage_unshareable_weights_on_cpu(weights)
+        assert staged is weights
+
+    @pytest.mark.skipif(
+        not torch.backends.mps.is_available(), reason="needs an MPS device"
+    )
+    def test_multicollector_mps_policy(self):
+        device = torch.device("mps")
+        env_maker = ContinuousActionVecMockEnv
+        policy = TensorDictModule(
+            nn.LazyLinear(7), in_keys=["observation"], out_keys=["action"]
+        )
+        env = env_maker()
+        policy(env.reset())
+        policy = policy.to(device)
+        weights_before = TensorDict.from_module(policy).data.cpu().clone()
+        with pytest.warns(UserWarning, match="staged through CPU shared memory"):
+            collector = MultiSyncCollector(
+                [env_maker, env_maker],
+                policy=policy,
+                frames_per_batch=16,
+                total_frames=32,
+                policy_device=device,
+                storing_device="cpu",
+                env_device="cpu",
+            )
+        try:
+            for data in collector:
+                assert data.numel() == 16
+                break
+            weights_after = TensorDict.from_module(policy).data
+            assert_allclose_td(weights_before, weights_after.cpu())
+        finally:
+            collector.shutdown()
+            del collector
+
+
 def test_ray_weight_sync_copy_preserves_config_and_resets_runtime():
     with pytest.warns(UserWarning, match="state_dict strategy is experimental"):
         scheme = RayWeightSyncScheme(strategy="state_dict", backend="nccl")

torchrl/collectors/utils.py

@@ -5,6 +5,8 @@
 from __future__ import annotations
 
 import contextlib
+import platform
+import sys
 import warnings
 from collections.abc import Callable, Sequence
 from typing import Literal
@@ -416,6 +418,113 @@ def _map_to_cpu_if_needed(x):
     return x
 
 
+def _platform_supports_cuda_ipc() -> bool:
+    """Whether torch.multiprocessing can share CUDA tensors with other processes.
+
+    CUDA tensors are exchanged between processes through CUDA IPC handles,
+    which the CUDA driver only implements on native Linux. On Windows and on
+    WSL2 the exchange fails silently: the receiving process observes zeroed
+    tensors and the sender's CUDA memory can be corrupted as well. See
+    https://github.com/pytorch/pytorch/issues/149155 and
+    https://github.com/pytorch/rl/issues/3985.
+    """
+    if sys.platform != "linux":
+        return False
+    return "microsoft" not in platform.uname().release.lower()
+
+
+def _device_shareable_across_processes(device: torch.device | str | int) -> bool:
+    """Whether tensors on ``device`` can be sent to another process without a copy.
+
+    CPU and meta tensors can be shared on every platform. CUDA tensors rely on
+    CUDA IPC, which is only available on native Linux. Other backends
+    (MPS, NPU, ...) cannot be shared across processes at all.
+    """
+    if not isinstance(device, torch.device):
+        device = torch.device(device)
+    if device.type in ("cpu", "meta"):
+        return True
+    if device.type == "cuda":
+        return _platform_supports_cuda_ipc()
+    return False
+
+
+def _unshareable_devices(
+    device: torch.device | str | int | None, weights: TensorDictBase | None
+) -> list[torch.device]:
+    """Return the devices of ``weights`` that cannot cross process boundaries.
+
+    ``device`` is the target device the weights are about to be moved to; when
+    it is ``None`` the weights are shipped on their current devices, so each
+    leaf tensor's device is checked instead.
+    """
+    if device is not None:
+        if _device_shareable_across_processes(device):
+            return []
+        return [torch.device(device)]
+    if weights is None:
+        return []
+    return sorted(
+        {
+            tensor.device
+            for tensor in weights.values(True, True)
+            if not _device_shareable_across_processes(tensor.device)
+        },
+        key=str,
+    )
+
+
+def _stage_unshareable_weights_on_cpu(
+    weights: TensorDictBase | dict | None,
+) -> TensorDictBase | dict | None:
+    """Return ``weights`` moved to CPU when any leaf cannot cross process boundaries.
+
+    Accepts a ``TensorDictBase`` or a state-dict mapping and returns the input
+    unchanged when every leaf tensor can be shared with other processes.
+    A warning is emitted when staging happens.
+    """
+    if isinstance(weights, TensorDictBase):
+        offending = _unshareable_devices(None, weights)
+        if offending:
+            _warn_cpu_staging(offending)
+            weights = weights.to("cpu")
+        return weights
+    if isinstance(weights, dict):
+        offending = sorted(
+            {
+                value.device
+                for value in weights.values()
+                if isinstance(value, torch.Tensor)
+                and not _device_shareable_across_processes(value.device)
+            },
+            key=str,
+        )
+        if offending:
+            _warn_cpu_staging(offending)
+            weights = {
+                key: value.cpu() if isinstance(value, torch.Tensor) else value
+                for key, value in weights.items()
+            }
+        return weights
+    return weights
+
+
+def _warn_cpu_staging(devices: Sequence[torch.device]) -> None:
+    """Warn that weights bound for ``devices`` are staged through CPU shared memory."""
+    device_names = sorted({str(d) for d in devices})
+    warnings.warn(
+        f"Policy weights on device(s) {device_names} cannot be shared across processes "
+        "on this platform: sharing CUDA tensors requires CUDA IPC, which is only "
+        "available on native Linux (not Windows or WSL2), and backends such as MPS "
+        "cannot be shared at all. Sending such tensors to another process silently "
+        "corrupts them (they are received as zeros and the sender's memory can be "
+        "zeroed too, see https://github.com/pytorch/rl/issues/3985). The weights will "
+        "be staged through CPU shared memory instead and moved back to the policy "
+        "device inside each worker, adding a host-device copy to every weight sync.",
+        UserWarning,
+    )
+
+
 def _make_meta_params(param):
     is_param = isinstance(param, Parameter)
 

torchrl/weight_update/_mp.py

@@ -301,6 +301,8 @@ def send(
         Note: If sync=True (default), this is a blocking call that ensures
         specified workers are updated before returning.
         """
+        from torchrl.collectors.utils import _stage_unshareable_weights_on_cpu
+
         if not self.initialized_on_sender:
             raise RuntimeError("Must be initialized on sender before sending weights")
         if not self.synchronized_on_sender:
@@ -316,6 +318,7 @@ def send(
             strategy=self._strategy,
             context=context,
         )
+        prepared_weights = _stage_unshareable_weights_on_cpu(prepared_weights)
 
         transports = list(self._iterate_transports(worker_ids))
 

torchrl/weight_update/_shared.py

@@ -70,6 +70,14 @@ class SharedMemTransport:
 
     Both CPU and CUDA tensors maintain shared references when sent through mp.Queue.
 
+    .. note::
+        Sharing CUDA tensors relies on CUDA IPC, which is only available on
+        native Linux (not Windows or WSL2), and backends such as MPS cannot be
+        shared across processes at all. On those platforms the schemes stage
+        the weights through CPU shared memory instead (see
+        :meth:`SharedMemWeightSyncScheme._get_params_map`), and each worker
+        moves them back to its own device.
+
     """
 
     def __init__(self):
@@ -513,7 +521,11 @@ def _get_params_map(
     ):
         """Get the params_map for init_on_sender()."""
         # Import _cast locally to avoid circular imports
-        from torchrl.collectors.utils import _cast
+        from torchrl.collectors.utils import (
+            _cast,
+            _unshareable_devices,
+            _warn_cpu_staging,
+        )
 
         if params_map is not None:
             # Sanity check: params_map must be a dict[int, TensorDictBase]
@@ -576,6 +588,7 @@ def _get_params_map(
                     else:
                         # Distinct factories - create per-worker weights directly
                         params_map = {}
+                        staged_devices = []
                         for worker_idx, factory in enumerate(factories):
                             if factory is not None:
                                 worker_model = factory()
@@ -586,12 +599,19 @@ def _get_params_map(
                                     _cast, worker_weights
                                 )
                                 # Move to appropriate device
+                                device = None
                                 if devices and worker_idx < len(devices):
                                     device = devices[worker_idx]
-                                    if device is not None:
-                                        worker_weights = worker_weights.to(device)
+                                offending = _unshareable_devices(device, worker_weights)
+                                if offending:
+                                    staged_devices.extend(offending)
+                                    worker_weights = worker_weights.to("cpu")
+                                elif device is not None:
+                                    worker_weights = worker_weights.to(device)
                                 worker_weights = worker_weights.share_memory_()
                                 params_map[worker_idx] = worker_weights
+                        if staged_devices:
+                            _warn_cpu_staging(staged_devices)
                         # Set per_worker_weights flag on the scheme
                         self.per_worker_weights = True
                         return params_map
@@ -631,15 +651,26 @@ def _get_params_map(
             # Create device map with proper Parameter handling using _cast
             # _cast ensures Parameters stay as Parameters (with requires_grad=False)
             device_map = {}
+            staged_weights = None
+            staged_devices = []
             for d in devices_set:
-                if d != weights_device:
+                offending = _unshareable_devices(d, weights)
+                if offending:
+                    if staged_weights is None:
+                        staged_weights = weights.to("cpu").apply(_cast, weights)
+                        staged_weights.share_memory_()
+                    staged_devices.extend(offending)
+                    device_map[d] = staged_weights
+                elif d != weights_device:
                     # Move to device and apply _cast to preserve Parameter/Buffer types
                     weights_on_device = weights.to(d)
                     weights_on_device = weights_on_device.apply(_cast, weights)
                     device_map[d] = weights_on_device
                 else:
                     # Already on correct device, just apply _cast
                     device_map[d] = weights.apply(_cast, weights)
+            if staged_devices:
+                _warn_cpu_staging(staged_devices)
 
             # Create the map
             params_map = {