Fix ModelPruning iterative reuse RuntimeError on non-leaf tensors

When ModelPruning with use_lottery_ticket_hypothesis=True is reused
across multiple trainer.fit() calls, setup() is called again each run.
After the first pass, _copy_param sets dst.data = src.data, making the
parameters non-leaf tensors. The subsequent deepcopy(module) then raises:

  RuntimeError: Only Tensors created explicitly by the user (graph
  leaves) support the deepcopy protocol at the moment

Fix by adding _deepcopy_for_pruning(), a helper that temporarily
replaces non-leaf parameters with detached clones before deepcopy and
restores the originals afterward. Also explicitly set _original_layers
to None before re-populating in setup() so the previous run's tensor
references are released before new copies are allocated.

Fixes #8542

Author: Priyanshu-byte-coder

src/lightning/pytorch/callbacks/pruning.py

@@ -275,6 +275,31 @@ def make_pruning_permanent(self, module: nn.Module) -> None:
                     hook.remove(module)
                     del module._forward_pre_hooks[k]
 
+    @staticmethod
+    def _deepcopy_for_pruning(module: nn.Module) -> nn.Module:
+        """Deep-copy a module safely when parameters may be non-leaf tensors.
+
+        After a pruning pass with ``use_lottery_ticket_hypothesis=True``, the
+        module's parameters are rewritten via ``_copy_param`` (``dst.data =
+        src.data``).  This makes them non-leaf tensors, and a plain
+        ``deepcopy`` raises ``RuntimeError: Only Tensors created explicitly by
+        the user (graph leaves) support the deepcopy protocol``.
+
+        This helper temporarily replaces every non-leaf parameter with a
+        detached clone, performs the deep-copy, then restores the originals so
+        the live model is unchanged.
+        """
+        non_leaf: dict[str, Tensor] = {}
+        for param_name, param in list(module._parameters.items()):
+            if param is not None and not param.is_leaf:
+                non_leaf[param_name] = param
+                module._parameters[param_name] = param.detach().clone()
+        try:
+            return deepcopy(module)
+        finally:
+            for param_name, original in non_leaf.items():
+                module._parameters[param_name] = original
+
     @staticmethod
     def _copy_param(new: nn.Module, old: nn.Module, name: str) -> None:
         # Check if the parameter has been pruned (has _orig suffix)
@@ -376,12 +401,18 @@ def setup(self, trainer: "pl.Trainer", pl_module: LightningModule, stage: str) -
         self._parameters_to_prune = self.filter_parameters_to_prune(parameters_to_prune)
 
         if self._use_lottery_ticket_hypothesis:
+            # Release references from any previous run so their tensors can be
+            # garbage-collected before we allocate the new copies.
+            self._original_layers = None
             # group modules by id. Each entry has a copy of the initial data
             # and a list of the associated parameter names to prune
             self._original_layers = {}
             for i, (module, name) in enumerate(self._parameters_to_prune):
                 id_ = id(module)
-                self._original_layers.setdefault(id_, _LayerRef(data=deepcopy(module), names=[]))
+                # Use the detach-safe helper so that iterative pruning (where
+                # parameters may already be non-leaf tensors from a previous
+                # pruning cycle) does not raise a RuntimeError.
+                self._original_layers.setdefault(id_, _LayerRef(data=self._deepcopy_for_pruning(module), names=[]))
                 self._original_layers[id_]["names"].append((i, name))
 
     def _run_pruning(self, current_epoch: int) -> None:

tests/tests_pytorch/callbacks/test_pruning.py

@@ -551,3 +551,38 @@ def forward(self, x):
     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
+
+
+def test_iterative_pruning_no_runtime_error(tmp_path):
+    """Reusing a ModelPruning callback with use_lottery_ticket_hypothesis across multiple trainer.fit()
+    calls must not raise RuntimeError due to non-leaf tensors.
+
+    Regression test for https://github.com/Lightning-AI/pytorch-lightning/issues/8542
+    """
+    seed_everything(42)
+
+    model = BoringModel()
+    pruning_callback = ModelPruning(
+        "l1_unstructured",
+        use_lottery_ticket_hypothesis=True,
+        use_global_unstructured=True,
+        make_pruning_permanent=False,
+        amount=0.2,
+    )
+
+    for _ in range(3):
+        trainer = Trainer(
+            default_root_dir=tmp_path,
+            enable_progress_bar=False,
+            enable_model_summary=False,
+            enable_checkpointing=False,
+            logger=False,
+            limit_train_batches=2,
+            limit_val_batches=1,
+            max_epochs=1,
+            accelerator="cpu",
+            callbacks=[pruning_callback],
+        )
+        # Must not raise RuntimeError: "Only Tensors created explicitly by the
+        # user (graph leaves) support the deepcopy protocol"
+        trainer.fit(model)