fix(FSDP): pass explicit CPU device to satisfy torch>=2.5 device_id contract (#21774)

* fix(strategy): pass CPU device to FSDP device_id on CPU

Passing `device_id=None` (which is `root_device.index` for CPU) triggers a guard in PyTorch >= 2.5 ("FSDP needs a non-CPU accelerator device"). Passing the actual `torch.device("cpu")` avoids this guard and allows FSDP to run on CPU.

- Update `FSDPStrategy` in both Fabric and PyTorch to pass `self.root_device` if it is a CPU device, otherwise `self.root_device.index`.
- Add unit tests to verify `device_id` selection for CPU and GPU devices.

CONV=516b4463-4a64-48ac-aabe-20fc589f838a
TAG=agy

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* chore: update CHANGELOG

* refactor(FSDP): address review feedback on device_id CPU fix

- common out `device_id` into a local so the debug log reflects the
  device actually passed (was logging root_device.index = None on CPU)
- add call-site comment noting CPU is not a supported FSDP path; the
  branch only honors the torch>=2.5 device_id contract
- replace the parametrized CPU/GPU device_id tests with a single
  minimal CPU-only assertion per strategy: guards the device_id=None
  regression without codifying the GPU path as a tested contract
- reword CHANGELOG to frame this as a narrow robustness fix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* apply copilot suggestion

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deependu <deependujha21@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 4 people

src/lightning/fabric/CHANGELOG.md

@@ -20,6 +20,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `FSDPStrategy` raising `RuntimeError` under PyTorch 2.5+ when `root_device` is CPU, by passing an explicit `torch.device("cpu")` instead of `device_id=None` (relevant only when the GPU-accelerator guard is bypassed) ([#21774](https://github.com/Lightning-AI/pytorch-lightning/pull/21774))
+
 - Fixed inconsistent FLOPs reporting on NVIDIA H100/H200 GPUs by defaulting to dense FLOPs, with sparse FLOPs now requiring an explicit opt-in. ([#21743](https://github.com/Lightning-AI/pytorch-lightning/pull/21743))
 
 - Fixed AccumulateGrad stream mismatch warning when using DDP with Fabric ([#21746](https://github.com/Lightning-AI/pytorch-lightning/pull/21746))

src/lightning/fabric/strategies/fsdp.py

@@ -296,12 +296,15 @@ def setup_module(self, module: Module) -> Module:
                 del self._fsdp_kwargs["auto_wrap_policy"]
         else:
             _warn_if_shared_params_across_fsdp_units(module, self._fsdp_kwargs.get("auto_wrap_policy"))
+            # CPU is not a supported FSDP target; this branch only honors the torch>=2.5 contract,
+            # which rejects device_id=None (root_device.index is None on CPU). The GPU path is unchanged.
+            device_id = self.root_device if self.root_device.type == "cpu" else self.root_device.index
             module = FullyShardedDataParallel(
                 module=module,
                 cpu_offload=self.cpu_offload,
                 mixed_precision=self.mixed_precision_config,
                 sharding_strategy=self.sharding_strategy,
-                device_id=self.root_device.index,
+                device_id=device_id,
                 **self._fsdp_kwargs,
             )
 
@@ -354,12 +357,15 @@ def module_sharded_context(self) -> AbstractContextManager:
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
         from torch.distributed.fsdp.wrap import enable_wrap
 
+        # CPU is not a supported FSDP target; this branch only honors the torch>=2.5 contract,
+        # which rejects device_id=None (root_device.index is None on CPU). The GPU path is unchanged.
+        device_id = self.root_device if self.root_device.type == "cpu" else self.root_device.index
         return enable_wrap(
             wrapper_cls=FullyShardedDataParallel,
             cpu_offload=self.cpu_offload,
             mixed_precision=self.mixed_precision_config,
             sharding_strategy=self.sharding_strategy,
-            device_id=self.root_device.index,
+            device_id=device_id,
             **self._fsdp_kwargs,
         )
 

src/lightning/pytorch/CHANGELOG.md

@@ -24,6 +24,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `FSDPStrategy` raising `RuntimeError` under PyTorch 2.5+ when `root_device` is CPU, by passing an explicit `torch.device("cpu")` instead of `device_id=None` (relevant only when the GPU-accelerator guard is bypassed) ([#21774](https://github.com/Lightning-AI/pytorch-lightning/pull/21774))
+
 - Fixed non-zero process exits in `CombinedLoader.reset()` with large tensors and persistent spawned workers by avoiding explicit `_shutdown_workers()` calls and relying on iterator cleanup via `del` [#21708](https://github.com/Lightning-AI/pytorch-lightning/issues/21708)
 
 - Fixed `SIGTERMException` producing a zero exit code instead of 143 (128 + SIGTERM) ([#21623](https://github.com/Lightning-AI/pytorch-lightning/issues/21623))

src/lightning/pytorch/strategies/fsdp.py

@@ -308,13 +308,16 @@ def _setup_model(self, model: Module) -> Module:
                 del self.kwargs["auto_wrap_policy"]
         else:
             _warn_if_shared_params_across_fsdp_units(model, self.kwargs.get("auto_wrap_policy"))
-            log.debug(f"setting up FSDP model with device id: {self.root_device.index}, kwargs: {self.kwargs}")
+            # CPU is not a supported FSDP target; this branch only honors the torch>=2.5 contract,
+            # which rejects device_id=None (root_device.index is None on CPU). The GPU path is unchanged.
+            device_id = self.root_device if self.root_device.type == "cpu" else self.root_device.index
+            log.debug(f"setting up FSDP model with device id: {device_id}, kwargs: {self.kwargs}")
             model = FullyShardedDataParallel(
                 module=model,
                 cpu_offload=self.cpu_offload,
                 mixed_precision=self.mixed_precision_config,
                 sharding_strategy=self.sharding_strategy,
-                device_id=self.root_device.index,
+                device_id=device_id,
                 **self.kwargs,
             )
 
@@ -403,12 +406,15 @@ def model_sharded_context(self) -> Generator[None, None, None]:
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
         from torch.distributed.fsdp.wrap import enable_wrap
 
+        # CPU is not a supported FSDP target; this branch only honors the torch>=2.5 contract,
+        # which rejects device_id=None (root_device.index is None on CPU). The GPU path is unchanged.
+        device_id = self.root_device if self.root_device.type == "cpu" else self.root_device.index
         with enable_wrap(
             wrapper_cls=FullyShardedDataParallel,
             cpu_offload=self.cpu_offload,
             mixed_precision=self.mixed_precision_config,
             sharding_strategy=self.sharding_strategy,
-            device_id=self.root_device.index,
+            device_id=device_id,
             **self.kwargs,
         ):
             yield

tests/tests_fabric/strategies/test_fsdp.py

@@ -188,6 +188,46 @@ def __init__(self):
     apply_mock.assert_called_with(wrapped, checkpoint_wrapper_fn=ANY, **strategy._activation_checkpointing_kwargs)
 
 
+def test_setup_module_device_id_cpu():
+    """``setup_module`` passes an explicit ``torch.device('cpu')`` (not ``device_id=None``) on CPU.
+
+    ``root_device.index`` is ``None`` on CPU; ``device_id=None`` trips torch>=2.5's "FSDP needs a
+    non-CPU accelerator device" guard. Only reachable when the GPU-accelerator guard is bypassed.
+
+    """
+    captured = {}
+
+    class FakeFSDP(nn.Module):
+        def __init__(self, module, **kwargs):
+            super().__init__()
+            captured.update(kwargs)
+            self.module = module
+
+    strategy = FSDPStrategy()
+    strategy._parallel_devices = [torch.device("cpu")]
+    with mock.patch("torch.distributed.fsdp.FullyShardedDataParallel", FakeFSDP):
+        strategy.setup_module(nn.Linear(2, 2))
+    assert captured["device_id"] == torch.device("cpu")
+
+
+def test_module_sharded_context_device_id_cpu():
+    """``module_sharded_context`` passes an explicit ``torch.device('cpu')`` (not ``device_id=None``) on CPU."""
+    from contextlib import contextmanager
+
+    captured = {}
+
+    @contextmanager
+    def fake_enable_wrap(*args, **kwargs):
+        captured.update(kwargs)
+        yield
+
+    strategy = FSDPStrategy()
+    strategy._parallel_devices = [torch.device("cpu")]
+    with mock.patch("torch.distributed.fsdp.wrap.enable_wrap", fake_enable_wrap), strategy.module_sharded_context():
+        pass
+    assert captured["device_id"] == torch.device("cpu")
+
+
 def test_forbidden_precision_raises():
     with pytest.raises(TypeError, match="can only work with the `FSDPPrecision"):
         FSDPStrategy(precision=HalfPrecision())

tests/tests_pytorch/strategies/test_fsdp.py

@@ -446,6 +446,49 @@ def __init__(self):
     apply_mock.assert_called_with(wrapped, checkpoint_wrapper_fn=ANY, **strategy._activation_checkpointing_kwargs)
 
 
+def test_setup_model_device_id_cpu():
+    """``_setup_model`` passes an explicit ``torch.device('cpu')`` (not ``device_id=None``) on CPU.
+
+    ``root_device.index`` is ``None`` on CPU; ``device_id=None`` trips torch>=2.5's "FSDP needs a
+    non-CPU accelerator device" guard. Only reachable when the GPU-accelerator guard is bypassed.
+
+    """
+    captured = {}
+
+    class FakeFSDP(nn.Module):
+        def __init__(self, module, **kwargs):
+            super().__init__()
+            captured.update(kwargs)
+            self.module = module
+
+    strategy = FSDPStrategy()
+    model = nn.Linear(2, 2)
+    strategy._parallel_devices = [torch.device("cpu")]
+    strategy._lightning_module = model
+    strategy._process_group = Mock()
+    with mock.patch("torch.distributed.fsdp.FullyShardedDataParallel", FakeFSDP):
+        strategy._setup_model(model)
+    assert captured["device_id"] == torch.device("cpu")
+
+
+def test_model_sharded_context_device_id_cpu():
+    """``model_sharded_context`` passes an explicit ``torch.device('cpu')`` (not ``device_id=None``) on CPU."""
+    from contextlib import contextmanager
+
+    captured = {}
+
+    @contextmanager
+    def fake_enable_wrap(*args, **kwargs):
+        captured.update(kwargs)
+        yield
+
+    strategy = FSDPStrategy()
+    strategy._parallel_devices = [torch.device("cpu")]
+    with mock.patch("torch.distributed.fsdp.wrap.enable_wrap", fake_enable_wrap), strategy.model_sharded_context():
+        pass
+    assert captured["device_id"] == torch.device("cpu")
+
+
 def test_strategy_cpu_offload():
     """Test the different ways cpu offloading can be enabled."""
     # bool