weights_only=True default for torch>=2.6

Author: matsumotosan

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

@@ -48,18 +48,18 @@ def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_optio
 
     @abstractmethod
     def load_checkpoint(
-        self, path: _PATH, map_location: Optional[Any] = None, weights_only: bool = True
+        self, path: _PATH, map_location: Optional[Any] = None, weights_only: Optional[bool] = None
     ) -> 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
+            weights_only: Defaults to ``None``. 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,14 +59,22 @@ 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, weights_only: bool = True
+        self,
+        path: _PATH,
+        map_location: Optional[Callable] = lambda storage, loc: storage,
+        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Loads checkpoint using :func:`torch.load`, with additional handling for ``fsspec`` remote loading of files.
 
         Args:
             path: Path to checkpoint
             map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage
                 locations.
+            weights_only: Defaults to ``None``. 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/utilities/cloud_io.py

@@ -17,7 +17,7 @@
 import io
 import logging
 from pathlib import Path
-from typing import IO, Any, Union
+from typing import IO, Any, Optional, Union
 
 import fsspec
 import fsspec.utils
@@ -26,6 +26,7 @@
 from fsspec.implementations.local import AbstractFileSystem
 from lightning_utilities.core.imports import module_available
 
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 
 log = logging.getLogger(__name__)
@@ -34,7 +35,7 @@
 def _load(
     path_or_url: Union[IO, _PATH],
     map_location: _MAP_LOCATION_TYPE = None,
-    weights_only: bool = True,
+    weights_only: Optional[bool] = None,
 ) -> Any:
     """Loads a checkpoint.
 
@@ -48,6 +49,11 @@ def _load(
             `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
 
     """
+    # default to `weights_only=True` for torch>=2.6
+    if weights_only is None and _TORCH_GREATER_EQUAL_2_6:
+        log.debug("Defaulting to `weights_only=True` for torch>=2.6.")
+        weights_only = True
+
     if not isinstance(path_or_url, (str, Path)):
         # any sort of BytesIO or similar
         return torch.load(

src/lightning/fabric/utilities/imports.py

@@ -35,6 +35,7 @@
 _TORCH_GREATER_EQUAL_2_4 = compare_version("torch", operator.ge, "2.4.0")
 _TORCH_GREATER_EQUAL_2_4_1 = compare_version("torch", operator.ge, "2.4.1")
 _TORCH_GREATER_EQUAL_2_5 = compare_version("torch", operator.ge, "2.5.0")
+_TORCH_GREATER_EQUAL_2_6 = compare_version("torch", operator.ge, "2.6.0")
 _TORCH_LESS_EQUAL_2_6 = compare_version("torch", operator.le, "2.6.0")
 _TORCHMETRICS_GREATER_EQUAL_1_0_0 = compare_version("torchmetrics", operator.ge, "1.0.0")
 _PYTHON_GREATER_EQUAL_3_10_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 10)

src/lightning/pytorch/core/saving.py

@@ -61,10 +61,6 @@ def _load_from_checkpoint(
 ) -> 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, weights_only=weights_only)
 

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -46,12 +46,7 @@ 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]
 
-        # 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)
+        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24)
         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)

tests/tests_pytorch/models/test_hparams.py

@@ -94,7 +94,7 @@ def __init__(self, hparams, *my_args, **my_kwargs):
 # -------------------------
 # STANDARD TESTS
 # -------------------------
-def _run_standard_hparams_test(tmp_path, model, cls, datamodule=None, try_overwrite=False, weights_only=True):
+def _run_standard_hparams_test(tmp_path, model, cls, datamodule=None, try_overwrite=False):
     """Tests for the existence of an arg 'test_arg=14'."""
     obj = datamodule if issubclass(cls, LightningDataModule) else model
 
@@ -108,20 +108,20 @@ def _run_standard_hparams_test(tmp_path, model, cls, datamodule=None, try_overwr
 
     # make sure the raw checkpoint saved the properties
     raw_checkpoint_path = _raw_checkpoint_path(trainer)
-    raw_checkpoint = torch.load(raw_checkpoint_path, weights_only=weights_only)
+    raw_checkpoint = torch.load(raw_checkpoint_path)
 
     assert cls.CHECKPOINT_HYPER_PARAMS_KEY in raw_checkpoint
     assert raw_checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY]["test_arg"] == 14
 
     # verify that model loads correctly
-    obj2 = cls.load_from_checkpoint(raw_checkpoint_path, weights_only=weights_only)
+    obj2 = cls.load_from_checkpoint(raw_checkpoint_path)
     assert obj2.hparams.test_arg == 14
 
     assert isinstance(obj2.hparams, hparam_type)
 
     if try_overwrite:
         # verify that we can overwrite the property
-        obj3 = cls.load_from_checkpoint(raw_checkpoint_path, test_arg=78, weights_only=weights_only)
+        obj3 = cls.load_from_checkpoint(raw_checkpoint_path, test_arg=78)
         assert obj3.hparams.test_arg == 78
 
     return raw_checkpoint_path
@@ -176,7 +176,7 @@ def test_omega_conf_hparams(tmp_path, cls):
     assert isinstance(obj.hparams, Container)
 
     # run standard test suite
-    raw_checkpoint_path = _run_standard_hparams_test(tmp_path, model, cls, datamodule=datamodule, weights_only=False)
+    raw_checkpoint_path = _run_standard_hparams_test(tmp_path, model, cls, datamodule=datamodule)
     obj2 = cls.load_from_checkpoint(raw_checkpoint_path, weights_only=False)
 
     assert isinstance(obj2.hparams, Container)