Fix stateful dataloader state not restored on resume after estimated_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 #20550

Author: discobot

src/lightning/pytorch/CHANGELOG.md

@@ -32,6 +32,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed `LightningModule.toggle_optimizer` / `untoggle_optimizer` breaking under `torch.compile` by disabling Dynamo tracing on these bookkeeping helpers ([#21513](https://github.com/Lightning-AI/pytorch-lightning/issues/21513))
 
+- Fixed stateful dataloaders not restoring their state from a checkpoint when `trainer.estimated_stepping_batches` is accessed in `configure_optimizers` ([#20550](https://github.com/Lightning-AI/pytorch-lightning/issues/20550))
+
 ---
 
 ## [2.6.4] - 2026-05-20

src/lightning/pytorch/loops/fit_loop.py

@@ -224,6 +224,13 @@ def run(self) -> None:
 
     def setup_data(self) -> None:
         if self._combined_loader is not None and not self._should_reload_train_dl:
+            # the dataloaders may have been set up before the checkpoint was loaded, e.g., if the user accessed
+            # `trainer.estimated_stepping_batches` in `configure_optimizers`. in that case, restore the dataloader
+            # states now and recreate the iterator so the restored states take effect
+            if self.restarting and self._combined_loader_states_to_load:
+                self._load_combined_loader_states()
+                assert self._data_fetcher is not None
+                iter(self._data_fetcher)  # creates the iterator inside the fetcher
             return
 
         trainer = self.trainer

tests/tests_pytorch/loops/test_loops.py

@@ -1248,8 +1248,9 @@ def load_state_dict(self, state_dict):
         ),
     ],
 )
+@pytest.mark.parametrize("access_estimated_stepping_batches", [False, True])
 def test_fit_loop_save_and_restore_dataloaders(
-    train_dataloader_factory, has_state, batches_before, batches_after, tmp_path
+    train_dataloader_factory, has_state, batches_before, batches_after, access_estimated_stepping_batches, tmp_path
 ):
     """Test that the CheckpointConnector saves the state of stateful dataloaders."""
 
@@ -1258,6 +1259,12 @@ def __init__(self):
             super().__init__()
             self.seen_data = []
 
+        def configure_optimizers(self):
+            if access_estimated_stepping_batches:
+                # sets up the dataloaders before the loop state gets restored from the checkpoint (#20550)
+                _ = int(self.trainer.estimated_stepping_batches)
+            return super().configure_optimizers()
+
         def training_step(self, batch, batch_idx):
             self.seen_data.append(batch)
             print(batch)
@@ -1276,8 +1283,9 @@ def train_dataloader(self):
     }
 
     # Train for 2 steps
+    # `max_epochs` is set so that `estimated_stepping_batches` cannot take the infinite-training shortcut
     model = DummyModel()
-    trainer = Trainer(**trainer_kwargs, max_steps=2)
+    trainer = Trainer(**trainer_kwargs, max_steps=2, max_epochs=10)
     trainer.fit(model)
     assert model.seen_data == batches_before
 
@@ -1291,6 +1299,6 @@ def train_dataloader(self):
 
     # Restore training from step 2 and continue 2 more steps
     model = DummyModel()
-    trainer = Trainer(**trainer_kwargs, max_steps=4)
+    trainer = Trainer(**trainer_kwargs, max_steps=4, max_epochs=10)
     trainer.fit(model, ckpt_path=(tmp_path / "checkpoint.ckpt"))
     assert model.seen_data == batches_after