ModelCheckpoint overwrites the checkpoint it resumed from when using gradient accumulation #22033

import hashlib, logging, shutil, tempfile, warnings
import torch
import lightning.pytorch as L
from lightning.pytorch.callbacks import ModelCheckpoint
from torch.utils.data import DataLoader

warnings.filterwarnings("ignore"); logging.getLogger("lightning.pytorch").setLevel(logging.ERROR)
root = tempfile.mkdtemp()

class Model(L.LightningModule):
    def __init__(self, crash_at=None):
        super().__init__(); torch.manual_seed(0); self.layer = torch.nn.Linear(4, 1); self.crash_at = crash_at
    def training_step(self, batch, batch_idx):
        if batch_idx == self.crash_at:
            raise RuntimeError("simulated crash")
        return (self.layer(batch) - 3).pow(2).mean()
    def configure_optimizers(self): return torch.optim.SGD(self.parameters(), lr=0.05)
    def train_dataloader(self): return DataLoader(torch.ones(8, 4), batch_size=1)  # 8 identical batches

def fit(name, crash_at=None, ckpt_path=None):
    model = Model(crash_at)
    trainer = L.Trainer(accelerator="cpu", max_epochs=1, accumulate_grad_batches=2, logger=False,
                        enable_progress_bar=False, enable_model_summary=False,
                        callbacks=[ModelCheckpoint(dirpath=f"{root}/{name}", every_n_train_steps=2, save_top_k=-1)])
    try:
        trainer.fit(model, ckpt_path=ckpt_path)
    except RuntimeError:
        pass
    return model.layer.weight.detach().clone()

def describe(path):
    ckpt = torch.load(path, weights_only=False)
    done = ckpt["loops"]["fit_loop"]["epoch_loop.batch_progress"]["current"]["completed"]
    return f"sha={hashlib.sha256(open(path, 'rb').read()).hexdigest()[:10]} global_step={ckpt['global_step']} batches_completed={done}"

print("lightning", L.__version__, "torch", torch.__version__)
w_full = fit("full")                                  # reference: uninterrupted run
fit("run", crash_at=5)                                # crashes after the step-2 checkpoint was written
ckpt = f"{root}/run/epoch=0-step=2.ckpt"
shutil.copy(ckpt, f"{root}/original.ckpt")
print("before resume:", describe(ckpt))
fit("run", crash_at=5, ckpt_path=ckpt)                # resume, crashes again one batch later
print("after resume: ", describe(ckpt))
print("checkpoint we resumed from was overwritten:", open(ckpt, "rb").read() != open(f"{root}/original.ckpt", "rb").read())
w_orig = fit("a", ckpt_path=f"{root}/original.ckpt")
w_over = fit("b", ckpt_path=ckpt)
print("resume from original step-2 file    == uninterrupted:", torch.equal(w_orig, w_full))
print("resume from overwritten step-2 file == uninterrupted:", torch.equal(w_over, w_full))
shutil.rmtree(root)