Fix stateful dataloader state not restored on resume after estimated_stepping_batches#21765
Open
discobot wants to merge 1 commit into
Open
Fix stateful dataloader state not restored on resume after estimated_stepping_batches#21765discobot wants to merge 1 commit into
estimated_stepping_batches#21765discobot wants to merge 1 commit into
Conversation
discobot
requested review from
ethanwharris,
justusschock and
tchaton
as code owners
June 13, 2026 11:15
…_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
force-pushed
the
fix/20550-stateful-dataloader-resume
branch
from
July 18, 2026 14:45
5178d95 to
90a9316
Compare
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.
What does this PR do?
Fixes #20550.
Stateful dataloader states saved in a checkpoint are silently dropped on resume when
configure_optimizersaccessestrainer.estimated_stepping_batches. Building on the analysis inthe issue: the first
setup_data()runs duringstrategy.setup(), beforerestore_training_state(), so when_FitLoop.on_load_checkpointlater stashes the checkpoint'sdataloader states into
_combined_loader_states_to_load, the secondsetup_data()early-returns(
_combined_loaderalready 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 (torchdataStatefulDataLoader-style) — ora
_PrefetchDataFetcherthat already prefetched a batch — would otherwise keep stale pre-restorestate. 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
restartingplus pending states, so it is a no-op everywhere else.Testing: parametrized
test_fit_loop_save_and_restore_dataloaderswith a flag that accessesestimated_stepping_batchesinconfigure_optimizers— the stateful variants fail on master andpass with this change.
tests_pytorch/loops/,trainer/properties/test_estimated_stepping_batches.py, andutilities/test_combined_loader.pypass locally (CPU). CHANGELOG updated.Before submitting
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