Trainer: after resume_from_checkpoint=<older>, a crash while re-saving checkpoint-N can silently mix two runs #50677

System Info

- `transformers` version: 5.22.2
- Platform: Linux-6.8.0-50-generic-aarch64-with-glibc2.41
- Python version: 3.12.16
- Huggingface_hub version: 2.4.0
- Safetensors version: 0.8.0
- Accelerate version: 1.15.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.15.0+cpu (NA)
- Using distributed or parallel set-up in script?: no (one process, CPU)

• Also reproduced on main at acf8004 (2026-11-12), with the same output. CPU only, single process, local file system.

Who can help?

@SunMarc

Reproduction

Summary. After a rollback, resuming from an older checkpoint-K in the same output_dir, the rolled-back run later saves checkpoint-N into the abandoned run's directory of the same name, file by file, without invalidating it first. A crash during that save can leave a checkpoint whose weights come from one run and whose optimizer, scheduler and trainer state come from the other, and resume_from_checkpoint=True loads it without an error. Until the rolled-back run saves past the abandoned run's newest step, resume_from_checkpoint=True also resumes the abandoned run's newer checkpoints; that is how get_last_checkpoint is meant to work, but nothing warns about it at the rollback.

Both windows are short: the mix needs a crash between the files of one re-save, and the abandoned-run resume needs a crash before the rolled-back run saves past the abandoned run's newest step. A warning at the rollback, before either can happen, would cover both. We would like Trainer to log one when resume_from_checkpoint names a checkpoint-K inside output_dir while higher-numbered checkpoint-* directories exist there, and the docs to say that a rollback should use a new output_dir; we are happy to open a PR for both. Rolling back into a new output_dir avoids both problems today.

Relation to #35525. #35525 asked get_last_checkpoint to skip an incomplete checkpoint and was closed as not planned on 2026-10-27; the maintainers found it "quite niche" and noted that "manual deleting does resolve it". The rollback case is different: a partial checkpoint fails loudly, so the operator knows to delete it, but the mixed checkpoint-N loads without an error, so nothing tells the operator that anything needs deleting. A completeness check would not catch it either, because the abandoned run's trainer_state.json is still in the directory.

Impact. This needs an operator action: trainer.train(resume_from_checkpoint="<output_dir>/checkpoint-K") for a K older than the newest checkpoint, for example to go back before a loss spike or bad data. The Trainer recipes show both steps: resuming a particular checkpoint by path, and resume_from_checkpoint=True to resume an interrupted run from the latest checkpoint in output_dir, which is what the restart of a preempted job uses.

• Silent, mixed checkpoint: when the rolled-back run reaches a step the abandoned run had saved, it writes model.safetensors, then optimizer.pt and scheduler.pt, then trainer_state.json into the existing checkpoint-N. A crash in between leaves, for example, run B's weights with run A's optimizer moments, scheduler and step counters, and the restart loads it without error.
• Loud: a crash in the middle of writing one of these files can leave it truncated, and every restart fails to load until an operator deletes the checkpoint.
• Abandoned run (documented, but no warning): until the rolled-back run saves past the abandoned run's newest step, resume_from_checkpoint=True picks the abandoned run's newest checkpoint (the highest number) and continues the abandoned run, without an error or a warning.

Script. Run A trains on data drawn with seed 0 and saves every 2 steps. Run B is the rollback: it resumes from checkpoint-2 and trains on different data (seed 1). A different learning rate would not show anything here, because Trainer restores the optimizer and scheduler from the checkpoint on resume. Run B is killed (os._exit) at the point named on the command line; the restart is run B's command with resume_from_checkpoint=True.

"""Rolling back with resume_from_checkpoint=<older checkpoint> in the same output_dir.

  python hf_trainer_rollback.py torn        run B dies while re-saving checkpoint-4
  python hf_trainer_rollback.py abandoned   run B dies before its first save

Run A trains on data drawn with seed 0 and saves every 2 steps. Run B is the rollback: it resumes
from checkpoint-2 and trains on different data (drawn with seed 1). A new learning rate would not
do: on resume, Trainer restores the scheduler and optimizer from the checkpoint. Run B is killed
(os._exit) at the point named above. The restart is run B's command with
resume_from_checkpoint=True.
"""
import hashlib
import os
import subprocess
import sys
import tempfile

