[deep1blue]

What does this PR do?

Fixes a KeyError bug when saving checkpoints with ModelPruning(..., make_pruning_permanent=True) on models whose submodule names contain "_mask", and removes an $O(M^2)$ linear list scan in ModelPruning.sanitize_parameters_to_prune:

1. Strip only the trailing "_mask" suffix in _make_pruning_permanent_on_state_dict:
   • Previously, map_pruned_params = {k.replace("_mask", "") for k in state_dict if k.endswith("_mask")} replaced every occurrence of "_mask" in the key string rather than only the trailing "_mask" suffix added by torch.nn.utils.prune.
   • If any pruned submodule name contains "_mask" (for example, self.attention_mask_proj = nn.Linear(...), whose pruned state_dict keys are "attention_mask_proj.weight_orig" and "attention_mask_proj.weight_mask"), k.replace("_mask", "") corrupted tensor_name to "attention_proj.weight", causing state_dict.pop(tensor_name + "_orig") on the next line to crash with KeyError: 'attention_proj.weight_orig' during Trainer.save_checkpoint().
   • Replacing k.replace("_mask", "") with k.removesuffix("_mask") strips only the trailing "_mask" suffix.
2. Avoid $O(M^2)$ module membership checks and duplicate getattr calls in sanitize_parameters_to_prune:
   • When parameters_to_prune is provided as a list/tuple of (module, name) pairs, if module not in current_modules: previously performed an $O(M)$ linear list scan over current_modules for every tuple ($O(M^2)$ total). Converting current_modules to current_modules_set = set(current_modules) makes membership checks $O(1)$.
   • When parameters_to_prune is empty, isinstance(getattr(m, p, None), nn.Parameter) already evaluates to False when getattr returns None, avoiding calling getattr(m, p, None) twice per (p, m) pair.

Acknowledgments

Credit to @tchaton, @carmocca, @awaelchli, @Borda, @SkafteNicki (#21191, #21223), and @chanokin (#20783) for prior work on ModelPruning.

Before submitting

• Was this discussed/agreed via a GitHub issue? (no, self-discovered bug fix + optimization)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request? (none)
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified