Lightning-AI / Lightning-AI/pytorch-lightning
Display `max_epoch=xxx reached` message when EarlyStopping triggers
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
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
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 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