Exposeweights_only for loading checkpoints with Trainer, LightningModule, LightningDataModule

src/lightning/fabric/plugins/io/checkpoint_io.py

@@ -47,13 +47,20 @@ def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_optio
         """
 
     @abstractmethod
-    def load_checkpoint(self, path: _PATH, map_location: Optional[Any] = None) -> dict[str, Any]:
+    def load_checkpoint(
+        self, path: _PATH, map_location: Optional[Any] = None, weights_only: bool = True
+    ) -> dict[str, Any]:
         """Load checkpoint from a path when resuming or loading ckpt for test/validate/predict stages.
 
         Args:
             path: Path to checkpoint
             map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage
                 locations.
+            weights_only: If ``True``, restricts loading to ``state_dicts`` of plain ``torch.Tensor`` and other
+            primitive types. If loading a checkpoint from a trusted source that contains an ``nn.Module``, use
+                ``weights_only=False``. If loading checkpoint from an untrusted source, we recommend using
+                ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
 
         Returns: The loaded checkpoint.
 

src/lightning/fabric/plugins/io/torch_io.py

@@ -59,7 +59,7 @@ def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_optio
 
     @override
     def load_checkpoint(
-        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage
+        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage, weights_only: bool = True
     ) -> dict[str, Any]:
         """Loads checkpoint using :func:`torch.load`, with additional handling for ``fsspec`` remote loading of files.
 
@@ -80,7 +80,7 @@ def load_checkpoint(
         if not fs.exists(path):
             raise FileNotFoundError(f"Checkpoint file not found: {path}")
 
-        return pl_load(path, map_location=map_location)
+        return pl_load(path, map_location=map_location, weights_only=weights_only)
 
     @override
     def remove_checkpoint(self, path: _PATH) -> None:

src/lightning/fabric/utilities/cloud_io.py

@@ -34,13 +34,18 @@
 def _load(
     path_or_url: Union[IO, _PATH],
     map_location: _MAP_LOCATION_TYPE = None,
-    weights_only: bool = False,
+    weights_only: bool = True,
 ) -> Any:
     """Loads a checkpoint.
 
     Args:
         path_or_url: Path or URL of the checkpoint.
         map_location: a function, ``torch.device``, string or a dict specifying how to remap storage locations.
+        weights_only: If ``True``, restricts loading to ``state_dicts`` of plain ``torch.Tensor`` and other primitive
+            types. If loading a checkpoint from a trusted source that contains an ``nn.Module``, use
+            ``weights_only=False``. If loading checkpoint from an untrusted source, we recommend using
+            ``weights_only=True``. For more information, please refer to the
+            `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`__.
 
     """
     if not isinstance(path_or_url, (str, Path)):
@@ -70,7 +75,7 @@ def get_filesystem(path: _PATH, **kwargs: Any) -> AbstractFileSystem:
     return fs
 
 
-def _atomic_save(checkpoint: dict[str, Any], filepath: Union[str, Path]) -> None:
+def _atomic_save(checkpoint: dict[str, Any], filepath: _PATH) -> None:
     """Saves a checkpoint atomically, avoiding the creation of incomplete checkpoints.
 
     Args:

src/lightning/pytorch/core/module.py

@@ -1690,6 +1690,7 @@ def load_from_checkpoint(
         map_location: _MAP_LOCATION_TYPE = None,
         hparams_file: Optional[_PATH] = None,
         strict: Optional[bool] = None,
+        weights_only: Optional[bool] = None,
         **kwargs: Any,
     ) -> Self:
         r"""Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint it stores the arguments
@@ -1778,6 +1779,7 @@ def load_from_checkpoint(
             map_location,
             hparams_file,
             strict,
+            weights_only,
             **kwargs,
         )
         return cast(Self, loaded)

src/lightning/pytorch/core/saving.py

@@ -56,11 +56,17 @@ def _load_from_checkpoint(
     map_location: _MAP_LOCATION_TYPE = None,
     hparams_file: Optional[_PATH] = None,
     strict: Optional[bool] = None,
+    weights_only: Optional[bool] = None,
     **kwargs: Any,
 ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
     map_location = map_location or _default_map_location
+
+    if weights_only is None:
+        log.debug("`weights_only` not specified, defaulting to `True`.")
+        weights_only = True
+
     with pl_legacy_patch():
-        checkpoint = pl_load(checkpoint_path, map_location=map_location)
+        checkpoint = pl_load(checkpoint_path, map_location=map_location, weights_only=weights_only)
 
     # convert legacy checkpoints to the new format
     checkpoint = _pl_migrate_checkpoint(

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -414,7 +414,9 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
         """Creating a model checkpoint dictionary object from various component states.
 
         Args:
-            weights_only: saving model weights only
+            weights_only: If True, only saves model and loops state_dict objects. If False,
+            additionally saves callbacks, optimizers, schedulers, and precision plugin states.
+
         Return:
             structured dictionary: {
                 'epoch':                     training epoch

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+from packaging.version import Version
 
 import lightning.pytorch as pl
 from lightning.pytorch import Callback, Trainer
@@ -45,7 +46,12 @@ def test_load_legacy_checkpoints(tmp_path, pl_version: str):
         assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
         path_ckpt = path_ckpts[-1]
 
-        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24)
+        # legacy load utility added in 1.5.0 (see https://github.com/Lightning-AI/pytorch-lightning/pull/9166)
+        if pl_version == "local":
+            pl_version = pl.__version__
+        weights_only = not Version(pl_version) < Version("1.5.0")
+
+        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24, weights_only=weights_only)
         trainer = Trainer(default_root_dir=tmp_path)
         dm = ClassifDataModule(num_features=24, length=6000, batch_size=128, n_clusters_per_class=2, n_informative=8)
         res = trainer.test(model, datamodule=dm)
@@ -73,13 +79,18 @@ def test_legacy_ckpt_threading(pl_version: str):
     assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
     path_ckpt = path_ckpts[-1]
 
+    # legacy load utility added in 1.5.0 (see https://github.com/Lightning-AI/pytorch-lightning/pull/9166)
+    if pl_version == "local":
+        pl_version = pl.__version__
+    weights_only = not Version(pl_version) < Version("1.5.0")
+
     def load_model():
         import torch
 
         from lightning.pytorch.utilities.migration import pl_legacy_patch
 
         with pl_legacy_patch():
-            _ = torch.load(path_ckpt, weights_only=False)
+            _ = torch.load(path_ckpt, weights_only=weights_only)
 
     with patch("sys.path", [PATH_LEGACY] + sys.path):
         t1 = ThreadExceptionHandler(target=load_model)

tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py

@@ -32,7 +32,9 @@ class CustomCheckpointIO(CheckpointIO):
     def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
         torch.save(checkpoint, path)
 
-    def load_checkpoint(self, path: _PATH, storage_options: Optional[Any] = None) -> dict[str, Any]:
+    def load_checkpoint(
+        self, path: _PATH, storage_options: Optional[Any] = None, weights_only: bool = True
+    ) -> dict[str, Any]:
         return torch.load(path, weights_only=True)
 
     def remove_checkpoint(self, path: _PATH) -> None: