Merge branch 'master' into batch_size_scaler_newargs

deependujha · web-flow · commit 6264e8f1a010 · 2025-10-17T12:30:42.000+05:30
diff --git a/.github/workflows/_legacy-checkpoints.yml b/.github/workflows/_legacy-checkpoints.yml
@@ -58,7 +58,7 @@ jobs:
       - uses: actions/checkout@v5
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: "3.9"
           # TODO: Avoid activating environment like this
diff --git a/.github/workflows/ci-tests-fabric.yml b/.github/workflows/ci-tests-fabric.yml
@@ -71,7 +71,7 @@ jobs:
       - uses: actions/checkout@v5
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: ${{ matrix.config.python-version || '3.9' }}
           # TODO: Avoid activating environment like this
diff --git a/.github/workflows/ci-tests-pytorch.yml b/.github/workflows/ci-tests-pytorch.yml
@@ -76,7 +76,7 @@ jobs:
       - uses: actions/checkout@v5
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: ${{ matrix.config.python-version || '3.9' }}
           # TODO: Avoid activating environment like this
diff --git a/.github/workflows/code-checks.yml b/.github/workflows/code-checks.yml
@@ -31,7 +31,7 @@ jobs:
       - uses: actions/checkout@v5
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: "3.11"
           # TODO: Avoid activating environment like this
diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml
@@ -73,7 +73,7 @@ jobs:
           lfs: ${{ matrix.pkg-name == 'pytorch' }}
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: "3.10"
           # TODO: Avoid activating environment like this
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -23,7 +23,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: end-of-file-fixer
       - id: trailing-whitespace
@@ -70,7 +70,7 @@ repos:
       - id: sphinx-lint
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.2
+    rev: v0.13.3
     hooks:
       # try to fix what is possible
       - id: ruff
@@ -95,8 +95,8 @@ repos:
               README.md
           )$
 
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v3.1.0
+  - repo: https://github.com/JoC0de/pre-commit-prettier
+    rev: v3.6.2
     hooks:
       - id: prettier
         # https://prettier.io/docs/en/options.html#print-width
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/fabric/test.txt b/requirements/fabric/test.txt
@@ -1,9 +1,10 @@
-coverage ==7.10.6
+coverage ==7.10.7
 numpy >=1.21.0, <1.27.0
 pytest ==8.4.2
 pytest-cov ==6.3.0
 pytest-timeout ==2.4.0
-pytest-rerunfailures ==16.0.1
+pytest-rerunfailures ==16.0.1; python_version < "3.10"
+pytest-rerunfailures ==16.1; python_version >= "3.10"
 pytest-random-order ==1.2.0
 click ==8.1.8; python_version < "3.11"
 click ==8.3.0; python_version > "3.10"
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -1,8 +1,9 @@
-coverage ==7.10.6
+coverage ==7.10.7
 pytest ==8.4.2
 pytest-cov ==6.3.0
 pytest-timeout ==2.4.0
-pytest-rerunfailures ==16.0.1
+pytest-rerunfailures ==16.0.1; python_version < "3.10"
+pytest-rerunfailures ==16.1; python_version >= "3.10"
 pytest-random-order ==1.2.0
 
 # needed in tests
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/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/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
@@ -262,13 +262,13 @@ def test_multiple_pruning_callbacks(tmp_path, caplog, make_pruning_permanent: bo
     actual = [m for m in actual if m.startswith("Applied")]
     percentage = r"\(\d+(?:\.\d+)?%\)"
     expected = [
-        rf"Applied `L1Unstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}",
+        rf"Applied `L1Unstructured`. Pruned: \d+\/1088 {percentage} -> \d+\/1088 {percentage}",
         rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.5. Pruned: 0 \(0.00%\) -> \d+ {percentage}",  # noqa: E501
         rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.5. Pruned: 0 \(0.00%\) -> \d+ {percentage}",  # noqa: E501
-        rf"Applied `RandomUnstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}",
+        rf"Applied `RandomUnstructured`. Pruned: \d+\/1088 {percentage} -> \d+\/1088 {percentage}",
         rf"Applied `RandomUnstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.25. Pruned: \d+ {percentage} -> \d+ {percentage}",  # noqa: E501
         rf"Applied `RandomUnstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.25. Pruned: \d+ {percentage} -> \d+ {percentage}",  # noqa: E501
-        rf"Applied `L1Unstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}",
+        rf"Applied `L1Unstructured`. Pruned: \d+\/1088 {percentage} -> \d+\/1088 {percentage}",
         rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.5. Pruned: \d+ {percentage} -> \d+ {percentage}",  # noqa: E501
         rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.5. Pruned: \d+ {percentage} -> \d+ {percentage}",  # noqa: E501
     ]
@@ -329,9 +329,9 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint):
     actual = [m for m in actual if m.startswith("Applied")]
     percentage = r"\(\d+(?:\.\d+)?%\)"
     expected = [
-        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
-        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
-        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
+        rf"Applied `RandomUnstructured`. Pruned: \d+\/64 {percentage} -> \d+\/64 {percentage}",
+        rf"Applied `RandomUnstructured`. Pruned: \d+\/64 {percentage} -> \d+\/64 {percentage}",
+        rf"Applied `RandomUnstructured`. Pruned: \d+\/64 {percentage} -> \d+\/64 {percentage}",
     ]
     expected = [re.compile(s) for s in expected]
     assert all(regex.match(s) for s, regex in zip(actual, expected))
@@ -463,3 +463,91 @@ def __init__(self):
         f"Actual weight_orig: {weight_orig}\n"
         f"Max difference: {torch.max(torch.abs(weight_orig - original_weights))}"
     )
+
+
+@pytest.mark.parametrize("pruning_amount", [0.1, 0.2, 0.3, 0.5])
+@pytest.mark.parametrize("model_type", ["simple", "complex"])
+def test_sparsity_calculation(tmp_path, caplog, pruning_amount: float, model_type: str):
+    """Test that the sparsity calculation fix correctly reports percentages."""
+
+    class SimpleModel(BoringModel):
+        """Simple model with 66 parameters (64 weight + 2 bias)."""
+
+        def __init__(self):
+            super().__init__()
+            self.layer = nn.Linear(32, 2)  # 32*2 + 2 = 66 params
+
+    class ComplexModel(BoringModel):
+        """Complex model with multiple layers."""
+
+        def __init__(self):
+            super().__init__()
+            self.layer1 = nn.Linear(32, 64)  # 32*64 + 64 = 2112 params
+            self.layer2 = nn.Linear(64, 2)  # 64*2 + 2 = 130 params
+            # Total: 2112 + 130 = 2242 params (but only layer1 will be pruned)
+            # layer1 params: 2112
+
+        def forward(self, x):
+            x = torch.relu(self.layer1(x))
+            return self.layer2(x)
+
+    if model_type == "simple":
+        model = SimpleModel()
+        expected_total_params = 66
+        parameters_to_prune = None
+    else:
+        model = ComplexModel()
+        expected_total_params = 2112
+        parameters_to_prune = [(model.layer1, "weight"), (model.layer1, "bias")]
+
+    pruning = ModelPruning(
+        pruning_fn="l1_unstructured",
+        parameters_to_prune=parameters_to_prune,
+        amount=pruning_amount,
+        verbose=1,
+        use_global_unstructured=True,
+    )
+
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        enable_checkpointing=False,
+        logger=False,
+        limit_train_batches=1,
+        max_epochs=1,
+        accelerator="cpu",
+        callbacks=[pruning],
+    )
+
+    with caplog.at_level(INFO):
+        trainer.fit(model)
+
+    sparsity_logs = [msg for msg in caplog.messages if "Applied `L1Unstructured`. Pruned:" in msg]
+    assert len(sparsity_logs) == 1, f"Expected 1 sparsity log, got {len(sparsity_logs)}"
+    sparsity_log = sparsity_logs[0]
+    pattern = r"Applied `L1Unstructured`\. Pruned: \d+/(\d+) \(\d+\.\d+%\) -> (\d+)/(\d+) \((\d+\.\d+)%\)"
+    match = re.search(pattern, sparsity_log)
+    assert match, f"Could not parse sparsity log: {sparsity_log}"
+
+    total_params_before = int(match.group(1))
+    pruned_count = int(match.group(2))
+    total_params_after = int(match.group(3))
+    sparsity_percentage = float(match.group(4))
+    assert total_params_before == expected_total_params, (
+        f"Total parameter count mismatch for {model_type} model. "
+        f"Expected {expected_total_params}, got {total_params_before}"
+    )
+    assert total_params_after == expected_total_params, (
+        f"Total parameter count should be consistent. Before: {total_params_before}, After: {total_params_after}"
+    )
+
+    # Verify sparsity percentage is approximately correct
+    expected_sparsity = pruning_amount * 100
+    tolerance = 5.0
+    assert abs(sparsity_percentage - expected_sparsity) <= tolerance
+
+    # Verify the number of pruned parameters is reasonable
+    expected_pruned_count = int(expected_total_params * pruning_amount)
+    pruned_tolerance = max(1, int(expected_total_params * 0.05))
+    assert abs(pruned_count - expected_pruned_count) <= pruned_tolerance
diff --git a/tests/tests_pytorch/graveyard/test_precision.py b/tests/tests_pytorch/graveyard/test_precision.py
diff --git a/tests/tests_pytorch/utilities/test_model_helpers.py b/tests/tests_pytorch/utilities/test_model_helpers.py
