Skip to content

Fix stateful dataloader state not restored on resume after estimated_stepping_batches#21765

Open
discobot wants to merge 1 commit into
Lightning-AI:masterfrom
discobot:fix/20550-stateful-dataloader-resume
Open

Fix stateful dataloader state not restored on resume after estimated_stepping_batches#21765
discobot wants to merge 1 commit into
Lightning-AI:masterfrom
discobot:fix/20550-stateful-dataloader-resume

Conversation

@discobot

Copy link
Copy Markdown

What does this PR do?

Fixes #20550.

Stateful dataloader states saved in a checkpoint are silently dropped on resume when
configure_optimizers accesses trainer.estimated_stepping_batches. Building on the analysis in
the issue: the first setup_data() runs during strategy.setup(), before
restore_training_state(), so when _FitLoop.on_load_checkpoint later stashes the checkpoint's
dataloader states into _combined_loader_states_to_load, the second setup_data() early-returns
(_combined_loader already exists) and _load_combined_loader_states() never runs.

This PR loads the pending states in that early-return branch and re-creates the data-fetcher
iterator. The re-iteration matters: the existing iterator was created before the restore, so
loaders that snapshot their position at iter() time (torchdata StatefulDataLoader-style) — or
a _PrefetchDataFetcher that already prefetched a batch — would otherwise keep stale pre-restore
state. A states-only variant passes the issue's repro but fails an eager-iterator one; loading
then re-iterating matches the order of the normal resume path. The new branch is guarded on
restarting plus pending states, so it is a no-op everywhere else.

Testing: parametrized test_fit_loop_save_and_restore_dataloaders with a flag that accesses
estimated_stepping_batches in configure_optimizers — the stateful variants fail on master and
pass with this change. tests_pytorch/loops/,
trainer/properties/test_estimated_stepping_batches.py, and
utilities/test_combined_loader.py pass locally (CPU). CHANGELOG updated.

Before submitting
  • Was this discussed/agreed via a GitHub issue? (not for typos and docs)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary)
  • Did you write any new necessary tests? (not for typos and docs)
  • Did you verify new and existing tests pass locally with your changes?
  • Did you list all the breaking changes introduced by this pull request?
  • Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist
  • Is this pull request ready for review? (if not, please submit in draft mode)
  • Check that all items from Before submitting are resolved
  • Make sure the title is self-explanatory and the description concisely explains the PR
  • Add labels and milestones (and optionally projects) to the PR so it can be classified

…_stepping_batches`

Accessing trainer.estimated_stepping_batches in configure_optimizers (e.g. for OneCycleLR) runs FitLoop.setup_data() during strategy setup, before the checkpoint's loop state is restored. The second setup_data() call in fit_loop.run() then early-returns because the combined loader already exists, so the stateful dataloader states from the checkpoint are never loaded and resumed training replays already-seen batches. Load the pending dataloader states in the early-return path and recreate the fetcher iterator so the restored state takes effect.

Fixes Lightning-AI#20550
@discobot
discobot force-pushed the fix/20550-stateful-dataloader-resume branch from 5178d95 to 90a9316 Compare July 18, 2026 14:45
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.

stateful dataloaders do not load their state_dict if self.trainer.estimated_stepping_batches called beforehand

1 participant