Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#21,938 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug needs triage ver: 2.6.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.