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

if LOAD_FUNCTIONS.torch_artifacts.name in smash_config.load_fns:
load_torch_artifacts(model_path, **kwargs)
smash_config.load_fns.remove(LOAD_FUNCTIONS.torch_artifacts.name)

Copy link

Choose a reason for hiding this comment

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

Loading old models with torch artifacts crashes

High Severity

Models saved before this refactoring have torch_artifacts as an entry in load_fns. The old loading code would detect and remove this entry before processing. With the new code, LOAD_FUNCTIONS[smash_config.load_fns[0]] raises a KeyError when torch_artifacts is the first entry, since torch_artifacts was removed from LOAD_FUNCTIONS without adding migration handling in load_from_json.

Fix in Cursor Fix in Web

@gsprochette gsprochette requested review from llcnt and simlang January 14, 2026 17:17
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.

2 participants