Lightning-AI / Lightning-AI/pytorch-lightning

ModelCheckpoint Doesn't Delete Old Best Checkpoints When Resuming Training

Open
#18,687 3 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug callback: model checkpoint repro needed ver: 1.9.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

**Description**:
When using `ModelCheckpoint` with the parameters `top_k=1` and `monitor='val_loss'` during a singular training run, the behavior is as expected and only retains one 'best_val_confidence-epoch...' checkpoint.

However, in the context of cloud-based training where instances may be preempted or restarted from a checkpoint:

- The training resumes from a checkpoint labeled "last.ckpt", which was initially created by a different `ModelCheckpoint`.
- There aren't any explicit warnings indicating that the `ModelCheckpoint` state was restored incorrectly.
- Post-resumption, `ModelCheckpoint` creates a new checkpoint but fails to delete the old one. Thus, if there's a single preemption/restart during the training run, we end up with two 'best_val_loss' checkpoints.

It should be noted we load/write checkpoints to GCS with `fsspec`, which allows for checkpoints to be written to and loaded directly from Google Cloud Storage (GCS).

**Code Details**:

There are two current `ModelCheckpoint` callbacks in use:

1. The first is for saving the latest checkpoint:

```python
last_ckpt_callback = ModelCheckpoint(
save_top_k= -1,
save_last= True,
dirpath= self.checkpoint_dir,
)
last_ckpt_callback.CHECKPOINT_NAME_LAST = _CHECKPOINT_NAME_LAST
```

2. The second is for saving the best validation loss checkpoint:

```python
best_val_loss_ckpt_callback = ModelCheckpoint(
monitor=f'val_loss',
mode='min',
save_top_k=1,
auto_insert_metric_name=False,
filename='best_val_confidence-epoch{epoch}-val_loss{{val_loss:.4e}}',
dirpath=self.checkpoint_dir,
)
```
Environment:
- Lightning Component: ModelCheckpoint object
- PyTorch Lightning Version: 1.9.2
- PyTorch Version: 1.13.0
- Python Version: 3.10.12
- OS: Linux
- CUDA/cuDNN version: Build cuda_11.6.r11.6/compiler.31057947_0
- GPU models: Nvidia A100
- How you installed Lightning: Conda
- Cloud: Running on GCP Cluster

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

v1.9

### How to reproduce the bug

```python
1. Setup a training loop on the cloud with the aforementioned `ModelCheckpoint` callbacks.
2. Intentionally interrupt the training to simulate preemption.
3. Resume the training from the "last.ckpt".
4. Post-resumption, inspect the stored checkpoints. There should be two 'best_val_loss' checkpoints instead of one.

**Expected behavior**: Only one 'best_val_confidence-epoch...' checkpoint should remain after resumption.

**Actual behavior**: Multiple 'best_val_confidence-epoch...' checkpoints are observed after training preemption and resumption.
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
- Lightning Component: ModelCheckpoint object
- PyTorch Lightning Version: 1.9.2
- PyTorch Version: 1.13.0
- Python Version: 3.10.12
- OS: Linux
- CUDA/cuDNN version: Build cuda_11.6.r11.6/compiler.31057947_0
- GPU models: Nvidia A100
- How you installed Lightning: Conda
- Running environment of LightningApp: Cloud, Running on GCP A100 instance
```

### More info

_No response_

cc @carmocca @awaelchli

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 ModelCheckpoint state-restoration and checkpoint-deletion behavior, using the two callbacks and the last.ckpt resumption sequence described in the issue. Reproduce the interruption and restart with fsspec-backed GCS storage; done means resumption leaves only one best_val_confidence checkpoint.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.