pytorch-lightning 2.6.0__tar.gz → 2.6.0.dev0__tar.gz

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/.actions/assistant.py

@@ -46,14 +46,14 @@ _PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
 
 
 class _RequirementWithComment(Requirement):
-    strict_cmd = "strict"
+    strict_string = "# strict"
 
     def __init__(self, *args: Any, comment: str = "", pip_argument: Optional[str] = None, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.comment = comment
         assert pip_argument is None or pip_argument  # sanity check that it's not an empty str
         self.pip_argument = pip_argument
-        self.strict = self.strict_cmd in comment.lower()
+        self.strict = self.strict_string in comment.lower()
 
     def adjust(self, unfreeze: str) -> str:
         """Remove version restrictions unless they are strict.
@@ -62,26 +62,25 @@ class _RequirementWithComment(Requirement):
         'arrow<=1.2.2,>=1.2.0'
         >>> _RequirementWithComment("arrow<=1.2.2,>=1.2.0", comment="# strict").adjust("none")
         'arrow<=1.2.2,>=1.2.0  # strict'
-        >>> _RequirementWithComment('arrow<=1.2.2,>=1.2.0; python_version >= "3.10"', comment="# my name").adjust("all")
-        'arrow>=1.2.0; python_version >= "3.10"'
+        >>> _RequirementWithComment("arrow<=1.2.2,>=1.2.0", comment="# my name").adjust("all")
+        'arrow>=1.2.0'
         >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# strict").adjust("all")
         'arrow<=1.2.2,>=1.2.0  # strict'
-        >>> _RequirementWithComment('arrow; python_version >= "3.10"').adjust("all")
-        'arrow; python_version >= "3.10"'
+        >>> _RequirementWithComment("arrow").adjust("all")
+        'arrow'
         >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# cool").adjust("major")
         'arrow<2.0,>=1.2.0'
         >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# strict").adjust("major")
         'arrow<=1.2.2,>=1.2.0  # strict'
-        >>> _RequirementWithComment('arrow>=1.2.0; python_version >= "3.10"').adjust("major")
-        'arrow>=1.2.0; python_version >= "3.10"'
+        >>> _RequirementWithComment("arrow>=1.2.0").adjust("major")
+        'arrow>=1.2.0'
         >>> _RequirementWithComment("arrow").adjust("major")
         'arrow'
 
         """
         out = str(self)
         if self.strict:
-            return f"{out}  # {self.strict_cmd}"
-
+            return f"{out}  {self.strict_string}"
         specs = [(spec.operator, spec.version) for spec in self.specifier]
         if unfreeze == "major":
             for operator, version in specs:

{pytorch_lightning-2.6.0/src/pytorch_lightning.egg-info → pytorch_lightning-2.6.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytorch-lightning
-Version: 2.6.0
+Version: 2.6.0.dev0
 Summary: PyTorch Lightning is the lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.
 Home-page: https://github.com/Lightning-AI/lightning
 Download-URL: https://github.com/Lightning-AI/lightning
@@ -36,11 +36,6 @@ Requires-Dist: torchmetrics>0.7.0
 Requires-Dist: packaging>=20.0
 Requires-Dist: typing-extensions>4.5.0
 Requires-Dist: lightning-utilities>=0.10.0
-Provides-Extra: examples
-Requires-Dist: requests<2.33.0; extra == "examples"
-Requires-Dist: torchvision>=0.16.0; extra == "examples"
-Requires-Dist: ipython[all]>=8.0.0; extra == "examples"
-Requires-Dist: torchmetrics>=0.10.0; extra == "examples"
 Provides-Extra: extra
 Requires-Dist: matplotlib>3.1; extra == "extra"
 Requires-Dist: omegaconf>=2.2.3; extra == "extra"
@@ -49,11 +44,17 @@ Requires-Dist: jsonargparse[jsonnet,signatures]>=4.39.0; extra == "extra"
 Requires-Dist: rich>=12.3.0; extra == "extra"
 Requires-Dist: tensorboardX>=2.2; extra == "extra"
 Requires-Dist: bitsandbytes>=0.45.2; platform_system != "Darwin" and extra == "extra"
+Provides-Extra: strategies
+Requires-Dist: deepspeed<0.17.0,>=0.15.0; (platform_system != "Windows" and platform_system != "Darwin") and extra == "strategies"
+Provides-Extra: examples
+Requires-Dist: requests<2.33.0; extra == "examples"
+Requires-Dist: torchvision>=0.16.0; extra == "examples"
+Requires-Dist: ipython[all]>=8.0.0; extra == "examples"
+Requires-Dist: torchmetrics>=0.10.0; extra == "examples"
 Provides-Extra: test
-Requires-Dist: coverage==7.11.0; python_version >= "3.10" and extra == "test"
-Requires-Dist: coverage==7.10.7; python_version < "3.10" and extra == "test"
+Requires-Dist: coverage==7.10.7; extra == "test"
 Requires-Dist: pytest==8.4.2; extra == "test"
-Requires-Dist: pytest-cov==7.0.0; extra == "test"
+Requires-Dist: pytest-cov==6.3.0; extra == "test"
 Requires-Dist: pytest-timeout==2.4.0; extra == "test"
 Requires-Dist: pytest-rerunfailures==16.0.1; python_version < "3.10" and extra == "test"
 Requires-Dist: pytest-rerunfailures==16.1; python_version >= "3.10" and extra == "test"
@@ -64,15 +65,13 @@ Requires-Dist: numpy>1.20.0; extra == "test"
 Requires-Dist: onnx>1.12.0; extra == "test"
 Requires-Dist: onnxruntime>=1.12.0; extra == "test"
 Requires-Dist: onnxscript>=0.1.0; extra == "test"
-Requires-Dist: psutil<7.2.0; extra == "test"
+Requires-Dist: psutil<7.1.2; extra == "test"
 Requires-Dist: pandas>2.0; extra == "test"
 Requires-Dist: fastapi; extra == "test"
 Requires-Dist: uvicorn; extra == "test"
 Requires-Dist: tensorboard>=2.11; extra == "test"
 Requires-Dist: torch-tensorrt; (platform_system == "Linux" and python_version >= "3.12") and extra == "test"
 Requires-Dist: huggingface-hub; extra == "test"
-Provides-Extra: strategies
-Requires-Dist: deepspeed<0.17.0,>=0.15.0; (platform_system != "Windows" and platform_system != "Darwin") and extra == "strategies"
 Provides-Extra: deepspeed
 Requires-Dist: deepspeed<0.17.0,>=0.15.0; (platform_system != "Windows" and platform_system != "Darwin") and extra == "deepspeed"
 Provides-Extra: all
@@ -101,10 +100,9 @@ Requires-Dist: requests<2.33.0; extra == "dev"
 Requires-Dist: torchvision>=0.16.0; extra == "dev"
 Requires-Dist: ipython[all]>=8.0.0; extra == "dev"
 Requires-Dist: torchmetrics>=0.10.0; extra == "dev"
-Requires-Dist: coverage==7.11.0; python_version >= "3.10" and extra == "dev"
-Requires-Dist: coverage==7.10.7; python_version < "3.10" and extra == "dev"
+Requires-Dist: coverage==7.10.7; extra == "dev"
 Requires-Dist: pytest==8.4.2; extra == "dev"
-Requires-Dist: pytest-cov==7.0.0; extra == "dev"
+Requires-Dist: pytest-cov==6.3.0; extra == "dev"
 Requires-Dist: pytest-timeout==2.4.0; extra == "dev"
 Requires-Dist: pytest-rerunfailures==16.0.1; python_version < "3.10" and extra == "dev"
 Requires-Dist: pytest-rerunfailures==16.1; python_version >= "3.10" and extra == "dev"
@@ -115,7 +113,7 @@ Requires-Dist: numpy>1.20.0; extra == "dev"
 Requires-Dist: onnx>1.12.0; extra == "dev"
 Requires-Dist: onnxruntime>=1.12.0; extra == "dev"
 Requires-Dist: onnxscript>=0.1.0; extra == "dev"
-Requires-Dist: psutil<7.2.0; extra == "dev"
+Requires-Dist: psutil<7.1.2; extra == "dev"
 Requires-Dist: pandas>2.0; extra == "dev"
 Requires-Dist: fastapi; extra == "dev"
 Requires-Dist: uvicorn; extra == "dev"
@@ -166,7 +164,7 @@ ______________________________________________________________________
 [![DockerHub](https://img.shields.io/docker/pulls/pytorchlightning/pytorch_lightning.svg)](https://hub.docker.com/r/pytorchlightning/pytorch_lightning)
 [![codecov](https://codecov.io/gh/Lightning-AI/pytorch-lightning/graph/badge.svg?token=SmzX8mnKlA)](https://codecov.io/gh/Lightning-AI/pytorch-lightning)
 
-[![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=2.6.0)](https://lightning.ai/docs/pytorch/stable/)[![Discord](https://img.shields.io/discord/1077906959069626439?style=plastic)](https://discord.gg/VptPCZkGNa)
+[![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=2.6.0dev0)](https://lightning.ai/docs/pytorch/stable/)[![Discord](https://img.shields.io/discord/1077906959069626439?style=plastic)](https://discord.gg/VptPCZkGNa)
 [![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/Lightning-AI/pytorch-lightning/blob/master/LICENSE)
 
 <!--

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/pyproject.toml

@@ -92,6 +92,9 @@ ignore = [
 "src/**" = [
     "S101",  # todo: Use of `assert` detected
     "S113",  # todo: Probable use of requests call without timeout
+    "S301",  # todo: `pickle` and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue
+    "S324",  # todo: Probable use of insecure hash functions in `hashlib`
+    "S403",  # todo: `pickle`, `cPickle`, `dill`, and `shelve` modules are possibly insecure
     "S404",  # todo: `subprocess` module is possibly insecure
     "S602",  # todo: `subprocess` call with `shell=True` identified, security issue
     "S603",  # todo: `subprocess` call: check for execution of untrusted input
@@ -103,14 +106,18 @@ ignore = [
 "tests/**" = [
     "S101",  # Use of `assert` detected
     "S301",  # `pickle` and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue
+    "S113",  # todo: Probable use of requests call without timeout
     "S311",  # todo: Standard pseudo-random generators are not suitable for cryptographic purposes
     "S108",  # todo: Probable insecure usage of temporary file or directory: "/tmp/sys-customizations-sync"
+    "S202",  # Uses of `tarfile.extractall()`
     "S403",  # `pickle`, `cPickle`, `dill`, and `shelve` modules are possibly insecure
     "S404",  # `subprocess` module is possibly insecure
     "S602",  # todo: `subprocess` call with `shell=True` identified, security issue
     "S603",  # todo: `subprocess` call: check for execution of untrusted input
     "S605",  # todo: Starting a process with a shell: seems safe, but may be changed in the future; consider rewriting without `shell`
     "S607",  # todo: Starting a process with a partial executable path
+    "RET504",  # todo:Unnecessary variable assignment before `return` statement
+    "PT004",  # todo: Fixture `tmpdir_unittest_fixture` does not return anything, add leading underscore
     "PT012",  # todo: `pytest.raises()` block should contain a single simple statement
     "PT019",  # todo: Fixture `_` without value is injected as parameter, use `@pytest.mark.usefixtures` instead
 ]
@@ -172,7 +179,6 @@ filterwarnings = [
     # "error::DeprecationWarning",
     "error::FutureWarning",
     "ignore::FutureWarning:onnxscript", # Temporary ignore until onnxscript is updated
-    "ignore:You are using `torch.load` with `weights_only=False`.*:FutureWarning",
 ]
 xfail_strict = true
 junit_duration_report = "call"

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/requirements/fabric/test.txt

@@ -1,8 +1,7 @@
-coverage ==7.11.0; python_version >= "3.10"
-coverage ==7.10.7; python_version < "3.10"
+coverage ==7.10.7
 numpy >=1.21.0, <1.27.0
 pytest ==8.4.2
-pytest-cov ==7.0.0
+pytest-cov ==6.3.0
 pytest-timeout ==2.4.0
 pytest-rerunfailures ==16.0.1; python_version < "3.10"
 pytest-rerunfailures ==16.1; python_version >= "3.10"

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/requirements/pytorch/docs.txt

@@ -1,9 +1,9 @@
 -r ../docs.txt
 
 nbformat  # used for generate empty notebook
-ipython[notebook] <9.8.0
+ipython[notebook] <9.7.0
 setuptools<81.0  # workaround for `error in ipython setup command: use_2to3 is invalid.`
 
-onnxscript >= 0.2.2, < 0.6.0
+onnxscript >= 0.2.2, < 0.5.0
 
 #-r ../../_notebooks/.actions/requires.txt

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/requirements/pytorch/extra.txt

@@ -5,7 +5,7 @@
 matplotlib>3.1, <3.11.0
 omegaconf >=2.2.3, <2.4.0
 hydra-core >=1.2.0, <1.4.0
-jsonargparse[signatures,jsonnet] >=4.39.0, <4.44.0
-rich >=12.3.0, <14.3.0
+jsonargparse[signatures,jsonnet] >=4.39.0, <4.43.0
+rich >=12.3.0, <14.2.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.45.2,<0.47.0; platform_system != "Darwin"

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/requirements/pytorch/test.txt

@@ -1,7 +1,6 @@
-coverage ==7.11.0; python_version >= "3.10"
-coverage ==7.10.7; python_version < "3.10"
+coverage ==7.10.7
 pytest ==8.4.2
-pytest-cov ==7.0.0
+pytest-cov ==6.3.0
 pytest-timeout ==2.4.0
 pytest-rerunfailures ==16.0.1; python_version < "3.10"
 pytest-rerunfailures ==16.1; python_version >= "3.10"
@@ -13,8 +12,8 @@ scikit-learn >0.22.1, <1.8.0
 numpy >1.20.0, <1.27.0
 onnx >1.12.0, <1.20.0
 onnxruntime >=1.12.0, <1.24.0
-onnxscript >= 0.1.0, < 0.6.0
-psutil <7.2.0 # for `DeviceStatsMonitor`
+onnxscript >= 0.1.0, < 0.5.0
+psutil <7.1.2 # for `DeviceStatsMonitor`
 pandas >2.0, <2.4.0  # needed in benchmarks
 fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in App
 uvicorn  # for `ServableModuleValidator`  # not setting version as re-defined in App

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/CHANGELOG.md

@@ -5,42 +5,29 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
-## [Unreleased] - YYYY-MM-DD
+## [unreleased] - YYYY-MM-DD
 
 ### Added
 
 -
 
-### Changed
-
--
 
 ### Removed
 
 -
 
 
-## [2.6.0] - 2025-11-28
-
-### Added
-
-- Added kwargs-filtering for `Fabric.call` to support different callback method signatures ([#21258](https://github.com/Lightning-AI/pytorch-lightning/pull/21258))
+### Changed
 
+-
 
-### Changed
 
-- Expose `weights_only` argument for `Trainer.{fit,validate,test,predict}` and let `torch` handle default value ([#21072](https://github.com/Lightning-AI/pytorch-lightning/pull/21072))
 - Set `_DeviceDtypeModuleMixin._device` from torch's default device function ([#21164](https://github.com/Lightning-AI/pytorch-lightning/pull/21164))
 
 
 ### Fixed
 
-- Fixed issue in detecting MPIEnvironment with partial mpi4py installation ([#21353](https://github.com/Lightning-AI/pytorch-lightning/pull/21353))
-
-- Learning rate scheduler is stepped at the end of epoch when `on_train_batch_start` returns -1 ([#21296](https://github.com/Lightning-AI/pytorch-lightning/issues/21296)).
-
-
-- Fixed FSDP mixed precision semantics and added user warning ([#21361](https://github.com/Lightning-AI/pytorch-lightning/pull/21361))
+- Fixed `EADDRINUSE` errors in distributed tests with port manager and retry logic ([#21309](https://github.com/Lightning-AI/pytorch-lightning/pull/21309))
 
 
 ---

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/fabric.py

@@ -985,34 +985,6 @@ class Fabric:
             )
         return self._wrap_and_launch(function, self, *args, **kwargs)
 
-    def _filter_kwargs_for_callback(self, method: Callable, kwargs: dict[str, Any]) -> dict[str, Any]:
-        """Filter keyword arguments to only include those that match the callback method's signature.
-
-        Args:
-            method: The callback method to inspect
-            kwargs: The keyword arguments to filter
-
-        Returns:
-            A filtered dictionary of keyword arguments that match the method's signature
-
-        """
-        try:
-            sig = inspect.signature(method)
-        except (ValueError, TypeError):
-            # If we can't inspect the signature, pass all kwargs to maintain backward compatibility
-            return kwargs
-
-        filtered_kwargs = {}
-        for name, param in sig.parameters.items():
-            # If the method accepts **kwargs, pass all original kwargs directly
-            if param.kind == inspect.Parameter.VAR_KEYWORD:
-                return kwargs
-            # If the parameter exists in the incoming kwargs, add it to filtered_kwargs
-            if name in kwargs:
-                filtered_kwargs[name] = kwargs[name]
-
-        return filtered_kwargs
-
     def call(self, hook_name: str, *args: Any, **kwargs: Any) -> None:
         r"""Trigger the callback methods with the given name and arguments.
 
@@ -1022,9 +994,7 @@ class Fabric:
         Args:
             hook_name: The name of the callback method.
             *args: Optional positional arguments that get passed down to the callback method.
-            **kwargs: Optional keyword arguments that get passed down to the callback method. Keyword arguments
-                that are not present in the callback's signature will be filtered out automatically, allowing
-                callbacks to have different signatures for the same hook.
+            **kwargs: Optional keyword arguments that get passed down to the callback method.
 
         Example::
 
@@ -1046,8 +1016,13 @@ class Fabric:
                 )
                 continue
 
-            filtered_kwargs = self._filter_kwargs_for_callback(method, kwargs)
-            method(*args, **filtered_kwargs)
+            method(*args, **kwargs)
+
+            # TODO(fabric): handle the following signatures
+            # method(self, fabric|trainer, x, y=1)
+            # method(self, fabric|trainer, *args, x, y=1)
+            # method(self, *args, y=1)
+            # method(self, *args, **kwargs)
 
     def log(self, name: str, value: Any, step: Optional[int] = None) -> None:
         """Log a scalar to all loggers that were added to Fabric.

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/plugins/environments/lightning.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 
 import os
-import socket
 
 from typing_extensions import override
 
 from lightning_fabric.plugins.environments.cluster_environment import ClusterEnvironment
+from lightning_fabric.utilities.port_manager import get_port_manager
 from lightning_fabric.utilities.rank_zero import rank_zero_only
 
 
@@ -104,21 +104,38 @@ class LightningEnvironment(ClusterEnvironment):
         if "WORLD_SIZE" in os.environ:
             del os.environ["WORLD_SIZE"]
 
+        if self._main_port != -1:
+            get_port_manager().release_port(self._main_port)
+            self._main_port = -1
+
+        os.environ.pop("MASTER_PORT", None)
+        os.environ.pop("MASTER_ADDR", None)
+
 
 def find_free_network_port() -> int:
     """Finds a free port on localhost.
 
-    If the environment variable `STANDALONE_PORT` is set, its value is used as the port number.
-
     It is useful in single-node training when we don't want to connect to a real main node but have to set the
     `MASTER_PORT` environment variable.
 
+    The allocated port is reserved and won't be returned by subsequent calls until it's explicitly released.
+
+    Returns:
+        A port number that is reserved and free at the time of allocation
+
     """
-    if "STANDALONE_PORT" in os.environ:
-        _port = os.environ["STANDALONE_PORT"]
-        return int(_port)
-    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    s.bind(("", 0))
-    port = s.getsockname()[1]
-    s.close()
-    return port
+    # If an external launcher already specified a MASTER_PORT (for example, torch.distributed.spawn or
+    # multiprocessing helpers), reserve it through the port manager so no other test reuses the same number.
+    if "MASTER_PORT" in os.environ:
+        master_port_str = os.environ["MASTER_PORT"]
+        try:
+            existing_port = int(master_port_str)
+        except ValueError:
+            pass
+        else:
+            port_manager = get_port_manager()
+            if port_manager.reserve_existing_port(existing_port):
+                return existing_port
+
+    port_manager = get_port_manager()
+    return port_manager.allocate_port()

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/plugins/environments/mpi.py

@@ -73,11 +73,7 @@ class MPIEnvironment(ClusterEnvironment):
         if not _MPI4PY_AVAILABLE:
             return False
 
-        try:
-            # mpi4py may be installed without MPI being present
-            from mpi4py import MPI
-        except ImportError:
-            return False
+        from mpi4py import MPI
 
         return MPI.COMM_WORLD.Get_size() > 1
 

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/plugins/io/checkpoint_io.py

@@ -47,20 +47,13 @@ class CheckpointIO(ABC):
         """
 
     @abstractmethod
-    def load_checkpoint(
-        self, path: _PATH, map_location: Optional[Any] = None, weights_only: Optional[bool] = None
-    ) -> dict[str, Any]:
+    def load_checkpoint(self, path: _PATH, map_location: Optional[Any] = 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: 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.
 

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/plugins/io/torch_io.py

@@ -59,10 +59,7 @@ class TorchCheckpointIO(CheckpointIO):
 
     @override
     def load_checkpoint(
-        self,
-        path: _PATH,
-        map_location: Optional[Callable] = lambda storage, loc: storage,
-        weights_only: Optional[bool] = None,
+        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage
     ) -> dict[str, Any]:
         """Loads checkpoint using :func:`torch.load`, with additional handling for ``fsspec`` remote loading of files.
 
@@ -70,11 +67,6 @@ class TorchCheckpointIO(CheckpointIO):
             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.
 
@@ -88,7 +80,7 @@ class TorchCheckpointIO(CheckpointIO):
         if not fs.exists(path):
             raise FileNotFoundError(f"Checkpoint file not found: {path}")
 
-        return pl_load(path, map_location=map_location, weights_only=weights_only)
+        return pl_load(path, map_location=map_location)
 
     @override
     def remove_checkpoint(self, path: _PATH) -> None:

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/plugins/precision/fsdp.py

@@ -24,7 +24,6 @@ from typing_extensions import get_args, override
 from lightning_fabric.plugins.precision.amp import _optimizer_handles_unscaling
 from lightning_fabric.plugins.precision.precision import Precision
 from lightning_fabric.plugins.precision.utils import _convert_fp_tensor, _DtypeContextManager
-from lightning_fabric.utilities import rank_zero_warn
 from lightning_fabric.utilities.types import Optimizable
 
 if TYPE_CHECKING:
@@ -85,18 +84,19 @@ class FSDPPrecision(Precision):
     def mixed_precision_config(self) -> "TorchMixedPrecision":
         from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision as TorchMixedPrecision
 
-        if self.precision in ("16-true", "bf16-true"):
-            rank_zero_warn(
-                f"FSDP with `{self.precision}` enables computation in lower precision. "
-                "FSDP will always retain a full-precision copy of the model parameters for sharding."
-            )
-
-        if self.precision in ("16-true", "16-mixed"):
+        if self.precision == "16-mixed":
+            param_dtype = torch.float32
+            reduce_dtype = buffer_dtype = torch.float16
+        elif self.precision == "bf16-mixed":
+            param_dtype = torch.float32
+            reduce_dtype = buffer_dtype = torch.bfloat16
+        elif self.precision == "16-true":
             param_dtype = reduce_dtype = buffer_dtype = torch.float16
-        elif self.precision in ("bf16-true", "bf16-mixed"):
+        elif self.precision == "bf16-true":
             param_dtype = reduce_dtype = buffer_dtype = torch.bfloat16
         elif self.precision == "32-true":
-            param_dtype = reduce_dtype = buffer_dtype = torch.float32
+            param_dtype = torch.float32
+            reduce_dtype = buffer_dtype = torch.float32
         else:
             raise ValueError(f"Was unable to infer precision type, received {self.precision!r}.")
 

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/strategies/deepspeed.py

@@ -473,7 +473,6 @@ class DeepSpeedStrategy(DDPStrategy, _Sharded):
         path: _PATH,
         state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
-        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Load the contents from a checkpoint and restore the state of the given objects.
 
@@ -499,7 +498,7 @@ class DeepSpeedStrategy(DDPStrategy, _Sharded):
             # This code path to enables loading a checkpoint from a non-deepspeed checkpoint or from
             # a consolidated checkpoint
             path = self.broadcast(path)
-            return super().load_checkpoint(path=path, state=state, strict=strict, weights_only=weights_only)
+            return super().load_checkpoint(path=path, state=state, strict=strict)
 
         if not state:
             raise ValueError(

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/strategies/fsdp.py

@@ -516,7 +516,6 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
         path: _PATH,
         state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
-        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Load the contents from a checkpoint and restore the state of the given objects."""
         if not state:
@@ -587,7 +586,7 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
                         optim.load_state_dict(flattened_osd)
 
             # Load metadata (anything not a module or optimizer)
-            metadata = torch.load(path / _METADATA_FILENAME, weights_only=weights_only)
+            metadata = torch.load(path / _METADATA_FILENAME)
             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:

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/strategies/model_parallel.py

@@ -275,7 +275,6 @@ class ModelParallelStrategy(ParallelStrategy):
         path: _PATH,
         state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
-        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Load the contents from a checkpoint and restore the state of the given objects."""
         if not state:
@@ -296,7 +295,7 @@ class ModelParallelStrategy(ParallelStrategy):
                 f"Loading a single optimizer object from a checkpoint is not supported yet with {type(self).__name__}."
             )
 
-        return _load_checkpoint(path=path, state=state, strict=strict, weights_only=weights_only)
+        return _load_checkpoint(path=path, state=state, strict=strict)
 
     def _setup_distributed(self) -> None:
         reset_seed()
@@ -412,7 +411,6 @@ def _load_checkpoint(
     state: dict[str, Union[Module, Optimizer, Any]],
     strict: bool = True,
     optimizer_states_from_list: bool = False,
-    weights_only: Optional[bool] = None,
 ) -> dict[str, Any]:
     from torch.distributed.checkpoint.state_dict import (
         StateDictOptions,
@@ -451,7 +449,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, weights_only=weights_only)
+        metadata = torch.load(path / _METADATA_FILENAME)
         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:
@@ -463,7 +461,7 @@ def _load_checkpoint(
         return metadata
 
     if _is_full_checkpoint(path):
-        checkpoint = torch.load(path, mmap=True, map_location="cpu", weights_only=weights_only)
+        checkpoint = torch.load(path, mmap=True, map_location="cpu", weights_only=False)
         _load_raw_module_state(checkpoint.pop(module_key), module, strict=strict)
 
         state_dict_options = StateDictOptions(

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/strategies/strategy.py

@@ -310,7 +310,6 @@ class Strategy(ABC):
         path: _PATH,
         state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
-        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Load the contents from a checkpoint and restore the state of the given objects.
 
@@ -331,7 +330,7 @@ class Strategy(ABC):
 
         """
         torch.cuda.empty_cache()
-        checkpoint = self.checkpoint_io.load_checkpoint(path, weights_only=weights_only)
+        checkpoint = self.checkpoint_io.load_checkpoint(path)
         if not state:
             return checkpoint
 

{pytorch_lightning-2.6.0 → pytorch_lightning-2.6.0.dev0}/src/lightning_fabric/strategies/xla_fsdp.py

@@ -516,7 +516,6 @@ class XLAFSDPStrategy(ParallelStrategy, _Sharded):
         path: _PATH,
         state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
         strict: bool = True,
-        weights_only: Optional[bool] = None,
     ) -> dict[str, Any]:
         """Given a folder, load the contents from a checkpoint and restore the state of the given objects.
 
@@ -609,7 +608,7 @@ class XLAFSDPStrategy(ParallelStrategy, _Sharded):
                 )
             if "model" not in state or not isinstance(model := state["model"], torch.nn.Module):
                 raise NotImplementedError("XLAFSDP only supports a single model instance with 'model' as the key.")
-            full_ckpt = torch.load(path, weights_only=weights_only)
+            full_ckpt = torch.load(path)
             model.load_state_dict(full_ckpt.pop("model"), strict=strict)
             return full_ckpt