Store the version inside checkpoints and add tools to upgrade checkpoints - #580
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR introduces checkpoint versioning by adding a version key to default hyperparameters and updating the model loading logic to handle version mismatches.
- Added tests to verify that loading a checkpoint with a mismatched version raises a NotImplementedError.
- Updated the model_from_checkpoint logic to convert version 0 to version 1 and to trigger checkpoint upgrade via the model’s upgrade_checkpoint method.
- Added a new find_architectures_version utility function and updated various model and default-hypers.yaml files to include a default version number.
Reviewed Changes
Copilot reviewed 11 out of 15 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/utils/test_io.py | Added test case verifying error for wrong version checkpoint. |
| src/metatrain/utils/io.py | Updated model loading logic to handle version checking and upgrades. |
| src/metatrain/utils/architectures.py | Introduced find_architectures_version to retrieve current version value. |
| src/metatrain/soap_bpnn/model.py & default-hypers.yaml | Added upgrade_checkpoint stub and version config for SOAP-BPNN model. |
| src/metatrain/pet/model.py & default-hypers.yaml | Added upgrade_checkpoint stub and version config for PET model. |
| src/metatrain/gap/model.py & default-hypers.yaml | Added upgrade_checkpoint stub and version config for GAP model. |
| src/metatrain/deprecated/pet/model.py & default-hypers.yaml | Added upgrade_checkpoint stub and version config for deprecated PET. |
Files not reviewed (4)
- src/metatrain/deprecated/pet/schema-hypers.json: Language not supported
- src/metatrain/gap/schema-hypers.json: Language not supported
- src/metatrain/pet/schema-hypers.json: Language not supported
- src/metatrain/soap_bpnn/schema-hypers.json: Language not supported
DavideTisi
force-pushed
the
versioning
branch
2 times, most recently
from
May 7, 2025 14:28
500c12b to
7d41219
Compare
Luthaf
previously requested changes
May 7, 2025
DavideTisi
force-pushed
the
versioning
branch
2 times, most recently
from
May 9, 2025 07:45
040cc3c to
56c94f3
Compare
4 tasks
4 tasks
Luthaf
force-pushed
the
versioning
branch
5 times, most recently
from
June 24, 2025 15:04
77b628e to
2d5f466
Compare
Luthaf
changed the title
PicoCentauri
left a comment
PicoCentauri
left a comment
Contributor
There was a problem hiding this comment.
This looks very good. I have some minor things that we might do. We could also consider putting the Python code for the Model and the trainer from the new-architecture.rst in a Python file and run it in the test suite. Linking these files in the docs should be easy and should also look good.
| @@ -158,19 +189,26 @@ def save_checkpoint(self, model, path: Union[str, Path]): | |||
| Save a checkoint of both the ``model`` and trainer state to the given ``path`` | |||
Contributor
There was a problem hiding this comment.
Should we do some checks that the version is actually injected or do we not care?
Luthaf
force-pushed
the
versioning
branch
3 times, most recently
from
July 7, 2025 15:08
3a139d4 to
d04e442
Compare
This comment was marked as outdated.
This comment was marked as outdated.
Luthaf
force-pushed
the
versioning
branch
2 times, most recently
from
July 8, 2025 09:04
24b0280 to
ded41ea
Compare
PicoCentauri
approved these changes
Jul 9, 2025
PicoCentauri
left a comment
PicoCentauri
left a comment
Contributor
There was a problem hiding this comment.
Thanks for this important work @DavideTisi and @Luthaf!
Comment on lines
+204
to
+208
| f"Unable to load the model checkpoint from '{path}' for " | ||
| f"the '{architecture_name}' architecture: the checkpoint is using " | ||
| f"version {model_ckpt_version}, while the current version is " | ||
| f"{architecture.__model__.__checkpoint_version__}; and trying to " | ||
| "upgrade the checkpoint failed." |
Contributor
There was a problem hiding this comment.
This is a very motivating error message. TLDR: Please go home, I tried everything I can xD
Luthaf
force-pushed
the
versioning
branch
2 times, most recently
from
July 10, 2025 16:23
9c6da66 to
d06e9cc
Compare
Co-Authored-By: Davide Tisi <davide.tisi93@gmail.com>
Luthaf
added this pull request to the merge queue
Jul 11, 2025
Luthaf
removed this pull request from the merge queue due to a manual request
Jul 11, 2025
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
add checkpoint versioning. Default is 1. each arch have to implement a transfer function to a version to the other, now it just throw an error
Contributor (creator of pull-request) checklist
Reviewer checklist
📚 Documentation preview 📚: https://metatrain--580.org.readthedocs.build/en/580/