import torch
from torch import nn
from transformers import (PreTrainedConfig, PreTrainedModel, Trainer, TrainerCallback,
                          TrainingArguments)
from transformers.trainer_utils import get_last_checkpoint


class TinyConfig(PreTrainedConfig):
    model_type = "tiny"


class TinyModel(PreTrainedModel):
    config_class = TinyConfig
    main_input_name = "x"

    def __init__(self, config):
        super().__init__(config)
        self.fc = nn.Linear(8, 1)
        self.post_init()

    def forward(self, x=None, labels=None):
        out = self.fc(x).squeeze(-1)
        return {"loss": nn.functional.mse_loss(out, labels), "logits": out}


class Data(torch.utils.data.Dataset):
    def __init__(self, seed):
        g = torch.Generator().manual_seed(seed)
        self.x, self.y = torch.randn(64, 8, generator=g), torch.randn(64, generator=g)

    def __len__(self):
        return 64

    def __getitem__(self, i):
        return {"x": self.x[i], "labels": self.y[i]}


class Start(TrainerCallback):
    def on_train_begin(self, args, state, control, **kwargs):
        print(f"  training begins at global_step={state.global_step}", flush=True)


def train(out, data_seed, max_steps, resume):
    torch.manual_seed(0)
    args = TrainingArguments(output_dir=out, max_steps=max_steps, save_steps=2, learning_rate=1e-2,
                             per_device_train_batch_size=4, use_cpu=True, report_to="none",
                             logging_strategy="no", disable_tqdm=True)
    Trainer(model=TinyModel(TinyConfig()), args=args, train_dataset=Data(data_seed),
            callbacks=[Start()]).train(resume_from_checkpoint=resume)


def digests(ckpt):
    return {f: hashlib.sha1(open(os.path.join(ckpt, f), "rb").read()).hexdigest()[:8]
            for f in ("model.safetensors", "optimizer.pt", "scheduler.pt", "trainer_state.json")}


if __name__ == "__main__":
    sys.stdout.reconfigure(line_buffering=True)  # keep the order of run A, run B and restart lines
    if sys.argv[1] == "run-b":
        out, mode, max_steps = sys.argv[2], sys.argv[3], int(sys.argv[4])
        if mode == "abandoned":
            real_step = Trainer.training_step

            def training_step(self, *args, **kwargs):
                if self.state.global_step == 3:
                    os._exit(137)
                return real_step(self, *args, **kwargs)

            Trainer.training_step = training_step
        else:
            Trainer._save_optimizer_and_scheduler = lambda self, output_dir: os._exit(137)
        train(out, data_seed=1, max_steps=max_steps, resume=os.path.join(out, "checkpoint-2"))
        sys.exit(0)

    mode = sys.argv[1]
    assert mode in ("torn", "abandoned"), __doc__
    max_steps = 4 if mode == "torn" else 6  # the same for run A and run B
    out = tempfile.mkdtemp()
    print("run A (data seed 0)")
    train(out, data_seed=0, max_steps=max_steps, resume=None)
    print(f"run A saved {sorted(os.listdir(out))}")
    before = digests(os.path.join(out, "checkpoint-4"))
    print("run B (data seed 1), resumed from checkpoint-2")
    rc = subprocess.run([sys.executable, __file__, "run-b", out, mode, str(max_steps)]).returncode
    print(f"run B killed, exit {rc}")
    if mode == "torn":
        after = digests(os.path.join(out, "checkpoint-4"))
        for f in before:
            print(f"  checkpoint-4/{f}: {'rewritten by run B' if before[f] != after[f] else 'still run A'}")
    print(f"restart (run B's command, resume_from_checkpoint=True) picks "
          f"{os.path.basename(get_last_checkpoint(out))}")
    train(out, data_seed=1, max_steps=8, resume=True)
    print("restart finished without error")

