Skip to content

torch.compile save restore fixes#71

Open
azrael417 wants to merge 22 commits intomainfrom
tkurth/compile-save-restore-fixes
Open

torch.compile save restore fixes#71
azrael417 wants to merge 22 commits intomainfrom
tkurth/compile-save-restore-fixes

Conversation

@azrael417
Copy link
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
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.

1 participant