Lightning-AI / Lightning-AI/pytorch-lightning

ModelCheckpoint tries to find monitored key on non-zero ranks

Open
#15,852 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

callback: model checkpoint feature logging
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Hi, we're using `ModelCheckpoint` with `save_top_k=2` to keep track of the best checkpoints based on `monitor=val_auc`.

We're computing and logging `val_auc` on `validation_epoch_end`, however we only do it on *rank zero*.

```python
def validation_epoch_end(self, outputs):
gathered_outputs = gloo_gather(outputs)
if rank_zero():
auc = compute_auc(gathered_outputs)
self.log('val_auc': auc)
```

If we run this model over >1 GPUs, we get an error:
```
pytorch_lightning.utilities.exceptions.MisconfigurationException: `ModelCheckpoint(monitor='val_auc')` could not find the monitored key in the returned metrics: ['train_loss', 'val_loss', 'epoch', 'step']. HINT: Did you call `log('val_auc, value)` in the `LightningModule`?
```

This tells us that `ModelCheckpoint` ran on other ranks too, not just on rank zero. This looks like a wrong behaviour - model checkpoint shouldn't run on non-zero ranks, in our understanding.

Is our configuration wrong or is this a bug? Thanks.

Related error was posted already in discussions but not answered: https://github.com/Lightning-AI/lightning/discussions/14806

### How to reproduce the bug

```python
Log a metric only on rank zero and run training with `ModelCheckpoint` callback and `save_top_k=2`.
```

### Error messages and logs

```
pytorch_lightning.utilities.exceptions.MisconfigurationException: `ModelCheckpoint(monitor='val_auc')` could not find the monitored key in the returned metrics: ['train_loss', 'val_loss', 'epoch', 'step']. HINT: Did you call `log('val_auc, value)` in the `LightningModule`?
```

### Environment

```
#- PyTorch Lightning Version (e.g., 1.5.0): 1.7.7
#- PyTorch Version (e.g., 1.10): 1.12.1+cu116
#- Python version (e.g., 3.9): 3.8.10
#- OS (e.g., Linux): Linux
#- How you installed Lightning(`conda`, `pip`, source): pip
```

### More info

_No response_

cc @borda @awaelchli @carmocca @Blaizzy

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 ModelCheckpoint callback entry point and reproduce the reported multi-GPU case with save_top_k=2, monitor=val_auc, and the metric logged only on rank zero. Trace how the callback handles monitored metrics on non-zero ranks. Done means the reproduction no longer raises the missing-key MisconfigurationException while checkpoint monitoring still works.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.