Honor strict=False when loading sharded (distributed) checkpoints

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))

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

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:

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))

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)