Skip to content

ModelCheckpoint deletes *previous run's* checkpoint when remote filesystem #21813

Description

@parhamfh

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

def _should_remove_checkpoint(self, trainer: "pl.Trainer", previous: str, current: str) -> bool:
"""Checks if the previous checkpoint should be deleted.
A checkpoint won't be deleted if any of the cases apply:
- The previous checkpoint is the same as the current checkpoint (means the old was already overwritten by new)
- The previous checkpoint is not in the current checkpoint directory and the filesystem is local
- The previous checkpoint is the checkpoint the Trainer resumed from and the filesystem is local

it seems intentional since it implies the protection is only for local paths. This PR #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:

Details
  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:

  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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions