From 95add982e6e530fa013f2f6e021ac2f50a333f8c Mon Sep 17 00:00:00 2001 From: zhixiangli Date: Fri, 17 Jul 2026 01:17:17 +0000 Subject: [PATCH 1/4] Allow FSDP strategy on CPU accelerator This enables using FSDP strategy with `accelerator="cpu"`, which is useful for CPU-based FSDP checkpoint benchmarking and testing. Also updated CHANGELOG.md and added tests. TAG=agy CONV=0b443e3e-6fbb-468a-9f11-b99f211cf34b --- src/lightning/pytorch/CHANGELOG.md | 2 + .../connectors/accelerator_connector.py | 10 +- tests/tests_pytorch/strategies/test_fsdp.py | 129 ++++++++++++++++-- .../connectors/test_accelerator_connector.py | 6 +- 4 files changed, 134 insertions(+), 13 deletions(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index f0201b14d9ec9..4ca2688758a22 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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 - diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py index f9d82708a9efd..a22932de16b0f 100644 --- a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py +++ b/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( diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 0610ce3b3b9e9..ce7311f0c012e 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/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,123 @@ 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) +@pytest.mark.parametrize("state_dict_type", ["full", "sharded"]) +def test_fsdp_cpu_trainable(state_dict_type, tmp_path): + """FSDP on CPU across 2 ranks trains and round-trips a checkpoint. + + Exercises both the ``full`` (single-file) and ``sharded`` (per-rank directory) checkpoint formats: 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=state_dict_type), + 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) + + # the two formats must produce genuinely different on-disk layouts + if state_dict_type == "sharded": + assert ckpt_path.is_dir() and _is_sharded_checkpoint(ckpt_path) + else: + assert ckpt_path.is_file() and not _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(): diff --git a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py b/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py index 826b818fbb9d5..d46b551f73747 100644 --- a/tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py +++ b/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 From 2b7f2f62988967cf8c2f3445a1457efcedcb349c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:19:45 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/tests_pytorch/strategies/test_fsdp.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index ce7311f0c012e..6548c8b3029b8 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -208,6 +208,7 @@ class _CPUFSDPTrainableModel(BoringModel): 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): @@ -262,6 +263,7 @@ def test_fsdp_cpu_trainable(state_dict_type, tmp_path): Exercises both the ``full`` (single-file) and ``sharded`` (per-rank directory) checkpoint formats: 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: @@ -299,9 +301,11 @@ def _make_trainer(max_epochs: int) -> Trainer: # the two formats must produce genuinely different on-disk layouts if state_dict_type == "sharded": - assert ckpt_path.is_dir() and _is_sharded_checkpoint(ckpt_path) + assert ckpt_path.is_dir() + assert _is_sharded_checkpoint(ckpt_path) else: - assert ckpt_path.is_file() and not _is_sharded_checkpoint(ckpt_path) + assert ckpt_path.is_file() + assert not _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())) From 3075d6225e093c194f090ab14ecfd1277495852e Mon Sep 17 00:00:00 2001 From: zhixiangli Date: Sun, 19 Jul 2026 11:17:07 +0000 Subject: [PATCH 3/4] Remove full state dict option from test_fsdp_cpu_trainable TAG=agy CONV=e53a2100-ddf2-4298-b123-6e1312841db7 --- tests/tests_pytorch/strategies/test_fsdp.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 6548c8b3029b8..8bfe5336d7add 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -256,11 +256,10 @@ def val_dataloader(self): @pytest.mark.filterwarnings("ignore::FutureWarning") @RunIf(standalone=True, skip_windows=True) -@pytest.mark.parametrize("state_dict_type", ["full", "sharded"]) -def test_fsdp_cpu_trainable(state_dict_type, tmp_path): +def test_fsdp_cpu_trainable(tmp_path): """FSDP on CPU across 2 ranks trains and round-trips a checkpoint. - Exercises both the ``full`` (single-file) and ``sharded`` (per-rank directory) checkpoint formats: the module is + 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. @@ -270,7 +269,7 @@ def _make_trainer(max_epochs: int) -> Trainer: return Trainer( accelerator="cpu", devices=2, - strategy=FSDPStrategy(state_dict_type=state_dict_type), + strategy=FSDPStrategy(state_dict_type="sharded"), max_epochs=max_epochs, default_root_dir=tmp_path, enable_checkpointing=False, @@ -299,13 +298,7 @@ def _make_trainer(max_epochs: int) -> Trainer: ckpt_path = Path(trainer.strategy.broadcast(str(tmp_path / "checkpoint"))) trainer.save_checkpoint(ckpt_path) - # the two formats must produce genuinely different on-disk layouts - if state_dict_type == "sharded": - assert ckpt_path.is_dir() - assert _is_sharded_checkpoint(ckpt_path) - else: - assert ckpt_path.is_file() - assert not _is_sharded_checkpoint(ckpt_path) + assert ckpt_path.is_dir() and _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())) From a3e129a771c7f7cf828b733e48facefe6e6b9d80 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 11:17:56 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/tests_pytorch/strategies/test_fsdp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 8bfe5336d7add..a2eb9b4304b8b 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -298,7 +298,8 @@ def _make_trainer(max_epochs: int) -> Trainer: ckpt_path = Path(trainer.strategy.broadcast(str(tmp_path / "checkpoint"))) trainer.save_checkpoint(ckpt_path) - assert ckpt_path.is_dir() and _is_sharded_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()))