Lightning-AI / Lightning-AI/pytorch-lightning

Validation metrics not available when resuming training from checkpoint

Open
#18,595 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Bug description

I am training a lot of models in a preemtible environment where training is often stopped unexpectedly due to other jobs with higher priority requesting the GPU the model is being trained on. In these cases, I deploy resume jobs which resume from the last model checkpoint as follows:
```
trainer = Trainer(...)
datamodule = MyPLDataModule(...)
checkpoint_path = ...
module = MyPLModule.load_from_checkpoint(checkpoint_path)
trainer.fit(module, datamodule=datamodule, ckpt_path=checkpoint_path)
```

In ~90% of cases, this works perfectly fine and training resumes as expected. However, in the remaining 10%, training crashes as soon as the EarlyStoppingCallback is invoked since the monitored metric is not available.
Note that I am using a modified version of the EarlyStoppingCallback that allows for a warm-up grace-period:
```
class EarlyStoppingWithWarmUp(EarlyStopping):
def __init__(
self,
monitor: str,
min_delta: float = 0.0,
patience: int = 3,
verbose: bool = False,
mode: str = "min",
strict: bool = True,
check_finite: bool = True,
stopping_threshold: Optional[float] = None,
divergence_threshold: Optional[float] = None,
check_on_train_epoch_end: Optional[bool] = None,
log_rank_zero_only: bool = False,
warmup_epochs: int = 0,
):
super().__init__(
monitor,
min_delta,
patience,
verbose,
mode,
strict,
check_finite,
stopping_threshold,
divergence_threshold,
check_on_train_epoch_end,
log_rank_zero_only,
)
self._warmup_epochs = warmup_epochs

def _is_warmup_epoch(self, trainer: pl.Trainer) -> bool:
return trainer.current_epoch < self._warmup_epochs

def _should_skip_check_due_to_warmup(self, trainer: pl.Trainer) -> bool:
return self._is_warmup_epoch(trainer)

def on_train_epoch_end(
self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
logs = trainer.callback_metrics
if (
trainer.fast_dev_run
or not self._validate_condition_metric( # disable early_stopping with fast_dev_run
logs
)
): # short circuit if metric not present
return
if self.check_finite:
# If the monitor is not finite, run the check anyway to stop training
current = logs[self.monitor].squeeze()
if not torch.isfinite(current):
return super().on_train_epoch_end(trainer, pl_module)
if self._should_skip_check_due_to_warmup(trainer):
return
return super().on_train_epoch_end(trainer, pl_module)

def on_validation_end(
self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
if self._should_skip_check_due_to_warmup(trainer):
return
return super().on_validation_end(trainer, pl_module)

```
Note that this custom callback does not produce any issues when training a model from start to finish without resuming.

Any insights would be appreciated!

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

v2.0

### How to reproduce the bug

_No response_

### Error messages and logs

```
Early stopping conditioned on metric `val_TimeSeriesMAE` which is not available. Pass in or modify your `EarlyStopping` callback to use any of the following: `train_loss`, `train_TimeSeriesMAE`, `train_TimeSeriesMAPE`, `train_TimeSeriesRMSE`, `train_TimeSeriesMSE`
Traceback (most recent call last):
File "/app/pbim_models/cli/resume.py", line 199, in resume
trainer.fit(
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 529, in fit
call._call_and_handle_interrupt(
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 568, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 973, in _run
results = self._run_stage()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1016, in _run_stage
self.fit_loop.run()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
self.on_advance_end()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 369, in on_advance_end
call._call_callback_hooks(trainer, "on_train_epoch_end", monitoring_callbacks=True)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 193, in _call_callback_hooks
fn(trainer, trainer.lightning_module, *args, **kwargs)
File "/app/models/util.py", line 110, in on_train_epoch_end
or not self._validate_condition_metric( # disable early_stopping with fast_dev_run
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/callbacks/early_stopping.py", line 146, in _validate_condition_metric
raise RuntimeError(error_msg)
RuntimeError: Early stopping conditioned on metric `val_TimeSeriesMAE` which is not available. Pass in or modify your `EarlyStopping` callback to use any of the following: `train_loss`, `train_TimeSeriesMAE`, `train_TimeSeriesMAPE`, `train_TimeSeriesRMSE`, `train_TimeSeriesMSE`
```

### Environment

Current environment
* Lightning:
- lightning-utilities: 0.9.0
- pytorch-lightning: 2.0.6
- torch: 2.0.0
- torch-cluster: 1.6.1
- torchmetrics: 0.11.4

### More info

_No response_

cc @lantiga

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 with the trainer.fit resume path and the EarlyStopping._validate_condition_metric call shown in the traceback, comparing normal training with training resumed from a checkpoint. The report names no repository test or reproducible setup; the work is done when the resumed case reliably handles the monitored validation metric without the reported crash.

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
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.