Lightning-AI / Lightning-AI/pytorch-lightning

Display `max_epoch=xxx reached` message when EarlyStopping triggers

Open
#21,031 6 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

When EarlyStopping callback signals Trainer to stop, the message displayed is always wrong.
## How to reproduce
Any case where EarlyStopping is triggered can reproduce this.
## Cause
In fit_loop.done, the order of if statements are wrong.
```python
@property
def done(self) -> bool:
"""Evaluates when to leave the loop."""
if self.max_batches == 0:
rank_zero_info("`Trainer.fit` stopped: No training batches.")
return True

# TODO: Move track steps inside training loop and move part of these condition inside training loop
stop_steps = _is_max_limit_reached(self.epoch_loop.global_step, self.max_steps)
if stop_steps:
rank_zero_info(f"`Trainer.fit` stopped: `max_steps={self.max_steps!r}` reached.")
return True

# `processed` is increased before `on_train_epoch_end`, the hook where checkpoints are typically saved.
# we use it here because the checkpoint data won't have `completed` increased yet
assert isinstance(self.max_epochs, int)
stop_epochs = _is_max_limit_reached(self.epoch_progress.current.processed, self.max_epochs)
if stop_epochs:
# in case they are not equal, override so `trainer.current_epoch` has the expected value
self.epoch_progress.current.completed = self.epoch_progress.current.processed
rank_zero_info(f"`Trainer.fit` stopped: `max_epochs={self.max_epochs!r}` reached.")
return True

if self.trainer.should_stop and self._can_stop_early:
rank_zero_debug("`Trainer.fit` stopped: `trainer.should_stop` was set.")
return True

return False
```
## Solution
Modify the order of if statements.

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

v2.5

### Reproduced in studio

_No response_

### How to reproduce the bug
#### Code of Module
```python
class IrisClassifier(MLPModule):
def __init__(self, **kwargs):
super().__init__()
self.model = torch.nn.Sequential(
torch.nn.Linear(4, 16), torch.nn.ReLU(), torch.nn.Linear(16, 3)
)
self.loss_fn = torch.nn.CrossEntropyLoss()
self.automatic_optimization = False # Disable automatic optimization

def on_train_start(self):
"""Check if running in a distributed environment."""
comprehensive_distributed_check()
print("distributed strategy:", self.trainer.strategy)

def forward(self, x):
return self.model(x)

def training_step(self, batch, batch_idx):

x = torch.stack(
[
batch["sepal_length"],
batch["sepal_width"],
batch["petal_length"],
batch["petal_width"],
],
dim=1,
).float() # shape: [batch, 4]

species_map = {"setosa": 0, "versicolor": 1, "virginica": 2}
y = torch.tensor(
[species_map[s] for s in batch["species"]],
dtype=torch.long,
device=self.device,
)

logits = self(x)
loss = self.loss_fn(logits, y)
acc = (logits.argmax(dim=1) == y).float().mean()

opt = self.optimizers()
opt.zero_grad()

self.manual_backward(loss)

opt.step()

return loss

def validation_step(self, batch, batch_idx):
x = torch.stack(
[
batch["sepal_length"],
batch["sepal_width"],
batch["petal_length"],
batch["petal_width"],
],
dim=1,
).float()
species_map = {"setosa": 0, "versicolor": 1, "virginica": 2}
y = torch.tensor(
[species_map[s] for s in batch["species"]],
dtype=torch.long,
device=self.device,
)

logits = self(x)
loss = self.loss_fn(logits, y)
acc = (logits.argmax(dim=1) == y).float().mean()
print(f"on Rank{self.trainer.global_rank} the batch_idx is {batch_idx} the val_acc is {acc}")
self.log("Validation/val_loss", loss, prog_bar=True)
self.log("Validation/val_acc", acc, prog_bar=True, sync_dist=True)
```
####
Configuration I used
```
max_epochs=1,
EarlyStopping(monitor="Validation/val_acc", mode="max", min_delta=100, patience=2),
val_check_interval=2,
limit_val_batches=10,
```

### Error messages and logs

```
# Error messages and logs here please
Monitored metric Validation/val_acc did not improve in the last 2 records. Best score: 0.317. Signaling Trainer to stop.
`Trainer.fit` stopped: `max_epochs=1` reached.
```
The message is supposed to be `Trainer.fit` stopped: `trainer.should_stop` was set.

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.5.0): 2.5.1.post0
#- PyTorch Version (e.g., 2.5): 2.7.0
#- Python version (e.g., 3.12): 3.9.6
#- OS (e.g., Linux): macOS
#- CUDA/cuDNN version: null
#- GPU models and configuration: null
#- How you installed Lightning(`conda`, `pip`, source): uv pip
```

### More info

_No response_

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 at the Trainer fit_loop.done property shown in the issue and inspect the stop-condition order around max_epochs and trainer.should_stop. Reproduce with EarlyStopping using min_delta=100, patience=2, max_epochs=1, val_check_interval=2, and limit_val_batches=10. Done means an EarlyStopping-triggered run reports that trainer.should_stop was set rather than max_epochs being reached.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.