From 44b502ae81146580e415fe9830378ff709197ee6 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Thu, 22 Oct 2026 16:16:05 +0100 Subject: [PATCH] Fix ModelPruning._make_pruning_permanent_on_state_dict suffix stripping and O(M^2) sanitization 1. In ModelPruning._make_pruning_permanent_on_state_dict(), k.replace("_mask", "") replaced all occurrences of "_mask" in the state_dict key rather than only the trailing "_mask" suffix. When any pruned submodule or parameter path contains "_mask" (e.g., self.attention_mask_proj = nn.Linear(...) with keys attention_mask_proj.weight_orig and attention_mask_proj.weight_mask), k.replace("_mask", "") produced attention_proj.weight and caused state_dict.pop(tensor_name + "_orig") to crash with KeyError during checkpoint saving. Use k.removesuffix("_mask") so only the trailing suffix is stripped. 2. In ModelPruning.sanitize_parameters_to_prune(), convert current_modules to a set before validating explicit (module, name) pairs (O(1) lookup per pair instead of O(M) list scan), and avoid calling getattr(m, p, None) twice per (p, m) pair when parameters_to_prune is defaulted. Signed-off-by: Parth Vadera Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/lightning/pytorch/callbacks/pruning.py | 10 +++---- tests/tests_pytorch/callbacks/test_pruning.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/src/lightning/pytorch/callbacks/pruning.py b/src/lightning/pytorch/callbacks/pruning.py index f0e1bcbe49f99..e94a15238dd50 100644 --- a/src/lightning/pytorch/callbacks/pruning.py +++ b/src/lightning/pytorch/callbacks/pruning.py @@ -420,7 +420,7 @@ def _make_pruning_permanent_on_state_dict(self, pl_module: LightningModule) -> d state_dict = pl_module.state_dict() # find the mask and the original weights. - map_pruned_params = {k.replace("_mask", "") for k in state_dict if k.endswith("_mask")} + map_pruned_params = {k.removesuffix("_mask") for k in state_dict if k.endswith("_mask")} for tensor_name in map_pruned_params: orig = state_dict.pop(tensor_name + "_orig") mask = state_dict.pop(tensor_name + "_mask") @@ -459,10 +459,7 @@ def sanitize_parameters_to_prune( if not parameters_to_prune: parameters_to_prune = [ - (m, p) - for p in parameters - for m in current_modules - if getattr(m, p, None) is not None and isinstance(getattr(m, p, None), nn.Parameter) + (m, p) for p in parameters for m in current_modules if isinstance(getattr(m, p, None), nn.Parameter) ] elif ( isinstance(parameters_to_prune, (list, tuple)) @@ -470,9 +467,10 @@ def sanitize_parameters_to_prune( and all(len(p) == 2 for p in parameters_to_prune) and all(isinstance(a, nn.Module) and isinstance(b, str) for a, b in parameters_to_prune) ): + current_modules_set = set(current_modules) missing_modules, missing_parameters = [], [] for module, name in parameters_to_prune: - if module not in current_modules: + if module not in current_modules_set: missing_modules.append(module) continue if not hasattr(module, name): diff --git a/tests/tests_pytorch/callbacks/test_pruning.py b/tests/tests_pytorch/callbacks/test_pruning.py index 1a23efd919171..206577886946f 100644 --- a/tests/tests_pytorch/callbacks/test_pruning.py +++ b/tests/tests_pytorch/callbacks/test_pruning.py @@ -551,3 +551,29 @@ 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_make_pruning_permanent_on_state_dict_with_mask_in_module_name(): + """Test that `_make_pruning_permanent_on_state_dict` only strips the trailing `_mask` suffix.""" + + class MaskedLayerModel(BoringModel): + def __init__(self): + super().__init__() + self.attention_mask_proj = nn.Linear(32, 2, bias=False) + + model = MaskedLayerModel() + pruning = ModelPruning( + pruning_fn="l1_unstructured", + parameters_to_prune=[(model.attention_mask_proj, "weight")], + amount=0.5, + make_pruning_permanent=True, + ) + pytorch_prune.l1_unstructured(model.attention_mask_proj, name="weight", amount=0.5) + assert "attention_mask_proj.weight_orig" in model.state_dict() + assert "attention_mask_proj.weight_mask" in model.state_dict() + + permanent_state_dict = pruning._make_pruning_permanent_on_state_dict(model) + assert "attention_mask_proj.weight" in permanent_state_dict + assert "attention_mask_proj.weight_orig" not in permanent_state_dict + assert "attention_mask_proj.weight_mask" not in permanent_state_dict + assert torch.equal(permanent_state_dict["attention_mask_proj.weight"], model.attention_mask_proj.weight.cpu())