Skip to content

Conversation

@gsprochette
Copy link
Collaborator

Description

Formalize artifact saving and expose it to more algorithms, and add artifact hooks into Smash config

  • Generalized artifact saving/loading: Extracted and formalized the artifact-saving logic so it is no longer tied specifically to the Torch Compile workflow. Other Pruna algorithms can now register and reuse the same standardized artifact save/load utilities.
  • Config-driven hooks: Extended the Smash configuration to include save_artifacts_fns and load_artifacts_fns, working in a way similar to save_fns and load_fns
  • Improved composability: This structure makes it easy to save artifacts (such as parameters estimated from a warmup, or checkpoints for finetuning or distillation).

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

How Has This Been Tested?

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

Additional Notes

Copy link

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comment @cursor review or bugbot run to trigger another review on this PR

@gsprochette gsprochette requested review from llcnt and simlang January 14, 2026 17:17
Copy link
Collaborator

@llcnt llcnt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work, this is a very cool feature :) (I was personally doing a custom save/load artifacts in my ongoing PR )

Copy link
Member

@simlang simlang left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good! Thanks for taking care of that, super useful for literally everything 🚀

Just had a question about the backward compatibility

if SAVE_FUNCTIONS.torch_artifacts.name in smash_config.save_fns:
save_torch_artifacts(model, model_path, smash_config)
smash_config.save_fns.remove(SAVE_FUNCTIONS.torch_artifacts.name)
# Backward compatibility with torch artifacts save function
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i see why this is necessary, but it would be nice to remove this as soon as possible again - should we add a deprecation warning or something, and then actually remove it again at some point?

@gsprochette gsprochette force-pushed the refactor/formalize-artifact-saving branch from 4bcb4cd to 87a422d Compare January 21, 2026 14:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants