Lightning-AI / Lightning-AI/pytorch-lightning

`ModelCheckpoint` deletes *previous run's* checkpoint when remote filesystem

Open
#21,813 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Bug description

`ModelCheckpoint` deletes the *previous run's* checkpoint, including the exact file the trainer resumed from, when the checkpoint dirpath is on a remote (fsspec) filesystem. Authored with the help of an agent but I detected the bug myself.

I continued a training using `Trainer.fit` passing in `ckpt_path`. That run crashed and noticed that the checkpoint I resumed the training from was no longer present in the bucket (R2/S3). It was a `topk` checkpoint. Looking at the code https://github.com/Lightning-AI/pytorch-lightning/blob/4819088b0c6838f0d878b30050d59003305924b8/src/lightning/pytorch/callbacks/model_checkpoint.py#L1007-L1014

it seems intentional since it implies the protection is only for local paths. This PR https://github.com/Lightning-AI/pytorch-lightning/pull/19023 claims to have fixed it but not sure why the guard leaves out non-local file systems. Is it intentional to delete remote, pre-existing checkpoints?
The resuming run subsequently crashed and I cant start it again since the original checkpoint is gone 😬

Sequence:

1. Run 1 trains with `ModelCheckpoint(dirpath="s3://bucket/run1", every_n_train_steps=N, save_top_k=1, monitor=None)` and dies, leaving
`run1/periodic-step=X.ckpt`.
2. Run 2 resumes: `trainer.fit(..., ckpt_path="s3://bucket/run1/periodic-step=X.ckpt")` with a **new** dirpath `s3://bucket/run2`.
3. At run 2's first periodic save, `run1/periodic-step=X.ckpt` is silently deleted. If it was the only copy (`save_top_k=1`), the previous run's data is gone.

Cause chain:

- `ModelCheckpoint.load_state_dict` restores `best_model_path` **unconditionally**, even when it warns that the dirpath changed and skips the other fields.
- For `monitor=None`, `_save_none_monitor_checkpoint` uses that restored `best_model_path` as `previous` and consults `_should_remove_checkpoint`.
- `_should_remove_checkpoint` returns `True` for any non-local `previous` (`if not _is_local_file_protocol(previous): return True`) *before* the two safety guards. Per its own docstring, "not in the current checkpoint directory" and "the checkpoint the Trainer resumed from" only protect when "the filesystem is local".

##### Expected behavior

Same as local: a checkpoint outside the callback's own `dirpath` is never deleted.

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

v2.6

### How to reproduce the bug

Self-contained — `memory://` is a non-local protocol, so no cloud credentials are needed:

```python
import fsspec, torch
from torch.utils.data import DataLoader, TensorDataset
import lightning as L
from lightning.pytorch.callbacks import ModelCheckpoint

class Dummy(L.LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(4, 1)
def training_step(self, batch, _):
x, y = batch
return torch.nn.functional.mse_loss(self.layer(x), y)
def configure_optimizers(self):
return torch.optim.SGD(self.parameters(), lr=0.01)

def loader():
return DataLoader(TensorDataset(torch.randn(512, 4), torch.randn(512, 1)), batch_size=8)

def cb(dirpath):
return ModelCheckpoint(dirpath=dirpath, filename="periodic-step={step}",
every_n_train_steps=5, save_top_k=1,
auto_insert_metric_name=False, save_on_train_epoch_end=False)

def trainer(c, steps):
return L.Trainer(max_steps=steps, callbacks=[c], logger=False, enable_progress_bar=False,
enable_model_summary=False, accelerator="cpu")

c1 = cb("memory://ckpts/run1")
trainer(c1, 12).fit(Dummy(), loader()) # leaves run1/periodic-step=10.ckpt

c2 = cb("memory://ckpts/run2") # resume into a DIFFERENT dirpath
trainer(c2, 20).fit(Dummy(), loader(), ckpt_path=c1.best_model_path)

fs = fsspec.filesystem("memory")
try:
print("run1:", fs.ls("/ckpts/run1", detail=False))
except FileNotFoundError:
print("run1: []")
print("run2:", fs.ls("/ckpts/run2", detail=False))
```

Output:

```
run1: [] # resumed-from checkpoint was deleted
run2: ['/ckpts/run2/periodic-step=20.ckpt']
```

Replace the two `memory://` dirs with local paths and `run1/periodic-step=10.ckpt` survives, as the docstring promises.

```
run1: ['periodic-step=10.ckpt']
run2: ['periodic-step=20.ckpt']
```
run 1's checkpoint survives.

### Environment

Current environment

- lightning: 2.6.0
- pytorch: 2.10.0
- fsspec: 2025.12.0
- Python 3.10, Linux

### More info

[Suggested fix by agent]

The "different directory" guard does not need path resolution for remote URIs — both strings are absolute URIs. Replacing the early return with a normalized containment check fixes it while preserving `save_top_k` retention inside the run's own dir:

```python
if not _is_local_file_protocol(previous):
return previous.startswith(self.dirpath.rstrip("/") + "/")
```

This errs toward keeping (worst case: a stale checkpoint accumulates) rather than deleting.

cc @ethanwharris

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 in src/lightning/pytorch/callbacks/model_checkpoint.py, especially the linked lines around _should_remove_checkpoint, _save_none_monitor_checkpoint, and load_state_dict. Run the self-contained memory:// reproduction to confirm the resumed checkpoint is removed when the callback dirpath changes, then verify that an equivalent local-path run preserves it. Done means the remote-filesystem case matches the documented local safety behavior without breaking checkpoint retention within the active run.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.