Output (transformers 5.22.2; warnings and Trainer's end-of-training metrics lines removed):

$ python hf_trainer_rollback.py torn
run A (data seed 0)
  training begins at global_step=0
run A saved ['checkpoint-2', 'checkpoint-4']
run B (data seed 1), resumed from checkpoint-2
  training begins at global_step=2
run B killed, exit 137
  checkpoint-4/model.safetensors: rewritten by run B
  checkpoint-4/optimizer.pt: still run A
  checkpoint-4/scheduler.pt: still run A
  checkpoint-4/trainer_state.json: still run A
restart (run B's command, resume_from_checkpoint=True) picks checkpoint-4
  training begins at global_step=4
restart finished without error

$ python hf_trainer_rollback.py abandoned
run A (data seed 0)
  training begins at global_step=0
run A saved ['checkpoint-2', 'checkpoint-4', 'checkpoint-6']
run B (data seed 1), resumed from checkpoint-2
  training begins at global_step=2
run B killed, exit 137
restart (run B's command, resume_from_checkpoint=True) picks checkpoint-6
  training begins at global_step=6
restart finished without error

The torn mode stops run B right after save_model (by replacing _save_optimizer_and_scheduler with os._exit), which is where a crash leaves this state.

Expected behavior

• Expected: a crash while the rolled-back run re-saves checkpoint-N leaves the old checkpoint-N, a complete new one, or one that resuming refuses; and resuming an older checkpoint while newer ones exist in output_dir logs a warning.
• Actual: checkpoint-4 can hold run B's weights with run A's optimizer, scheduler and trainer_state.json, and it loads without an error; the rollback itself gives no warning.

Root cause. Links are to main at 9e7f5cd (the same lines at acf8004); the 5.22.2 lines are in parentheses.

• get_last_checkpoint (same lines) returns the highest-numbered checkpoint-* directory, with no notion of which run wrote it; resume_from_checkpoint=True uses it (trainer.py#L1554-L1557; L1537-L1540).
• _save_checkpoint (L3318-L3384) writes into output_dir/checkpoint-{global_step} directly: save_model (line 3369; 3332), then _save_optimizer_and_scheduler (line 3392; 3355), then trainer_state.json (line 3409; 3372). If the directory already exists, its files are replaced one at a time.
• Nothing records the rollback, so the abandoned run's newer checkpoints stay eligible.

Suggested change.

• Warn at the rollback (the cheapest change): when resume_from_checkpoint names a checkpoint-K inside output_dir and higher-numbered checkpoint-* directories exist, log a warning that names them and says that the run will re-save over the ones it reaches and that resume_from_checkpoint=True picks the highest of them after a crash. Add one sentence to the docs saying that a rollback should use a new output_dir. We can open a PR for both.
• Invalidate before overwriting. In _save_checkpoint, when checkpoint-N already exists, remove it or move it aside before writing, so that a crash leaves no checkpoint-N instead of a mixed one. In a multi-process run other ranks also write into that directory (rng_state_{rank}.pth), so this has to happen on one process before the others write, with a barrier.
• Writing to a temporary directory and renaming it covers this too. PR Save checkpoint to temporary directory to handle partial saves during failures #35580 did that; PR Revert checkpoint tmp dir #36112 reverted it after it broke multi-GPU saves (Trainer saving checkpoint crashes in multi-gpu (DDP): OSError: [Errno 39] Directory not empty #36076: several processes moved the temporary directory at the same time), so it needs the same coordination.

Related issues. We found no prior report of the rollback case.

• #35525 "Trainer: Use second last checkpoint if last checkpoint loading fails" (closed as not planned, 2026-10-27): see above.
• #39599 "Fix: check TrainerState file exists before loading during resume" (closed unmerged, 2026-10-17): it turned the crash into a silent restart from step 0.
• #35782 (closed as completed): partial checkpoints after an unattended crash. Addressed by PR #35580, which PR #36112 (merged 2025-02-10) reverted.
• Earlier history: PR #27820 ("fix: non-atomic checkpoint save", merged), PR #28364 ("Fix for checkpoint rename race condition", merged), PR #29370 ("Fully revert atomic checkpointing", merged).

How this was found. We found this with a crash-injection tool that kills the process with SIGKILL just before or just after each file operation of a checkpoint save, restarts it and checks, through per-run markers in the saved state, which run each restored file came from. The standalone script above reproduces it without that tool.

AI assistance

This report, its reproduction script and the analysis were prepared with an AI coding agent. I ran the reproduction on the versions listed above and checked the analysis and the cited code.