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)