Skip to content

torch.compile save restore fixes#71

Merged
bonevbs merged 1 commit intomainfrom
tkurth/compile-save-restore-fixes
Mar 27, 2026
Merged

torch.compile save restore fixes#71
bonevbs merged 1 commit intomainfrom
tkurth/compile-save-restore-fixes

Conversation

@azrael417
Copy link
Copy Markdown
Collaborator

torch.compile prepends more prefixes to parameter names causing the current logic to fail loading a checkpoint into a compiled model which was written by an uncompiled model and vice versa. Also, torch.compile breaks some of the logic we had in place for dealing with data parallel prefixes. This MR fixes all those issues by detecting all combinations of those prefixes and making sure to strip them for saving and prepend them for loading.

@azrael417 azrael417 requested a review from bonevbs March 24, 2026 08:50
@azrael417 azrael417 force-pushed the tkurth/compile-save-restore-fixes branch from c673e9d to 531c81e Compare March 26, 2026 07:04
Copy link
Copy Markdown
Collaborator

@bonevbs bonevbs left a comment

Choose a reason for hiding this comment

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

LGTM

@bonevbs bonevbs merged commit c6c68d8 into main Mar 27, 2026
1 check passed
@bonevbs bonevbs deleted the tkurth/compile-save-restore-fixes branch March 27, 2026 09:25
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