First epoch after resuming from checkpoint uses sampler epoch 0 (regression from #20775, num_workers > 0)

Bug description

Since 2.5.3, the first training epoch after resuming from an end-of-epoch checkpoint (trainer.fit(..., ckpt_path=...)) iterates the training dataloader with the sampler's epoch 0 permutation instead of the restored epoch, whenever the DataLoader uses worker processes (num_workers > 0). Later epochs of the resumed run are correct.

With use_distributed_sampler=True the Trainer controls DistributedSampler.set_epoch, so users cannot work around this from LightningModule code; every chained/preemptible DDP job replays the epoch-0 shuffle on its first epoch. Combined with limit_train_batches this becomes a data-coverage bug; each restart re-trains on the same fixed slice of the dataset.

Cause

FitLoop.setup_data() creates the training iterator before any sampler epoch is set:

# src/lightning/pytorch/loops/fit_loop.py, setup_data()
self._data_fetcher = _select_data_fetcher(trainer, RunningStage.TRAINING)
self._data_fetcher.setup(combined_loader)
iter(self._data_fetcher)  # creates the iterator inside the fetcher

_set_sampler_epoch(dl, self.epoch_progress.current.processed) only runs later, in FitLoop.on_advance_start(). A multi-process DataLoader iterator pulls its first indices from the sampler at creation, so the permutation is already drawn with sampler.epoch == 0 (a fresh process) by the time set_epoch is called.

Before #20775 this did not matter for map-style datasets, because _TrainingEpochLoop.on_run_start re-created the iterator (after set_epoch) at the start of every epoch of a resumed run:

# 2.5.2
if self.trainer.current_epoch > 0 and not self.restarting:
    iter(data_fetcher)

#20775 (2.5.3, "Fix double iteration bug when resumed from a checkpoint") added and not self.trainer.fit_loop.is_resuming to stop IterableDatasets from skipping data. As a side effect, the epoch-0 iterator from setup_data() is now the one consumed by the first resumed epoch:

# 2.5.3 … master
if self.trainer.current_epoch > 0 and not self.trainer.fit_loop.is_resuming and not self.restarting:
    iter(data_fetcher)

FitLoop.restarting is coerced to False for end-of-epoch checkpoints, so is_resuming is the deciding term.

Proposed fix

Set the sampler epoch before the initial iterator is created, so the fix in #20775 keeps its intent:

# fit_loop.py, setup_data()
for dl in combined_loader.flattened:
    _set_sampler_epoch(dl, self.epoch_progress.current.processed)
self._data_fetcher = _select_data_fetcher(trainer, RunningStage.TRAINING)
self._data_fetcher.setup(combined_loader)
iter(self._data_fetcher)

(epoch_progress.current.processed is already the epoch about to train at that point; it is the value on_advance_start uses.)

What version are you seeing the problem on?

v2.6

Reproduced in studio

https://lightning.ai/sam-lab/templates/first-epoch-after-resuming-from-checkpoint-uses-sampler-epoch-0-regression-from-20775-num-workers-0~01m20cggj23c2dpnztwxz4ab7n

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

#- PyTorch Lightning Version (e.g., 2.6.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

cc @ethanwharris