Lightning-AI / Lightning-AI/pytorch-lightning
ModelCheckpoint Doesn't Delete Old Best Checkpoints When Resuming Training
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
**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
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 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