diff --git a/test/test_collectors.py b/test/test_collectors.py index 35c7e6e735b..6baaf799a46 100644 --- a/test/test_collectors.py +++ b/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") diff --git a/torchrl/collectors/utils.py b/torchrl/collectors/utils.py index f0b6e05adcd..721ec7450cd 100644 --- a/torchrl/collectors/utils.py +++ b/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) diff --git a/torchrl/weight_update/_mp.py b/torchrl/weight_update/_mp.py index 5c495ba9ac8..63483d25e59 100644 --- a/torchrl/weight_update/_mp.py +++ b/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)) diff --git a/torchrl/weight_update/_shared.py b/torchrl/weight_update/_shared.py index af3353a847a..e3102f6c2f9 100644 --- a/torchrl/weight_update/_shared.py +++ b/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,8 +651,17 @@ 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) @@ -640,6 +669,8 @@ def _get_params_map( 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 = {