fix: proper checkpoint restore weights-only vs full training state by Mani212005 · Pull Request #240 · mllam/neural-lam

Mani212005 · 2026-02-23T13:57:59Z

Describe your changes

Replaces the fragile on_load_checkpoint optimizer hack in ARModel with proper PyTorch
Lightning patterns for checkpoint loading. This separates two distinct behaviours:

Weights-only (--load without --restore_opt): calls ModelClass.load_from_checkpoint(path, ...)
to instantiate the model, then trainer.fit(model, ckpt_path=None) — epoch/step/scheduler reset to 0.
Full resume (--load + --restore_opt): uses the normal constructor, then
trainer.fit(model, ckpt_path=path) — restores optimizer, epoch, LR scheduler, and callbacks.

ar_model.py

Removed self.restore_opt = args.restore_opt attribute.
Removed the on_load_checkpoint block that manually replaced optimizer/scheduler/epoch state — kept only the key-migration logic for older checkpoints.

train_model.py

Updated --restore_opt help text to say "full training state" (not just "optimizer state").
Added assertion: --restore_opt requires --load (fails fast with a clear message).

test_checkpoint_restore.py (new)

test_restore_opt_without_load_raises — assertion fires without --load.
test_weights_only_uses_load_from_checkpoint — load_from_checkpoint called, ckpt_path=None in fit.
test_restore_opt_passes_ckpt_path — normal constructor used, ckpt_path passed to fit.
test_eval_always_passes_ckpt_path — ckpt_path always passed to test.
test_no_load_creates_fresh_model — no --load, normal constructor, ckpt_path=None.

Other test files

Removed the now-unused restore_opt = False from ModelArgs in test_training.py,
test_clamping.py, test_plotting.py, and test_datasets.py.

Behaviour Table

CLI flags	Model instantiation	`trainer.fit` `ckpt_path`
(none)	`ModelClass(args, ...)`	`None`
`--load PATH`	`load_from_checkpoint(PATH, ...)`	`None`
`--load PATH --restore_opt`	`ModelClass(args, ...)`	`PATH`
`--eval test --load PATH`	`load_from_checkpoint(PATH, ...)`	N/A (`trainer.test`)

Issue Link

Closes #131

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
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.

Mani212005 · 2026-02-25T04:47:57Z

@observingClouds @joeloskarsson Kindly review this code. If you think something should be implemented differently. Please let me know

joeloskarsson

Thanks! I think this looks really great. Had a couple small thoughts that I commented on here

joeloskarsson

Sorry for taking so long to look at this again, it drowned below a lot of other things in my inbox :/
Anyway, I think this looks pretty good, just had one documentation suggestion. And unfortunately (since I took so long here) I think the branch needs to be updated with all the recent changes to main, before we can merge. In particular make sure that the changelog entry ends up in the right place with all the additions there.

Mani212005 · 2026-04-20T04:04:41Z

No worries :)
I would do the required changes and push the updates

joeloskarsson · 2026-04-20T19:35:57Z

+    # In tests parse_args may be patched with MagicMock. Normalize the flag
+    # to avoid truthy MagicMock attributes changing control flow.
+    load_training_state = bool(vars(args).get("load_training_state", False))


I don't fully follow this, and this seems quite hard to read now. It would be better to not adapt the code to the testing framework. If testing does not work without this then we need to reconsider the testing approach, not change the actual codebase IMO.

Will the refactoring in #208 make this easier? It should relieve some need for mocking argparse, and it will be merged before this PR.

Mani212005 · 2026-04-23T18:36:42Z

Sorry for taking such a long time to revert. I’ve been busy. would update the required changes on weekend.
Thanks!

joeloskarsson · 2026-05-14T14:18:18Z

#208 is merged now, which means 1) sadly some conflicts and this needing to be adapted, but also 2) maybe #240 (comment) will be easier to sort out?

@Mani212005

…llam#240) Re-applies @Mani212005's PR mllam#240 onto current main. The original PR was pre-mllam#208 and pre-mllam#641; this version applies the same Lightning-native pattern to current main with a much smaller surface area because the machinery is already in place via `load_forecaster_module_from_checkpoint` (introduced as part of the mllam#208 refactor). Source change: - Remove the `if not self.restore_opt: ...` fresh-optimizer hack from `ForecasterModule.on_load_checkpoint`. The hack only reset the optimizer state and silently left epoch / scheduler / callbacks restored from the checkpoint, which is not what `--restore_opt=False` is supposed to mean. - In `train_model.py`, branch on `args.restore_opt` for the `--load` path: weights-only goes through `load_forecaster_module_from_checkpoint` + `trainer.fit(model)` (no `ckpt_path`), full-resume keeps `trainer.fit(..., ckpt_path=args.load)`. Skipped the 239-line tests/test_checkpoint_restore.py from the original PR: it targeted the pre-mllam#208 API and would require a full rewrite. The existing `tests/test_training.py` smoke tests pass on the new code path. Verified: 18 training-related tests pass locally (test_training.py, test_cli.py, test_train_model_warnings.py). Co-Authored-By: Mani <Mani212005@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

@Mani212005

