Fix ModelPruning._make_pruning_permanent_on_state_dict suffix stripping and O(M^2) sanitization

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,20 +459,18 @@ 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))
             and len(parameters_to_prune) > 0
             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):

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())