Lightning-AI / Lightning-AI/pytorch-lightning

Training speed unstable significantly during the first epoch (increasingly accelerates)

Open
#21,343 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.5.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Hi team,

I'm encountering a reproducible phenomenon where the training speed during the first epoch is unstable in a peculiar way. The training process starts very slow for the first ~30% of minibatches and then accelerates significantly, becoming much faster after roughly the 50% mark. This behavior is consistent across runs.

### What is the current behavior?

During the trainer.fit() call, the progress bar for the first training epoch shows a slow iteration speed (e.g., 2.5 it/s) for the initial few hundred batches. After this "warm-up" period, the speed dramatically increases (e.g., 25 it/s) and remains high for the rest of the epoch and subsequent validation stages.

### What is the expected behavior?

I would expect the training speed to be relatively stable throughout the entire first epoch, perhaps with a very brief initial warm-up period of a few batches for things like kernel compilation, but not a sustained slowdown that lasts for a significant portion of the epoch.

### Investigation and what I've ruled out

I have already investigated several common causes for this kind of slowdown:

**I/O and Page Caching**: My initial hypothesis was that the initial slow speed was due to uncached data reads, and the later speed-up was due to the dataset fitting into the system's page cache.
- Test: I isolated my custom Dataset and DataLoader from the Lightning training loop. I manually iterated through the DataLoader for a full epoch and measured the time for each next(iter(dataloader)) call.
- Result: The data fetching time was stable and consistently fast, with no significant slowdown at the beginning. This suggests that the bottleneck is not in the data loading pipeline itself.

**CUDA Kernel Compilation/JIT**: I considered one-time costs like CUDA kernel compilation.
- Reasoning: While some initial overhead is expected, this warm-up period seems far too long, persisting for hundreds of minibatches. A typical kernel compilation warm-up should be much shorter.

**Environment or Model Specificity**: To ensure this wasn't an issue with a particular setup, I have:
- Tried different, simpler models.
- Run the training on different machines.
- Result: The phenomenon (training speed increasingly speed up) is reproducible in all these scenarios.

This leads me to believe the overhead might be originating from within the PyTorch Lightning training loop itself during this initial phase. I'm not sure what could be causing such a sustained "warm-up" period.

I would greatly appreciate any insights, suggestions for further debugging, or pointers on what to investigate next. Has anyone else observed similar behavior?

Thank you for your time and for this fantastic library.

### What version are you seeing the problem on?

v2.5

### Reproduced in studio

_No response_

### 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.5.0): 2.5.5
#- PyTorch Version (e.g., 2.5): 2.9.0
#- Python version (e.g., 3.12): 3.12
#- OS (e.g., Linux): Ubuntu 24.04
#- CUDA/cuDNN version: 12.8
#- GPU models and configuration: might irrelative
#- How you installed Lightning(`conda`, `pip`, source): pip
```

### 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

No reproduction code or failing test is provided. Start at trainer.fit() and compare its timing with the already-tested standalone DataLoader to narrow the slowdown to the Lightning training loop. Done means a minimal reproducible case identifies the responsible behavior and a regression test captures the expected first-epoch timing.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.