…llam#240) Re-applies @Mani212005's PR mllam#240 onto current main, picking up the naming/rename agreed with @joeloskarsson in the review thread. Source changes: - Rename CLI flag `--restore_opt` -> `--load_training_state` per @joeloskarsson's review comment (the old name was misleading because the flag controls epoch / scheduler / callbacks too, not just the optimizer). The internal `ForecasterModule(restore_opt=...)` keyword stays for now to keep legacy ARModel checkpoints loading via the `args=` shim. - Remove the `if not self.restore_opt: ...` fresh-optimizer hack from `ForecasterModule.on_load_checkpoint`. The hack only reset the optimizer state and silently left epoch / scheduler / callbacks restored from the checkpoint. - In `train_model.py`, branch on `args.load_training_state` for the `--load` path: weights-only goes through `load_forecaster_module_from_checkpoint` + `trainer.fit(model)` with no `ckpt_path`, full-resume keeps `trainer.fit(..., ckpt_path=args.load)`. - Add CLI assertion that `--load_training_state` requires `--load`. Test changes: - Drop the 239-line tests/test_checkpoint_restore.py from the original PR (targeted pre-mllam#208 API; @joeloskarsson also flagged it as adapting code to fit tests rather than testing real behaviour). - Add `test_load_training_state_requires_load` covering the new CLI assertion (parser.error -> SystemExit before any datastore loading). - Update tests/test_train_model_warnings.py mocks to set `load_training_state=False` explicitly (MagicMock returns truthy by default, which would trip the new assertion). Co-Authored-By: Mani <Mani212005@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

sadamov

I rebased on main which made this PR much lighter. @Mani212005 and @joeloskarsson please have a look if this is still in line with your original intent.

joeloskarsson · 2026-06-11T17:20:11Z

+- Fix `--load` to truly restore weights only by default: instead of overriding the optimizer state inside `on_load_checkpoint` (which left epoch / scheduler / callbacks inherited from the checkpoint), `train_model.py` now reconstructs the model via `ForecasterModule.load_from_checkpoint` and calls `trainer.fit(model)` with no `ckpt_path`, so all non-weight training state starts fresh. Also rename the CLI flag from `--restore_opt` to `--load_training_state` to reflect that it controls more than just the optimizer (epoch / scheduler / callbacks too), and assert it requires `--load` [\#240](https://github.com/mllam/neural-lam/pull/240) @Mani212005
+


Please simplify this changelog entry. The point is to get an idea of the change at a glance, not for this to be a full description of the change done (that's what the PR description is for 😄).

joeloskarsson

Great! Yes, I am also happy with this now. Just added a small comment about simplifying the changelog, which can be done before merging this (does not affect code).

Mani212005 force-pushed the model_weights_and_full_training_state branch from 6771b0e to da70355 Compare February 23, 2026 14:01

observingClouds mentioned this pull request Feb 23, 2026

Properly handle difference between restoring only model weights and full training state #131

Open

joeloskarsson requested changes Feb 28, 2026

View reviewed changes

Comment thread tests/test_checkpoint_restore.py Outdated

Comment thread neural_lam/train_model.py Outdated

sadamov requested a review from joeloskarsson March 16, 2026 04:28

sadamov mentioned this pull request Mar 16, 2026

Docs/hello world danra notebook 69 #202

Closed

14 tasks

sadamov added the bug Something isn't working label Apr 13, 2026

joeloskarsson self-assigned this Apr 18, 2026

joeloskarsson requested changes Apr 18, 2026

View reviewed changes

Comment thread neural_lam/train_model.py Outdated

joeloskarsson reviewed Apr 20, 2026

View reviewed changes

sadamov force-pushed the model_weights_and_full_training_state branch from 0ef9cf3 to 73333bd Compare June 6, 2026 19:00

sadamov force-pushed the model_weights_and_full_training_state branch from 73333bd to ccfbbd7 Compare June 6, 2026 19:09

sadamov approved these changes Jun 6, 2026

View reviewed changes

joeloskarsson reviewed Jun 11, 2026

View reviewed changes

joeloskarsson approved these changes Jun 11, 2026

View reviewed changes

joeloskarsson added enhancement New feature or request ready Review complete - proposed for milestone and removed bug Something isn't working labels Jun 11, 2026

joeloskarsson added this to the v0.7.0 (proposed) milestone Jun 11, 2026

joeloskarsson modified the milestones: v0.7.0 (proposed), v0.8.0 (proposed) Jul 17, 2026

		- Fix `--load` to truly restore weights only by default: instead of overriding the optimizer state inside `on_load_checkpoint` (which left epoch / scheduler / callbacks inherited from the checkpoint), `train_model.py` now reconstructs the model via `ForecasterModule.load_from_checkpoint` and calls `trainer.fit(model)` with no `ckpt_path`, so all non-weight training state starts fresh. Also rename the CLI flag from `--restore_opt` to `--load_training_state` to reflect that it controls more than just the optimizer (epoch / scheduler / callbacks too), and assert it requires `--load` [\#240](https://github.com/mllam/neural-lam/pull/240) @Mani212005

Uh oh!

Conversation

Mani212005 commented Feb 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

Issue Link

Type of change

Checklist before requesting a review

Checklist for reviewers

Author checklist after completed review

Checklist for assignee

Uh oh!

Mani212005 commented Feb 25, 2026

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Mani212005 commented Apr 20, 2026

Uh oh!

joeloskarsson Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

joeloskarsson Apr 20, 2026

Choose a reason for hiding this comment

Uh oh!

Mani212005 commented Apr 23, 2026

Uh oh!

joeloskarsson commented May 14, 2026

Uh oh!

sadamov left a comment

Choose a reason for hiding this comment

Uh oh!

joeloskarsson Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Mani212005 commented Feb 23, 2026 •

edited

Loading