From 5e213a7b9d9c4a7b833ccc6e5a0875b29e385b63 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Sun, 18 Oct 2026 16:33:43 +0100 Subject: [PATCH 1/2] Honor `strict=False` when loading sharded (distributed) checkpoints `FSDPStrategy.load_checkpoint(..., strict=False)` (Fabric and Trainer) and the Fabric ModelParallel strategy only forwarded `strict` to `module.load_state_dict`, after `torch.distributed.checkpoint.load` had already raised `Missing key in checkpoint state_dict` for any model key absent from the sharded checkpoint. Pass a `DefaultLoadPlanner(allow_partial_load=not strict)` to the distributed checkpoint load so that missing keys are tolerated when `strict=False`, while `strict=True` keeps raising. Optimizer state loading is unchanged. Fixes #20274 Signed-off-by: Parth Vadera Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/lightning/fabric/strategies/fsdp.py | 7 +++--- .../fabric/strategies/model_parallel.py | 2 +- src/lightning/pytorch/strategies/fsdp.py | 2 +- tests/tests_fabric/strategies/test_fsdp.py | 22 ++++++++++++++++++- .../strategies/test_model_parallel.py | 4 ++-- 5 files changed, 29 insertions(+), 8 deletions(-) diff --git a/src/lightning/fabric/strategies/fsdp.py b/src/lightning/fabric/strategies/fsdp.py index ce913d361502e..44f9d71c93389 100644 --- a/src/lightning/fabric/strategies/fsdp.py +++ b/src/lightning/fabric/strategies/fsdp.py @@ -572,7 +572,7 @@ def load_checkpoint( with state_dict_ctx: module_state = {module_key: module.state_dict()} - _distributed_checkpoint_load(module_state, path) + _distributed_checkpoint_load(module_state, path, strict=strict) module.load_state_dict(module_state[module_key], strict=strict) if optimizers: @@ -958,7 +958,8 @@ def _distributed_checkpoint_save(converted_state: dict[str, Any], path: _PATH) - save(converted_state, checkpoint_id=path) -def _distributed_checkpoint_load(module_state: dict[str, Any], path: _PATH) -> None: +def _distributed_checkpoint_load(module_state: dict[str, Any], path: _PATH, strict: bool = True) -> None: from torch.distributed.checkpoint import load + from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner - load(module_state, checkpoint_id=path) + load(module_state, checkpoint_id=path, planner=DefaultLoadPlanner(allow_partial_load=not strict)) diff --git a/src/lightning/fabric/strategies/model_parallel.py b/src/lightning/fabric/strategies/model_parallel.py index 3c2dbf1df7868..9256f0f053b81 100644 --- a/src/lightning/fabric/strategies/model_parallel.py +++ b/src/lightning/fabric/strategies/model_parallel.py @@ -446,7 +446,7 @@ def _load_checkpoint( state_dict_options = StateDictOptions(cpu_offload=True) module_state = {module_key: get_model_state_dict(module)} - _distributed_checkpoint_load(module_state, path) + _distributed_checkpoint_load(module_state, path, strict=strict) module.load_state_dict(module_state[module_key], strict=strict) # the optimizer states must be loaded separately diff --git a/src/lightning/pytorch/strategies/fsdp.py b/src/lightning/pytorch/strategies/fsdp.py index f94b0973e5acb..7c945cb3e036d 100644 --- a/src/lightning/pytorch/strategies/fsdp.py +++ b/src/lightning/pytorch/strategies/fsdp.py @@ -613,7 +613,7 @@ def load_checkpoint(self, checkpoint_path: _PATH, weights_only: Optional[bool] = with state_dict_ctx: module_state = {"model": self.model.state_dict()} - _distributed_checkpoint_load(module_state, path) + _distributed_checkpoint_load(module_state, path, strict=self.lightning_module.strict_loading) self.model.load_state_dict(module_state["model"], strict=self.lightning_module.strict_loading) if self.lightning_module.trainer.state.fn == TrainerFn.FITTING and self.optimizers: diff --git a/tests/tests_fabric/strategies/test_fsdp.py b/tests/tests_fabric/strategies/test_fsdp.py index fac384ac8c3bf..961e322094cfb 100644 --- a/tests/tests_fabric/strategies/test_fsdp.py +++ b/tests/tests_fabric/strategies/test_fsdp.py @@ -647,7 +647,9 @@ def test_load_sharded_checkpoint_metadata_weights_only(monkeypatch): monkeypatch.setattr( "lightning.fabric.strategies.fsdp._get_sharded_state_dict_context", lambda module: mock.MagicMock() ) - monkeypatch.setattr("lightning.fabric.strategies.fsdp._distributed_checkpoint_load", lambda state, path: None) + monkeypatch.setattr( + "lightning.fabric.strategies.fsdp._distributed_checkpoint_load", lambda state, path, strict=True: None + ) captured = {} @@ -685,3 +687,21 @@ def test_get_distributed_checkpoint_reader_missing_fsspec_module(monkeypatch): monkeypatch.setitem(sys.modules, "torch.distributed.checkpoint._fsspec_filesystem", None) with pytest.raises(ImportError, match=r"Remote .fsspec. distributed checkpoints require"): _get_distributed_checkpoint_reader("memory:///w/ckpt") + + +def test_distributed_checkpoint_load_strict(tmp_path): + """`strict=False` must allow keys that are missing from the sharded checkpoint, `strict=True` must not.""" + from torch.distributed.checkpoint import save + from torch.distributed.checkpoint.api import CheckpointException + + from lightning.fabric.strategies.fsdp import _distributed_checkpoint_load + + save({"model": {"weight": torch.ones(2)}}, checkpoint_id=tmp_path) + + module_state = {"model": {"weight": torch.zeros(2), "extra": torch.zeros(2)}} + with pytest.raises(CheckpointException): + _distributed_checkpoint_load(module_state, tmp_path) + + _distributed_checkpoint_load(module_state, tmp_path, strict=False) + assert torch.equal(module_state["model"]["weight"], torch.ones(2)) + assert torch.equal(module_state["model"]["extra"], torch.zeros(2)) diff --git a/tests/tests_fabric/strategies/test_model_parallel.py b/tests/tests_fabric/strategies/test_model_parallel.py index 08a9b20b357f6..5dda6be6a2fff 100644 --- a/tests/tests_fabric/strategies/test_model_parallel.py +++ b/tests/tests_fabric/strategies/test_model_parallel.py @@ -385,7 +385,7 @@ def test_model_parallel_load_checkpoint_does_not_corrupt_remote_path(monkeypatch monkeypatch.setattr(mp, "_has_dtensor_modules", lambda m: isinstance(m, nn.Module)) monkeypatch.setattr(mp, "_is_sharded_checkpoint", lambda p: True) monkeypatch.setattr(mp, "_is_full_checkpoint", lambda p: False) - monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path: captured.update(path=path)) + monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path, strict=True: captured.update(path=path)) monkeypatch.setattr(mp, "_load", lambda path, weights_only=None: captured.update(meta=str(path)) or {}) monkeypatch.setattr("torch.distributed.checkpoint.state_dict.get_model_state_dict", lambda module: {}) @@ -420,7 +420,7 @@ def test_model_parallel_load_checkpoint_loads_non_tensor_metadata(monkeypatch, t assert _is_sharded_checkpoint(ckpt_dir) monkeypatch.setattr(mp, "_has_dtensor_modules", lambda m: isinstance(m, nn.Module)) - monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path: None) + monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path, strict=True: None) monkeypatch.setattr("torch.distributed.checkpoint.state_dict.get_model_state_dict", lambda module: {}) model = nn.Linear(2, 2) From 25024771c5927cf25a04c1e479b0bbcf1af9fdcb Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Fri, 30 Oct 2026 08:09:01 +0000 Subject: [PATCH 2/2] Address review feedback on #22028: update CHANGELOGs, test strict passthrough, and document optimizer state strictness Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera --- src/lightning/fabric/CHANGELOG.md | 2 ++ src/lightning/fabric/strategies/fsdp.py | 8 +++++++- src/lightning/fabric/strategies/model_parallel.py | 8 +++++++- src/lightning/pytorch/CHANGELOG.md | 2 ++ src/lightning/pytorch/strategies/fsdp.py | 7 +++++++ tests/tests_fabric/strategies/test_fsdp.py | 11 +++++++---- tests/tests_fabric/strategies/test_model_parallel.py | 11 +++++++++-- tests/tests_pytorch/strategies/test_fsdp.py | 3 +++ 8 files changed, 44 insertions(+), 8 deletions(-) diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md index de997cc4448e4..7002aa9166197 100644 --- a/src/lightning/fabric/CHANGELOG.md +++ b/src/lightning/fabric/CHANGELOG.md @@ -42,6 +42,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed `_atomic_save` swallowing `PermissionError`, which made checkpoint saving report success without writing a file ([#21799](https://github.com/Lightning-AI/pytorch-lightning/pull/21799)) +- Fixed `FSDPStrategy.load_checkpoint` and `ModelParallelStrategy.load_checkpoint` ignoring `strict=False` when loading sharded (distributed) checkpoints ([#22028](https://github.com/Lightning-AI/pytorch-lightning/pull/22028)) + - Fixed the `scheduler` type hint on `Fabric.setup` and `Strategy.setup_module_and_optimizers` to use the public `LRScheduler` instead of the private `_LRScheduler` ([#21934](https://github.com/Lightning-AI/pytorch-lightning/pull/21934)) --- diff --git a/src/lightning/fabric/strategies/fsdp.py b/src/lightning/fabric/strategies/fsdp.py index 44f9d71c93389..3e003d4c37389 100644 --- a/src/lightning/fabric/strategies/fsdp.py +++ b/src/lightning/fabric/strategies/fsdp.py @@ -527,7 +527,13 @@ def load_checkpoint( strict: bool = True, weights_only: Optional[bool] = None, ) -> dict[str, Any]: - """Load the contents from a checkpoint and restore the state of the given objects.""" + """Load the contents from a checkpoint and restore the state of the given objects. + + Note: + ``strict=False`` applies to module and metadata state loading. Loading sharded optimizer state + remains strict and requires the model parameters to match the checkpoint. + + """ if not state: raise ValueError( f"Got FSDPStrategy.load_checkpoint(..., state={state!r}) but a state with at least " diff --git a/src/lightning/fabric/strategies/model_parallel.py b/src/lightning/fabric/strategies/model_parallel.py index 9256f0f053b81..2e0e8a91b809b 100644 --- a/src/lightning/fabric/strategies/model_parallel.py +++ b/src/lightning/fabric/strategies/model_parallel.py @@ -282,7 +282,13 @@ def load_checkpoint( strict: bool = True, weights_only: Optional[bool] = None, ) -> dict[str, Any]: - """Load the contents from a checkpoint and restore the state of the given objects.""" + """Load the contents from a checkpoint and restore the state of the given objects. + + Note: + ``strict=False`` applies to module and metadata state loading. Loading sharded optimizer state + remains strict and requires the model parameters to match the checkpoint. + + """ if not state: raise ValueError( f"Got {type(self).__name__}.load_checkpoint(..., state={state!r}) but a state with at least " diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 1c73ea347751a..c386baa1c5e97 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -56,6 +56,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed `LightningModule.toggle_optimizer` / `untoggle_optimizer` breaking under `torch.compile` by disabling Dynamo tracing on these bookkeeping helpers ([#21513](https://github.com/Lightning-AI/pytorch-lightning/issues/21513)) +- Fixed `FSDPStrategy.load_checkpoint` ignoring `LightningModule.strict_loading=False` when loading sharded (distributed) checkpoints ([#22028](https://github.com/Lightning-AI/pytorch-lightning/pull/22028)) + - Fixed `CUDAAccelerator.setup_device` initializing CUDA on an unrelated device by calling `torch.cuda.set_device` before the matmul precision check ([#21726](https://github.com/Lightning-AI/pytorch-lightning/pull/21726)) --- diff --git a/src/lightning/pytorch/strategies/fsdp.py b/src/lightning/pytorch/strategies/fsdp.py index 7c945cb3e036d..36970c91d448b 100644 --- a/src/lightning/pytorch/strategies/fsdp.py +++ b/src/lightning/pytorch/strategies/fsdp.py @@ -598,6 +598,13 @@ def save_checkpoint( @override def load_checkpoint(self, checkpoint_path: _PATH, weights_only: Optional[bool] = None) -> dict[str, Any]: + """Load the contents from a checkpoint and restore the state of the given objects. + + Note: + ``LightningModule.strict_loading=False`` applies to module state loading. Loading sharded optimizer + state remains strict and requires the model parameters to match the checkpoint. + + """ # broadcast the path from rank 0 to ensure all the states are loaded from a common path path = _resolve_path(self.broadcast(checkpoint_path)) diff --git a/tests/tests_fabric/strategies/test_fsdp.py b/tests/tests_fabric/strategies/test_fsdp.py index 961e322094cfb..88bed1f0dd853 100644 --- a/tests/tests_fabric/strategies/test_fsdp.py +++ b/tests/tests_fabric/strategies/test_fsdp.py @@ -647,11 +647,12 @@ def test_load_sharded_checkpoint_metadata_weights_only(monkeypatch): monkeypatch.setattr( "lightning.fabric.strategies.fsdp._get_sharded_state_dict_context", lambda module: mock.MagicMock() ) - monkeypatch.setattr( - "lightning.fabric.strategies.fsdp._distributed_checkpoint_load", lambda state, path, strict=True: None - ) captured = {} + monkeypatch.setattr( + "lightning.fabric.strategies.fsdp._distributed_checkpoint_load", + lambda state, path, strict=True: captured.update(strict=strict), + ) def fake_load(path, weights_only=None): captured["weights_only"] = weights_only @@ -662,9 +663,11 @@ def fake_load(path, weights_only=None): model = nn.Linear(2, 2) strategy.load_checkpoint("memory:///x/sharded", state={"model": model}) assert captured["weights_only"] is False + assert captured["strict"] is True - strategy.load_checkpoint("memory:///x/sharded", state={"model": model}, weights_only=True) + strategy.load_checkpoint("memory:///x/sharded", state={"model": model}, weights_only=True, strict=False) assert captured["weights_only"] is True + assert captured["strict"] is False def test_get_distributed_checkpoint_writer_missing_fsspec_module(monkeypatch): diff --git a/tests/tests_fabric/strategies/test_model_parallel.py b/tests/tests_fabric/strategies/test_model_parallel.py index 5dda6be6a2fff..29952f6ac70d6 100644 --- a/tests/tests_fabric/strategies/test_model_parallel.py +++ b/tests/tests_fabric/strategies/test_model_parallel.py @@ -385,13 +385,16 @@ def test_model_parallel_load_checkpoint_does_not_corrupt_remote_path(monkeypatch monkeypatch.setattr(mp, "_has_dtensor_modules", lambda m: isinstance(m, nn.Module)) monkeypatch.setattr(mp, "_is_sharded_checkpoint", lambda p: True) monkeypatch.setattr(mp, "_is_full_checkpoint", lambda p: False) - monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path, strict=True: captured.update(path=path)) + monkeypatch.setattr( + mp, "_distributed_checkpoint_load", lambda state, path, strict=True: captured.update(path=path, strict=strict) + ) monkeypatch.setattr(mp, "_load", lambda path, weights_only=None: captured.update(meta=str(path)) or {}) monkeypatch.setattr("torch.distributed.checkpoint.state_dict.get_model_state_dict", lambda module: {}) model = nn.Linear(2, 2) mp._load_checkpoint(path="gs://bucket/run/ckpt", state={"model": model}, strict=False) assert captured["path"] == "gs://bucket/run/ckpt" + assert captured["strict"] is False assert captured["meta"] == "gs://bucket/run/ckpt/meta.pt" @@ -419,12 +422,16 @@ def test_model_parallel_load_checkpoint_loads_non_tensor_metadata(monkeypatch, t _atomic_save({"user_meta": _NonTensorMeta(42)}, ckpt_dir / _METADATA_FILENAME) assert _is_sharded_checkpoint(ckpt_dir) + captured = {} monkeypatch.setattr(mp, "_has_dtensor_modules", lambda m: isinstance(m, nn.Module)) - monkeypatch.setattr(mp, "_distributed_checkpoint_load", lambda state, path, strict=True: None) + monkeypatch.setattr( + mp, "_distributed_checkpoint_load", lambda state, path, strict=True: captured.update(strict=strict) + ) monkeypatch.setattr("torch.distributed.checkpoint.state_dict.get_model_state_dict", lambda module: {}) model = nn.Linear(2, 2) state = {"model": model, "user_meta": None} mp._load_checkpoint(path=ckpt_dir, state=state, strict=False) + assert captured["strict"] is False assert isinstance(state["user_meta"], _NonTensorMeta) assert state["user_meta"].value == 42 diff --git a/tests/tests_pytorch/strategies/test_fsdp.py b/tests/tests_pytorch/strategies/test_fsdp.py index 95aa3259e33c8..1d304d2fe6cd4 100644 --- a/tests/tests_pytorch/strategies/test_fsdp.py +++ b/tests/tests_pytorch/strategies/test_fsdp.py @@ -964,9 +964,12 @@ def test_load_sharded_checkpoint_metadata_weights_only(load_mock, _is_sharded_mo strategy.load_checkpoint(checkpoint_path="memory:///x/sharded") assert load_mock.call_args.kwargs["weights_only"] is False + assert _dist_load_mock.call_args.kwargs["strict"] is True + model.strict_loading = False strategy.load_checkpoint(checkpoint_path="memory:///x/sharded", weights_only=True) assert load_mock.call_args.kwargs["weights_only"] is True + assert _dist_load_mock.call_args.kwargs["strict"] is False @RunIf(min_cuda_gpus=2, skip_windows=True, standalone=True)