feat: enable save/load after recovery without additional steps by psyberck · Pull Request #619 · PrunaAI/pruna

psyberck · 2026-04-07T05:35:37Z

Description

Extend save and load to also allow saving and loading after recovery without additional steps.

Related Issue

Fixes #603

Type of Change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Refactor (no functional change)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Documentation update

Testing

I added or updated tests covering my changes
Existing tests pass locally (uv run pytest -m "cpu and not slow")

For full setup and testing instructions, see the Contributing Guide.

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my code, especially for agent-assisted changes
I updated the documentation where necessary

Thanks for contributing to Pruna! We're excited to review your work.

New to contributing? Check out our Contributing Guide for everything you need to get started.

Note:

Draft PRs or PRs without a clear and detailed overview may be delayed.

Please mark your PR as Ready for Review and ensure the sections above are filled out.

Contributions that are entirely AI-generated without meaningful human review are discouraged.

First Prune (1-year OSS anniversary)

First Prune marks one year of Pruna’s open-source work. During the initiative window, qualifying merged contributions count toward First Prune. You can earn credits for our performance models via our API.

If you’d like your contribution to count toward First Prune, here’s how it works:

Initiative window: First Prune starts on March 31.
Issue assignment: For your PR to count toward First Prune, the related issue must be assigned to the contributor opening the PR. Issues are labeled with first-prune.
Open for review: Please open your PR and mark it ready for review by April 30 (end of April).
Review priority: We’ll make our best effort to review quickly any PR that is open and has a review request before April 30.
Credits: Each qualifying merged PR earns 30 credits. We’ll be in touch after all qualifying PRs for First Prune have been merged.
To get started: Have a look at all models. You’ll need to sign up on the dashboard before you can redeem your credits.

…: get_save_before_smash_dir

…'s weights for "save_before_apply" algos

codacy-production · 2026-04-07T05:36:35Z

Not up to standards ⛔

🔴 Issues 3 high

Alerts:
⚠ 3 issues (≤ 0 issues of at least minor severity)

Results:
3 new issues

Category Results

Security 3 high

View in Codacy

🟢 Metrics 12 complexity · 0 duplication

Metric Results

Complexity 12

Duplication 0

View in Codacy

_{TIP This summary will be updated as you push new changes. Give us feedback}

sdiazlor · 2026-04-08T09:13:49Z

Hi @psyberck! Thank you for your contribution! I've run the tests. Please take a look before we jump into any review :)

psyberck · 2026-04-13T17:25:40Z

Hey @sdiazlor , can you help kindly re-run the tests? There seems some network failures during the tests:

ERROR tests/algorithms/test_algorithms.py::test_full_integration[TestIFW_whisper_tiny_random-cpu] - huggingface_hub.errors.LocalEntryNotFoundError: An error happened while trying to locate the files on the Hub and we cannot find the appropriate snapshot folder for the specified revision on the local disk. Please check your internet connection and try again.

sdiazlor · 2026-04-13T20:57:41Z

Sure thing! I ran them again :)

github-actions · 2026-05-02T00:22:58Z

This PR has been inactive for 10 days and is now marked as stale.

simlang

Thanks for looking at this! I left a view requests :)

simlang · 2026-05-04T13:35:27Z

+            smash_config.save_fns = [fn for fn in smash_config.save_fns if fn != SAVE_FUNCTIONS.save_before_apply.name]
+            # Re-save with recovered weights
+            shutil.rmtree(save_dir, ignore_errors=True)
+            save_dir.mkdir(parents=True)


this line can be removed, as this is already handled in save_pruna_model

simlang · 2026-05-04T13:57:00Z

+
+
+@pytest.mark.cpu
+def test_recovery_save_fn_is_none() -> None:


imo this test can be removed

simlang · 2026-05-04T13:59:37Z

+
+
+@pytest.mark.cpu
+def test_recovery_does_not_add_to_save_fns(tmp_path) -> None:


i'm not sure what this is testing, as recovery is not applied

simlang · 2026-05-04T14:01:33Z

+    # Simulate a save_before_apply algorithm having run before recovery:
+    # 1. Save original (pre-transformation) model to cache
+    save_dir = PrunaAlgorithmBase.get_save_before_smash_dir(config)
+    save_dir.mkdir(parents=True)


creating the directory is handled in the save

simlang · 2026-05-04T14:08:46Z

            adapter_smash_config = SmashConfigPrefixWrapper(smash_config, adapter.adapter_prefix + "_")
            adapter.pre_smash_hook(model_recovery, adapter_smash_config, seed=adapter_seed)

+    def apply(self, model: Any, smash_config: SmashConfig) -> Any:


imo it would make sense to add a new saving function instead of overriding the apply here :)

simlang · 2026-05-04T14:09:37Z

+    """Test that recovery refreshes a stale save_before_apply cache with recovered weights."""
+    from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
+
+    model = AutoModelForCausalLM.from_pretrained("yujiepan/opt-tiny-random")


since we only change a single weight, i don't think it's necessary to load a HF model, but instead we probably should use a dummy model like other tests above.

simlang · 2026-05-04T14:13:51Z

+    # 3. Simulate the transformation (e.g., half) + recovery modifying weights
+    model.lm_head.weight.data.fill_(0.99)  # "recovered" weights
+
+    # 4. Simulate what recovery's apply() does: refresh the stale cache


i don't think that simulating is a good idea, because if the actual recovery-apply-save-logic changes, it's not actually being tested anymore.

sdiazlor · 2026-05-11T16:13:17Z

Hi @psyberck! When you have time, could you take a look at the comments? :)

psyberck · 2026-05-11T18:09:03Z

Thanks for the comments @simlang. I will update accordingly tonight @sdiazlor :)

psyberck added 3 commits April 6, 2026 23:06

refactor: add prepare_saving as config property and get util function…

cea0bb9

…: get_save_before_smash_dir

refactor: save_fns won't append any new save_fn; save recovered model…

67880c1

…'s weights for "save_before_apply" algos

test: add unit tests to verify main ideas

c86401d

fix: remove unused imports from test_save.py

10627a3

psyberck changed the title ~~Feat/save load recovery~~ feat: enable save/load after recovery without additional steps Apr 8, 2026

docs: add docstring for get_save_before_smash_dir

e5aaccd

sdiazlor requested review from johannaSommer and simlang April 21, 2026 13:59

github-actions Bot added the stale label May 2, 2026

sdiazlor removed the stale label May 4, 2026

simlang requested changes May 4, 2026

View reviewed changes



		@pytest.mark.cpu
		def test_recovery_save_fn_is_none() -> None:



		@pytest.mark.cpu
		def test_recovery_does_not_add_to_save_fns(tmp_path) -> None:

Conversation

psyberck commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

Testing

Checklist

First Prune (1-year OSS anniversary)

Uh oh!

codacy-production Bot commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Not up to standards ⛔

Uh oh!

sdiazlor commented Apr 8, 2026

Uh oh!

psyberck commented Apr 13, 2026

Uh oh!

sdiazlor commented Apr 13, 2026

Uh oh!

github-actions Bot commented May 2, 2026

Uh oh!

simlang left a comment

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

simlang May 4, 2026

Choose a reason for hiding this comment

Uh oh!

sdiazlor commented May 11, 2026

Uh oh!

psyberck commented May 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

psyberck commented Apr 7, 2026 •

edited

Loading

codacy-production Bot commented Apr 7, 2026 •

edited

Loading