Skip to content

feat: add checkpoint migration registry#701

Open
Sharkyii wants to merge 6 commits into
mllam:mainfrom
Sharkyii:checkpoint-migration-registry
Open

feat: add checkpoint migration registry#701
Sharkyii wants to merge 6 commits into
mllam:mainfrom
Sharkyii:checkpoint-migration-registry

Conversation

@Sharkyii

@Sharkyii Sharkyii commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

Describe your changes

Old checkpoints fail to load after model refactors because there's no structured way to evolve the state dict schema. This adds a small migration registry (
migrations.py
) that applies ordered, versioned transformations to a checkpoint's state dict on load — replacing the two ad-hoc rename blocks that already existed in on_load_checkpoint.

Schema versioning uses a monotonic integer instead of the package semver, since Hatch's git-tag-derived versioning makes version unreliable in dev installs.

The two existing migrations ( prefix remap + grid_mlp rename) are the first entries in the registry. A standalone python -m neural_lam.migrate_checkpoint script lets you upgrade a checkpoint file on disk rather than migrating in memory on every load.

Dependencies: none beyond what's already in the project.

Issue Link

#48

Type of change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes
  • I have added tests that prove my fix is effective or that my feature works
  • [ x I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section
    reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

@Sharkyii Sharkyii force-pushed the checkpoint-migration-registry branch from 4f831a2 to 1b79303 Compare July 6, 2026 10:57
@Sharkyii Sharkyii changed the title Add Checkpoint migration registry feat: add checkpoint migration registry Jul 6, 2026
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