Skip to content

[DOC] Document checkpoint resuming on lightning 2.6 and use pytorch_forecasting Tuner in examples - #2464

Open
Elvand-Lie wants to merge 6 commits into
sktime:mainfrom
Elvand-Lie:fix/resume-ckpt-weights-only
Open

Elvand-Lie wants to merge 6 commits into
sktime:mainfrom
Elvand-Lie:fix/resume-ckpt-weights-only

Conversation

@Elvand-Lie

@Elvand-Lie Elvand-Lie commented Oct 13, 2026 •

Copy link
Copy Markdown

Fixes #2463. Related: #2000 and #2027.

Change

On Lightning >=2.6 with PyTorch >=2.6, Trainer.fit(..., ckpt_path=...) rejects v1 checkpoints because the default weights-only loader cannot restore the fitted encoders, normalizers, losses and metrics they contain.

  • Add a "Resuming from a checkpoint" section to the Models usage docs: resume with ckpt_path, and pass weights_only=False for trusted checkpoints. The FAQ entry now points there; the stallion tutorial note stays.
  • Switch the README.md, Getting Started and Models examples from lightning.pytorch.tuner.Tuner to pytorch_forecasting.tuning.Tuner. Lightning's Tuner.lr_find has no weights_only argument and fails on these checkpoints; ours patches the loader ([BUG] Solve the unpickling error in weight Loading #2000). The tutorials already use it.
  • Integration test: train one epoch, save, resume for a second epoch, then validate and test from that checkpoint. An epoch-counting callback verifies that only the second epoch runs after resuming. Skipped on Lightning <2.6, where the keyword is unavailable.

Scope

Documentation, example imports and a test of an existing workaround. Runtime loading behavior is unchanged; default loading still fails without the explicit keyword. The pytorch_forecasting.tuning.Tuner wrapper is still needed: in Lightning 2.6.6, the latest release, Tuner.lr_find has no weights_only argument. v2 uses a separate representation and is out of scope. Only trusted checkpoints should be loaded this way, as unpickling can execute code.

Validation

Lightning 2.6.6, PyTorch 2.14.1, Windows CPU:

  • New regression passes (1 passed); skipped on Lightning 2.5.6.
  • TFT, NBeats and DeepAR reproduce the default loading error for fit/validate/test; explicit weights_only=False passes.
  • TFT: Lightning's Tuner.lr_find fails with the same UnpicklingError and pytorch_forecasting.tuning.Tuner.lr_find passes. The LearningRateFinder callback fails even with trainer.fit(..., weights_only=False). A plain LightningModule checkpoint resumes with the default.
  • Pre-commit passes on all changed files. A standalone Sphinx build renders the new section and resolves the FAQ link; the full docs build was not run locally; the PR's Read the Docs build passes.
  • Model suite on the earlier base 0acb5a0: 201 passed and 5 skipped across a parallel run and serial retries of 23 environment failures (Tk backend, memory).

Review focus

Placement of the note in the Models usage section, and whether the Tuner import changes belong in this PR.

Document weights_only=False for trusted v1 checkpoints on Lightning >=2.6. Cover resume, validate and test through ckpt_path with an integration test.

Fixes sktime#2463
@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@codecov

codecov Bot commented Oct 13, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@c35b938). Learn more about missing BASE report.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2464   +/-   ##
=======================================
  Coverage        ?   88.80%           
=======================================
  Files           ?      218           
  Lines           ?    11688           
  Branches        ?        0           
=======================================
  Hits            ?    10380           
  Misses          ?     1308           
  Partials        ?        0           
Flag Coverage Δ
cpu 88.80% <ø> (?)
pytest 88.80% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@phoeenniixx phoeenniixx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I think we just need to pass weights_only=False? That should be enough no?
Also, I feel this is something that should be mentioned in lightning docs rather than ours, no?

But I think we can mention it in our docs, but it should be in a more "visible" place than nb - maybe in models or somewhere?

@phoeenniixx phoeenniixx added the documentation Improvements or additions to documentation label Oct 14, 2026
@Elvand-Lie

