Lightning-AI / Lightning-AI/pytorch-lightning
First epoch after resuming from checkpoint uses sampler epoch 0 (regression from #20775, num_workers > 0)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### 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:
```python
# 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:
```python
# 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 `IterableDataset`s from skipping data. As a side effect, the epoch-0 iterator from `setup_data()` is now the one consumed by the first resumed epoch:
```python
# 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:
```python
# 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
```python
```
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
Current 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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/pytorch/loops/fit_loop.py, comparing setup_data() with on_advance_start() and the resumed-run condition in _TrainingEpochLoop.on_run_start. Reproduce an end-of-epoch checkpoint resume with num_workers > 0 and inspect the sampler epoch before iterator creation. Done means the first resumed epoch uses the restored sampler epoch while later epochs and IterableDataset behavior remain correct.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100