From 498e11f34d9b04e58e624ab70326aaefce6e9c28 Mon Sep 17 00:00:00 2001 From: Priyanshu Doshi Date: Sun, 17 May 2026 13:31:03 +0530 Subject: [PATCH 1/3] 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 --- src/lightning/pytorch/callbacks/pruning.py | 33 ++++++++++++++++- tests/tests_pytorch/callbacks/test_pruning.py | 35 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/callbacks/pruning.py b/src/lightning/pytorch/callbacks/pruning.py index f0e1bcbe49f99..af0e3ed54459a 100644 --- a/src/lightning/pytorch/callbacks/pruning.py +++ b/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: diff --git a/tests/tests_pytorch/callbacks/test_pruning.py b/tests/tests_pytorch/callbacks/test_pruning.py index 1a23efd919171..5e02e7829e76d 100644 --- a/tests/tests_pytorch/callbacks/test_pruning.py +++ b/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) From 03ed64ae41bd60b527c509b7c2a4298e4eac5e2d Mon Sep 17 00:00:00 2001 From: Priyanshu Doshi Date: Sun, 17 May 2026 16:32:13 +0530 Subject: [PATCH 2/3] Fix _deepcopy_for_pruning to patch __dict__ not _parameters The actual non-leaf tensor after pruning is module.weight stored in module.__dict__ by the forward pre-hook (weight_orig * weight_mask), not in module._parameters. Patching _parameters had two problems: 1. Did not fix the RuntimeError (wrong dict being patched) 2. Caused mypy errors: assigning Tensor to Parameter | None Switch to iterating module.__dict__ instead, which correctly captures the hook-written non-leaf weight attribute. --- src/lightning/pytorch/callbacks/pruning.py | 33 ++++++++++++---------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/src/lightning/pytorch/callbacks/pruning.py b/src/lightning/pytorch/callbacks/pruning.py index af0e3ed54459a..3066c03ed3d68 100644 --- a/src/lightning/pytorch/callbacks/pruning.py +++ b/src/lightning/pytorch/callbacks/pruning.py @@ -277,28 +277,31 @@ def make_pruning_permanent(self, module: nn.Module) -> None: @staticmethod def _deepcopy_for_pruning(module: nn.Module) -> nn.Module: - """Deep-copy a module safely when parameters may be non-leaf tensors. + """Deep-copy a module that may contain 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``. + PyTorch pruning hooks write the masked weight (e.g. + ``weight_orig * weight_mask``) back onto the module as a plain + ``__dict__`` attribute (``module.weight``) after each forward pass. + That stored value is a non-leaf tensor, so a bare ``deepcopy`` raises:: - 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. + RuntimeError: Only Tensors created explicitly by the user + (graph leaves) support the deepcopy protocol at the moment. + + This helper temporarily replaces any non-leaf tensor found in + ``module.__dict__`` with a detached leaf clone, performs the + deep-copy, then restores the originals so the live module 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() + for attr_name, attr_val in list(module.__dict__.items()): + if isinstance(attr_val, Tensor) and not attr_val.is_leaf: + non_leaf[attr_name] = attr_val + module.__dict__[attr_name] = attr_val.detach().clone() try: return deepcopy(module) finally: - for param_name, original in non_leaf.items(): - module._parameters[param_name] = original + for attr_name, original in non_leaf.items(): + module.__dict__[attr_name] = original @staticmethod def _copy_param(new: nn.Module, old: nn.Module, name: str) -> None: From 14e140828fc23eb1b2dfbd010cd976bdda218b05 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 17 May 2026 11:02:53 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning/pytorch/callbacks/pruning.py | 1 + tests/tests_pytorch/callbacks/test_pruning.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/lightning/pytorch/callbacks/pruning.py b/src/lightning/pytorch/callbacks/pruning.py index 3066c03ed3d68..4ae254ee2e247 100644 --- a/src/lightning/pytorch/callbacks/pruning.py +++ b/src/lightning/pytorch/callbacks/pruning.py @@ -291,6 +291,7 @@ def _deepcopy_for_pruning(module: nn.Module) -> nn.Module: ``module.__dict__`` with a detached leaf clone, performs the deep-copy, then restores the originals so the live module is unchanged. + """ non_leaf: dict[str, Tensor] = {} for attr_name, attr_val in list(module.__dict__.items()): diff --git a/tests/tests_pytorch/callbacks/test_pruning.py b/tests/tests_pytorch/callbacks/test_pruning.py index 5e02e7829e76d..77a495dbed11f 100644 --- a/tests/tests_pytorch/callbacks/test_pruning.py +++ b/tests/tests_pytorch/callbacks/test_pruning.py @@ -554,10 +554,11 @@ def forward(self, x): 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. + """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)