Elvand-Lie commented Oct 14, 2026 •

Copy link
Copy Markdown
Author

I think passing weights_only=False is what this PR documents: the FAQ entry and the tutorial note say exactly that, and the test pins it. So we agree on the fix. The PR is where users find it.

It is not enough everywhere, though. On current main (lightning 2.6.6, torch 2.14.1):

  • Lightning's Tuner.lr_find, used in README.md and models.rst, has no weights_only argument and fails with the same error.
  • The LearningRateFinder callback fails even with trainer.fit(..., weights_only=False).

On Lightning docs: a plain LightningModule checkpoint resumes with no flag. Only ours fail, because of the encoders, normalizers and metrics we pickle into the checkpoint. We already handle this ourselves in load_from_checkpoint and our Tuner (#2000, #2027).

On placement I do agree. The PR also adds a faq.rst entry, not only the notebook. I'll move the note into the usage section of models.rst. That example also fails at lr_find on lightning 2.6. I can switch it to our Tuner here or in a separate PR, I will do which is best and is helpful.

@phoeenniixx

Copy link
Copy Markdown
Member

I think passing weights_only=False is what this PR documents

I meant do we need to pass ckpt_path at all? Only passing weights_only=False should suffice, no?

Regarding lr_find, I think we need to use this Tuner and then, it should not fail! I think it is not correctly documented anywhere though.

I have a question here though - we wrapped the Tuner because there was no way to actually pass weights_only in Tuner when lightning 2.5 was released. I have not looked at the recent releases, is this issue solved? I think then we can safely remove this Tuner?

…casting Tuner

Add a "Resuming from a checkpoint" section to the Models usage docs and reduce the FAQ entry to a pointer. Switch the README, Getting Started and Models examples to pytorch_forecasting.tuning.Tuner, as Lightning's Tuner.lr_find cannot pass weights_only and fails on v1 checkpoints with lightning 2.6.
@Elvand-Lie Elvand-Lie changed the title [DOC] Document trusted checkpoint loading for Trainer ckpt_path [DOC] Document checkpoint resuming on lightning 2.6 and use pytorch_forecasting Tuner in examples Oct 15, 2026
@Elvand-Lie

Copy link
Copy Markdown
Author

weights_only=False controls how a checkpoint is loaded; it does not select one. ckpt_path selects the checkpoint to resume from. Without it, fit trains the in-memory model and there is nothing to load. Lightning can also select a checkpoint implicitly, e.g. validate/test without a model use the best one.

The wrapper is still needed. Lightning 2.6.6 is the latest release; Tuner.lr_find still has no weights_only argument, and its temporary checkpoint restore does not pass one. On a TFT, Lightning's Tuner.lr_find fails with the same UnpicklingError; ours passes.

Updated in a690577:

  • note moved to a "Resuming from a checkpoint" section in the Models usage docs, with a short FAQ pointer
  • README.md, Getting Started and Models examples now use pytorch_forecasting.tuning.Tuner

@phoeenniixx phoeenniixx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Why do we have a test addition in a doc change PR?

@Elvand-Lie

Copy link
Copy Markdown
Author

The test checks the workflow documented here: resuming through Trainer.fit with ckpt_path and weights_only=False, then validating and testing from that checkpoint. The existing v1 model checkpoint-loading tests use load_from_checkpoint; I found no other test of Trainer's ckpt_path route.

It passes on main because this PR documents an existing workaround. Keeping the test helps catch future pytorch-forecasting or Lightning changes that break that documented workflow in CI.

The earlier review points are also addressed:

  • The detailed explanation is now in a "Resuming from a checkpoint" section in the Models usage docs, with a short FAQ pointer. The notebook note is retained (a690577).
  • The README.md, Getting Started and Models examples now use pytorch_forecasting.tuning.Tuner (a690577).
  • The Tuner wrapper remains needed: in Lightning 2.6.6, the latest PyPI release, Tuner.lr_find still has no weights_only argument.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] trainer.fit(ckpt_path=...) cannot resume v1 models on lightning>=2.6 with torch>=2.6 (UnpicklingError, weights_only)

2 participants