weights_only for fsdp load

Author: matsumotosan

src/lightning/fabric/strategies/model_parallel.py

@@ -411,6 +411,7 @@ def _load_checkpoint(
     state: dict[str, Union[Module, Optimizer, Any]],
     strict: bool = True,
     optimizer_states_from_list: bool = False,
+    weights_only: bool = False,
 ) -> dict[str, Any]:
     from torch.distributed.checkpoint.state_dict import (
         StateDictOptions,
@@ -449,7 +450,7 @@ def _load_checkpoint(
             set_optimizer_state_dict(module, optim, optim_state_dict=optim_state[optim_key], options=state_dict_options)
 
         # Load metadata (anything not a module or optimizer)
-        metadata = torch.load(path / _METADATA_FILENAME)
+        metadata = torch.load(path / _METADATA_FILENAME, weights_only=weights_only)
         requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
         _validate_keys_for_strict_loading(requested_metadata_keys, metadata.keys(), strict=strict)
         for key in requested_metadata_keys:
@@ -461,7 +462,7 @@ def _load_checkpoint(
         return metadata
 
     if _is_full_checkpoint(path):
-        checkpoint = torch.load(path, mmap=True, map_location="cpu", weights_only=False)
+        checkpoint = torch.load(path, mmap=True, map_location="cpu", weights_only=weights_only)
         _load_raw_module_state(checkpoint.pop(module_key), module, strict=strict)
 
         state_dict_options = StateDictOptions(

src/lightning/pytorch/strategies/model_parallel.py

@@ -329,7 +329,7 @@ def save_checkpoint(
             return super().save_checkpoint(checkpoint=checkpoint, filepath=path)
 
     @override
-    def load_checkpoint(self, checkpoint_path: _PATH) -> dict[str, Any]:
+    def load_checkpoint(self, checkpoint_path: _PATH, weights_only: bool = False) -> dict[str, Any]:
         # broadcast the path from rank 0 to ensure all the states are loaded from a common path
         path = Path(self.broadcast(checkpoint_path))
         state = {
@@ -342,6 +342,7 @@ def load_checkpoint(self, checkpoint_path: _PATH) -> dict[str, Any]:
             state=state,
             strict=self.lightning_module.strict_loading,
             optimizer_states_from_list=True,
+            weights_only=weights_only,
         )
 
     def _setup_distributed(self) -> None: