Merge branch 'master' into dependabot-pip-requirements-numpy-gte-1.21.0-and-lt-2.4.0

Borda · web-flow · commit 21c956567fba · 2025-10-21T14:06:53.000+02:00
diff --git a/Makefile b/Makefile
@@ -7,7 +7,11 @@ export SPHINX_MOCK_REQUIREMENTS=1
 # install only Lightning Trainer packages
 export PACKAGE_NAME=pytorch
 
+
+# In Lightning Studio, the `lightning` package comes pre-installed.
+# Uninstall it first to ensure the editable install works correctly.
 setup:
+	uv pip uninstall lightning pytorch-lightning lightning-fabric || true
 	uv pip install -r requirements.txt \
 	    -r requirements/pytorch/base.txt \
 	    -r requirements/pytorch/test.txt \
diff --git a/docs/source-pytorch/deploy/production_advanced_2.rst b/docs/source-pytorch/deploy/production_advanced_2.rst
@@ -7,15 +7,20 @@ Deploy models into production (advanced)
 
 ----
 
-*********************************
-Compile your model to TorchScript
-*********************************
-`TorchScript <https://pytorch.org/docs/stable/jit.html>`_ allows you to serialize your models in a way that it can be loaded in non-Python environments.
-The ``LightningModule`` has a handy method :meth:`~lightning.pytorch.core.LightningModule.to_torchscript` that returns a scripted module which you
-can save or directly use.
+************************************
+Export your model with torch.export
+************************************
+
+`torch.export <https://pytorch.org/docs/stable/export.html>`_ is the recommended way to capture PyTorch models for
+deployment in production environments. It produces a clean intermediate representation with strong soundness guarantees,
+making models suitable for inference optimization and cross-platform deployment.
+You can export any ``LightningModule`` using the ``torch.export.export()`` API.
 
 .. testcode:: python
 
+    import torch
+    from torch.export import export
+
     class SimpleModel(LightningModule):
         def __init__(self):
             super().__init__()
@@ -25,25 +30,27 @@ can save or directly use.
             return torch.relu(self.l1(x.view(x.size(0), -1)))
 
 
-    # create the model
+    # create the model and example input
     model = SimpleModel()
-    script = model.to_torchscript()
+    example_input = torch.randn(1, 64)
 
-    # save for use in production environment
-    torch.jit.save(script, "model.pt")
+    # export the model
+    exported_program = export(model, (example_input,))
 
-It is recommended that you install the latest supported version of PyTorch to use this feature without limitations.
+    # save for use in production environment
+    torch.export.save(exported_program, "model.pt2")
 
-Once you have the exported model, you can run it in PyTorch or C++ runtime:
+It is recommended that you install the latest supported version of PyTorch to use this feature without
+limitations. Once you have the exported model, you can load and run it:
 
 .. code-block:: python
 
     inp = torch.rand(1, 64)
-    scripted_module = torch.jit.load("model.pt")
-    output = scripted_module(inp)
+    loaded_program = torch.export.load("model.pt2")
+    output = loaded_program.module()(inp)
 
 
-If you want to script a different method, you can decorate the method with :func:`torch.jit.export`:
+For more complex models, you can also export specific methods by creating a wrapper:
 
 .. code-block:: python
 
@@ -54,7 +61,6 @@ If you want to script a different method, you can decorate the method with :func
             self.dropout = nn.Dropout()
             self.mc_iteration = mc_iteration
 
-        @torch.jit.export
         def predict_step(self, batch, batch_idx):
             # enable Monte Carlo Dropout
             self.dropout.train()
@@ -66,4 +72,11 @@ If you want to script a different method, you can decorate the method with :func
 
 
     model = LitMCdropoutModel(...)
-    script = model.to_torchscript(file_path="model.pt", method="script")
+    example_batch = torch.randn(32, 10)  # example input
+
+    # Export the predict_step method
+    exported_program = torch.export.export(
+        lambda batch, idx: model.predict_step(batch, idx),
+        (example_batch, 0)
+    )
+    torch.export.save(exported_program, "mc_dropout_model.pt2")
diff --git a/requirements/doctests.txt b/requirements/doctests.txt
@@ -1,2 +1,2 @@
 pytest ==8.4.2
-pytest-doctestplus ==1.4.0
+pytest-doctestplus ==1.5.0
diff --git a/requirements/fabric/strategies.txt b/requirements/fabric/strategies.txt
@@ -5,5 +5,5 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.14.1,<=0.15.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
+deepspeed >=0.15.0,<0.17.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
 bitsandbytes >=0.45.2,<0.47.0; platform_system != "Darwin"
diff --git a/requirements/pytorch/extra.txt b/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.42.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"
diff --git a/requirements/pytorch/strategies.txt b/requirements/pytorch/strategies.txt
@@ -3,4 +3,4 @@
 
 # note: is a bug around 0.10 with `MPS_Accelerator must implement all abstract methods`
 #  shall be resolved by https://github.com/microsoft/DeepSpeed/issues/4372
-deepspeed >=0.14.1,<=0.15.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
+deepspeed >=0.15.0,<0.17.0; platform_system != "Windows" and platform_system != "Darwin"  # strict
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -12,7 +12,7 @@ scikit-learn >0.22.1, <1.8.0
 onnx >1.12.0, <1.20.0
 onnxruntime >=1.12.0, <1.24.0
 onnxscript >= 0.1.0, < 0.5.0
-psutil <7.1.1 # for `DeviceStatsMonitor`
+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
diff --git a/src/lightning/fabric/strategies/deepspeed.py b/src/lightning/fabric/strategies/deepspeed.py
@@ -37,6 +37,7 @@
 from lightning.fabric.strategies.registry import _StrategyRegistry
 from lightning.fabric.strategies.strategy import _Sharded
 from lightning.fabric.utilities.distributed import log
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.fabric.utilities.load import _move_state_into
 from lightning.fabric.utilities.rank_zero import rank_zero_info, rank_zero_warn
 from lightning.fabric.utilities.seed import reset_seed
@@ -47,6 +48,7 @@
     from torch.optim.lr_scheduler import _LRScheduler
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
+_DEEPSPEED_GREATER_EQUAL_0_16 = RequirementCache("deepspeed>=0.16.0")
 
 
 # TODO(fabric): Links in the docstrings to PL-specific deepspeed user docs need to be replaced.
@@ -239,6 +241,19 @@ def __init__(
                 " Install it by running `pip install -U deepspeed`."
             )
 
+        if _TORCH_GREATER_EQUAL_2_6 and not _DEEPSPEED_GREATER_EQUAL_0_16:
+            # Starting with PyTorch 2.6, `torch.load` defaults to `weights_only=True` when loading full checkpoints.
+            # DeepSpeed added support for this behavior in version 0.16.0.
+            import deepspeed
+
+            deepspeed_version = deepspeed.__version__
+
+            raise ImportError(
+                f"PyTorch >= 2.6 requires DeepSpeed >= 0.16.0. "
+                f"Detected DeepSpeed version: {deepspeed_version}. "
+                "Please upgrade by running `pip install -U 'deepspeed>=0.16.0'`."
+            )
+
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,
diff --git a/src/lightning/fabric/utilities/imports.py b/src/lightning/fabric/utilities/imports.py
@@ -36,5 +36,6 @@
 _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_LESS_EQUAL_2_6 = compare_version("torch", operator.le, "2.6.0")
+_TORCH_GREATER_EQUAL_2_6 = compare_version("torch", operator.ge, "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)
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -52,6 +52,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed preventing recursive symlink creation iwhen `save_last='link'` and `save_top_k=-1` ([#21186](https://github.com/Lightning-AI/pytorch-lightning/pull/21186))
 
 
+- Fixed `ModelPruning` sparsity logging bug that caused incorrect sparsity percentages ([#21223](https://github.com/Lightning-AI/pytorch-lightning/pull/21223))
+
+
 - Fixed `LightningCLI` loading of hyperparameters from `ckpt_path` failing for subclass model mode ([#21246](https://github.com/Lightning-AI/pytorch-lightning/pull/21246))
 
 
diff --git a/src/lightning/pytorch/callbacks/pruning.py b/src/lightning/pytorch/callbacks/pruning.py
@@ -349,7 +349,7 @@ def apply_pruning(self, amount: Union[int, float]) -> None:
     def _log_sparsity_stats(
         self, prev: list[tuple[int, int]], curr: list[tuple[int, int]], amount: Union[int, float] = 0
     ) -> None:
-        total_params = sum(p.numel() for layer, _ in self._parameters_to_prune for p in layer.parameters())
+        total_params = sum(total for _, total in curr)
         prev_total_zeros = sum(zeros for zeros, _ in prev)
         curr_total_zeros = sum(zeros for zeros, _ in curr)
         log.info(
diff --git a/src/lightning/pytorch/strategies/deepspeed.py b/src/lightning/pytorch/strategies/deepspeed.py
@@ -35,10 +35,12 @@
 from lightning.fabric.strategies import _StrategyRegistry
 from lightning.fabric.strategies.deepspeed import (
     _DEEPSPEED_AVAILABLE,
+    _DEEPSPEED_GREATER_EQUAL_0_16,
     _format_precision_config,
     _validate_checkpoint_directory,
     _validate_device_index_selection,
 )
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.fabric.utilities.optimizer import _optimizers_to_device
 from lightning.fabric.utilities.seed import reset_seed
 from lightning.fabric.utilities.types import _PATH
@@ -262,6 +264,19 @@ def __init__(
                 " Install it by running `pip install -U deepspeed`."
             )
 
+        if _TORCH_GREATER_EQUAL_2_6 and not _DEEPSPEED_GREATER_EQUAL_0_16:
+            # Starting with PyTorch 2.6, `torch.load` defaults to `weights_only=True` when loading full checkpoints.
+            # DeepSpeed added support for this behavior in version 0.16.0.
+            import deepspeed
+
+            deepspeed_version = deepspeed.__version__
+
+            raise ImportError(
+                f"PyTorch >= 2.6 requires DeepSpeed >= 0.16.0. "
+                f"Detected DeepSpeed version: {deepspeed_version}. "
+                "Please upgrade by running `pip install -U 'deepspeed>=0.16.0'`."
+            )
+
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,
diff --git a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py
@@ -91,8 +91,7 @@ def _generate_sync_fn(self) -> None:
         """Used to compute the syncing function and cache it."""
         fn = self.no_op if self.fn is None or not self.should or self.rank_zero_only else self.fn
         # save the function as `_fn` as the meta are being re-created and the object references need to match.
-        # ignore typing, bad support for `partial`: mypy/issues/1484
-        self._fn: Callable = partial(fn, reduce_op=self.op, group=self.group)  # type: ignore[unused-ignore]
+        self._fn: Callable = partial(fn, reduce_op=self.op, group=self.group)
 
     @property
     def __call__(self) -> Any:
diff --git a/src/lightning/pytorch/utilities/deepspeed.py b/src/lightning/pytorch/utilities/deepspeed.py
@@ -93,10 +93,10 @@ def convert_zero_checkpoint_to_fp32_state_dict(
     ]
     checkpoint_dir = ds_checkpoint_dir(checkpoint_dir)
     optim_files = get_optim_files(checkpoint_dir)
-    optim_state = torch.load(optim_files[0], map_location=CPU_DEVICE)
+    optim_state = torch.load(optim_files[0], map_location=CPU_DEVICE, weights_only=False)
     zero_stage = optim_state["optimizer_state_dict"]["zero_stage"]
     model_file = get_model_state_file(checkpoint_dir, zero_stage)
-    client_state = torch.load(model_file, map_location=CPU_DEVICE)
+    client_state = torch.load(model_file, map_location=CPU_DEVICE, weights_only=False)
     client_state = {key: value for key, value in client_state.items() if key not in deepspeed_states}
     # State dict keys will include reference to wrapper _LightningModuleWrapperBase in old checkpoints created in
     # Lightning version < 2.1. Delete the `_forward_module` prefix before saving.
diff --git a/src/lightning/pytorch/utilities/model_helpers.py b/src/lightning/pytorch/utilities/model_helpers.py
@@ -19,7 +19,7 @@
 
 from lightning_utilities.core.imports import RequirementCache
 from torch import nn
-from typing_extensions import Concatenate, ParamSpec
+from typing_extensions import Concatenate, ParamSpec, override
 
 import lightning.pytorch as pl
 
@@ -104,26 +104,32 @@ def _check_mixed_imports(instance: object) -> None:
 _R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
 
 
-class _restricted_classmethod_impl(Generic[_T, _R_co, _P]):
+class _restricted_classmethod_impl(classmethod, Generic[_T, _P, _R_co]):
     """Drop-in replacement for @classmethod, but raises an exception when the decorated method is called on an instance
     instead of a class type."""
 
+    method: Callable[Concatenate[type[_T], _P], _R_co]
+
     def __init__(self, method: Callable[Concatenate[type[_T], _P], _R_co]) -> None:
+        super().__init__(method)
         self.method = method
 
-    def __get__(self, instance: Optional[_T], cls: type[_T]) -> Callable[_P, _R_co]:
+    @override
+    def __get__(self, instance: _T, cls: Optional[type[_T]] = None) -> Callable[_P, _R_co]:  # type: ignore[override]
         # The wrapper ensures that the method can be inspected, but not called on an instance
         @functools.wraps(self.method)
         def wrapper(*args: Any, **kwargs: Any) -> _R_co:
             # Workaround for https://github.com/pytorch/pytorch/issues/67146
             is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
+            cls_type = cls if cls is not None else type(instance)
             if instance is not None and not is_scripting:
                 raise TypeError(
-                    f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be called on an instance."
+                    f"The classmethod `{cls_type.__name__}.{self.method.__name__}` cannot be called on an instance."
                     " Please call it on the class type and make sure the return value is used."
                 )
-            return self.method(cls, *args, **kwargs)
+            return self.method(cls_type, *args, **kwargs)
 
+        wrapper.__func__ = self.method
         return wrapper
 
 
diff --git a/tests/tests_pytorch/callbacks/test_pruning.py b/tests/tests_pytorch/callbacks/test_pruning.py

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`pytest ==8.4.2`
`2`		`-pytest-doctestplus ==1.4.0`
	`2`	`+pytest-doctestplus ==1.5.0`