feat(FSDP): enable FSDP strategy on CPU and add trainable checkpointing tests

src/lightning/pytorch/CHANGELOG.md

@@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added support for remote storage (fsspec URLs) when saving and loading distributed checkpoints with `ModelParallelStrategy` ([#21797](https://github.com/Lightning-AI/pytorch-lightning/issues/21797))
 
+- Added support for using the `FSDPStrategy` on CPU (`accelerator="cpu"`) ([#21812](https://github.com/Lightning-AI/pytorch-lightning/issues/21812))
+
 ### Changed
 
 -

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -33,6 +33,7 @@
 from lightning.fabric.utilities.imports import _IS_INTERACTIVE
 from lightning.pytorch.accelerators import AcceleratorRegistry
 from lightning.pytorch.accelerators.accelerator import Accelerator
+from lightning.pytorch.accelerators.cpu import CPUAccelerator
 from lightning.pytorch.accelerators.cuda import CUDAAccelerator
 from lightning.pytorch.accelerators.mps import MPSAccelerator
 from lightning.pytorch.accelerators.xla import XLAAccelerator
@@ -430,11 +431,14 @@ def _check_strategy_and_fallback(self) -> None:
 
         if (
             strategy_flag in FSDPStrategy.get_registered_strategies() or type(self._strategy_flag) is FSDPStrategy
-        ) and not (self._accelerator_flag in ("cuda", "gpu") or isinstance(self._accelerator_flag, CUDAAccelerator)):
+        ) and not (
+            self._accelerator_flag in ("cuda", "gpu", "cpu")
+            or isinstance(self._accelerator_flag, (CUDAAccelerator, CPUAccelerator))
+        ):
             raise ValueError(
-                f"The strategy `{FSDPStrategy.strategy_name}` requires a GPU accelerator, but received "
+                f"The strategy `{FSDPStrategy.strategy_name}` requires a GPU or CPU accelerator, but received "
                 f"`accelerator={self._accelerator_flag!r}`. Please set `accelerator='cuda'`, `accelerator='gpu'`,"
-                " or pass a `CUDAAccelerator()` instance to use FSDP."
+                " `accelerator='cpu'`, or pass a `CUDAAccelerator()`/`CPUAccelerator()` instance to use FSDP."
             )
         if strategy_flag in _DDP_FORK_ALIASES and "fork" not in torch.multiprocessing.get_all_start_methods():
             raise ValueError(

tests/tests_pytorch/strategies/test_fsdp.py

@@ -14,15 +14,17 @@
 import torch.nn as nn
 from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel, MixedPrecision
 from torch.distributed.fsdp.wrap import ModuleWrapPolicy, always_wrap_policy, size_based_auto_wrap_policy, wrap
+from torch.utils.data import DataLoader
 from torchmetrics import Accuracy
 
 from lightning.fabric.plugins.environments import LightningEnvironment
 from lightning.fabric.strategies.fsdp import _is_sharded_checkpoint
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_2, _TORCH_GREATER_EQUAL_2_3
 from lightning.fabric.utilities.load import _load_distributed_checkpoint
-from lightning.pytorch import Trainer
+from lightning.pytorch import Trainer, seed_everything
+from lightning.pytorch.accelerators.cpu import CPUAccelerator
 from lightning.pytorch.callbacks import ModelCheckpoint
-from lightning.pytorch.demos.boring_classes import BoringModel
+from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
 from lightning.pytorch.plugins import HalfPrecision
 from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
 from lightning.pytorch.strategies import FSDPStrategy
@@ -194,12 +196,121 @@ def _assert_save_equality(trainer, ckpt_path, cls=TestFSDPModel):
             assert torch.equal(ddp_param, shard_param)
 
 
-def test_invalid_on_cpu(tmp_path, cuda_count_0):
-    """Test to ensure that we raise Misconfiguration for FSDP on CPU."""
-    with pytest.raises(ValueError, match="The strategy `fsdp` requires a GPU accelerator"):
-        trainer = Trainer(accelerator="cpu", default_root_dir=tmp_path, fast_dev_run=True, strategy="fsdp")
-        assert isinstance(trainer.strategy, FSDPStrategy)
-        trainer.strategy.setup_environment()
+def test_fsdp_on_cpu_allowed(tmp_path, cuda_count_0):
+    """FSDP on CPU is allowed (enables CPU-based FSDP checkpointing)."""
+    trainer = Trainer(accelerator="cpu", default_root_dir=tmp_path, fast_dev_run=True, strategy="fsdp")
+    assert isinstance(trainer.strategy, FSDPStrategy)
+    assert isinstance(trainer.accelerator, CPUAccelerator)
+
+
+class _CPUFSDPTrainableModel(BoringModel):
+    """Overfits a fixed dataset so that the training loss must drop if FSDP actually trains the model.
+
+    If ``params_to_compare`` is given, the (already resharded) parameters are asserted to match it at the start of
+    training. This is used to verify that a checkpoint written after training reads back with identical weights.
+
+    """
+
+    def __init__(self, params_to_compare: Optional[list[torch.Tensor]] = None):
+        super().__init__()
+        self.layer: Optional[nn.Module] = None
+        self.losses: list[float] = []
+        self.params_to_compare = params_to_compare
+
+    def configure_model(self) -> None:
+        if self.layer is None:
+            self.layer = torch.nn.Sequential(torch.nn.Linear(32, 4), torch.nn.ReLU(), torch.nn.Linear(4, 2))
+        if isinstance(self.layer, FullyShardedDataParallel):
+            return
+        # wrap the two Linear layers and the container so their parameters are really flattened/sharded by FSDP
+        self.layer[0] = wrap(self.layer[0])
+        self.layer[2] = wrap(self.layer[2])
+        self.layer = wrap(self.layer)
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.layer.parameters(), lr=0.15)
+
+    def training_step(self, batch, batch_idx):
+        loss = self.step(batch)
+        self.losses.append(loss.detach().item())
+        return {"loss": loss}
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        # FSDP must be genuinely engaged here, not a silent fallback to plain CPU training
+        assert isinstance(self.layer, FullyShardedDataParallel)
+
+    def on_train_start(self):
+        # when resuming, the checkpoint has been loaded and resharded by now: the local shards must match exactly
+        if self.params_to_compare is None:
+            return
+        for expected, actual in zip(self.params_to_compare, self.trainer.model.parameters()):
+            torch.testing.assert_close(expected, actual, atol=0, rtol=0, equal_nan=True)
+
+    def train_dataloader(self):
+        # a small, fixed dataset the model can overfit, so the loss drop is deterministic and not flaky
+        return DataLoader(RandomDataset(32, 16), batch_size=8)
+
+    def val_dataloader(self):
+        return DataLoader(RandomDataset(32, 16), batch_size=8)
+
+
+@pytest.mark.filterwarnings("ignore::FutureWarning")
+@RunIf(standalone=True, skip_windows=True)
+def test_fsdp_cpu_trainable(tmp_path):
+    """FSDP on CPU across 2 ranks trains and round-trips a checkpoint.
+
+    Exercises the ``sharded`` (per-rank directory) checkpoint format: the module is
+    genuinely FSDP-wrapped, the loss drops sharply over 15 epochs, and the checkpoint written after training reads back
+    with identical sharded parameters.
+
+    """
+
+    def _make_trainer(max_epochs: int) -> Trainer:
+        return Trainer(
+            accelerator="cpu",
+            devices=2,
+            strategy=FSDPStrategy(state_dict_type="sharded"),
+            max_epochs=max_epochs,
+            default_root_dir=tmp_path,
+            enable_checkpointing=False,
+            logger=False,
+            enable_progress_bar=False,
+            enable_model_summary=False,
+        )
+
+    seed_everything(42)
+    model = _CPUFSDPTrainableModel()
+    trainer = _make_trainer(max_epochs=15)
+    trainer.fit(model)
+
+    # the model was really trained under a 2-rank sharding regime with FSDP engaged (also asserted live per step)
+    assert trainer.strategy.world_size == 2
+    assert isinstance(model.layer, FullyShardedDataParallel)
+
+    # trainable: the loss dropped by more than 10x between the start and the end of training
+    assert len(model.losses) > 5
+    first = sum(model.losses[:3]) / 3
+    last = sum(model.losses[-3:]) / 3
+    assert last < 0.1 * first, f"loss did not drop enough: first={first:.4f} last={last:.4f}"
+
+    # --- checkpoint write ---
+    # broadcast one path so every rank writes to (and later reads from) the same location
+    ckpt_path = Path(trainer.strategy.broadcast(str(tmp_path / "checkpoint")))
+    trainer.save_checkpoint(ckpt_path)
+
+    assert ckpt_path.is_dir()
+    assert _is_sharded_checkpoint(ckpt_path)
+
+    # snapshot the trained local shards on this rank for a read-back comparison
+    trained_params = deepcopy(list(trainer.model.parameters()))
+
+    # --- checkpoint read ---
+    # a fresh model + trainer must load the checkpoint and recover the exact trained parameters; the comparison runs in
+    # `on_train_start`, once the checkpoint has been loaded and resharded (one extra epoch guarantees the hook fires)
+    seed_everything(42)
+    reloaded_model = _CPUFSDPTrainableModel(params_to_compare=trained_params)
+    resumed_trainer = _make_trainer(max_epochs=15)
+    resumed_trainer.fit(reloaded_model, ckpt_path=ckpt_path)
 
 
 def test_custom_mixed_precision():

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -563,8 +563,10 @@ def test_strategy_choice_ddp_cpu_slurm(cuda_count_0, strategy):
 
 
 def test_check_fsdp_strategy_and_fallback():
-    with pytest.raises(ValueError, match="The strategy `fsdp` requires a GPU accelerator"):
-        Trainer(accelerator="cpu", strategy="fsdp")
+    # FSDP on CPU is now allowed (enables CPU-based FSDP checkpoint benchmarking).
+    trainer = Trainer(accelerator="cpu", strategy="fsdp")
+    assert isinstance(trainer.strategy, FSDPStrategy)
+    assert isinstance(trainer.accelerator, CPUAccelerator)
 
     class FSDPStrategySubclass(FSDPStrategy):
         pass