Skip to content

Fix: check TrainerState file exists before loading during resume - #39599

Closed
Petecheco wants to merge 2 commits into
huggingface:mainfrom
Petecheco:fix_trainer_state_load_when_file_missing
Closed

Petecheco wants to merge 2 commits into
huggingface:mainfrom
Petecheco:fix_trainer_state_load_when_file_missing

Conversation

@Petecheco

Copy link
Copy Markdown

What does this PR do?

When resuming training from a checkpoint, the Trainer attempts to load `trainer_state.json` to recover the train batch size. However, if the file does not exist, a `FileNotFoundError` is raised, causing resume to fail.

This patch adds a check using `os.path.isfile` before loading the state, consistent with other parts of the codebase. If the file is missing, a warning is logged and batch size recovery is skipped, allowing training to continue.

Fixes a potential crash and improves robustness when resuming from incomplete or custom checkpoint directories.

Related PR #27568

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

trainer: @zach-huggingface, @SunMarc and @qgallouedec
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Petecheco and others added 2 commits July 23, 2025 17:27
When resuming training from a checkpoint, the Trainer attempts to load
\`trainer_state.json\` to recover the train batch size. However, if the file
does not exist, a \`FileNotFoundError\` is raised, causing resume to fail.

This patch adds a check using \`os.path.isfile\` before loading the state,
consistent with other parts of the codebase. If the file is missing, a warning
is logged and batch size recovery is skipped, allowing training to continue.

Fixes a potential crash and improves robustness when resuming from incomplete
or custom checkpoint directories.
@rkb32

rkb32 commented Oct 4, 2026

Copy link
Copy Markdown

Thanks for picking this up. I tested this change with a crash-injection harness and want to share what it did, because the outcome differs from the crash it replaces.

Setup: transformers v5.18.0 (get_last_checkpoint on main is identical), tiny GPT-2, CPU, resume_from_checkpoint=True. A healthy run has checkpoint-5 and checkpoint-10; the newest folder loses only trainer_state.json. That is the last file _save_checkpoint writes, so any kill earlier in the save leaves exactly this state.

  • Current code: FileNotFoundError, the resume stops.
  • With this PR's logic: the warning is logged and training continues, but the Trainer starts at global_step = 0 (the later load of the state is guarded by os.path.isfile, so the step counter, epoch and data position all restart from zero) on the weights of the incomplete folder. The final weights differ from an uninterrupted run (max abs diff 3.0e-4).

Across the 18 kill points I injected one by one in a single save (before and with the file cut in half, at every file event), the current code raises on 17. With this PR, 8 of those 17 turn into the resume-from-step-0 case and 9 still raise (a half-written model/optimizer file fails to load). So it trades a crash for a silent restart in the cases where the model file happens to load.

An alternative that held up in the same test: make get_last_checkpoint skip a checkpoint-N folder whose trainer_state.json is missing or unparseable and fall back to the previous one. With that, all 18 kill points resumed from checkpoint-5 with weights identical to the uninterrupted run (19 of 19 runs including a fault-free control). That is also what #35525 asks for.

Multi-rank caveat (2 gloo ranks): the same change still leaves 3 of 21 kill points, all when a rank other than 0 dies mid-save. Rank 0 writes trainer_state.json without waiting for the other ranks, so the folder looks complete without rng_state_1.pth (and #36076 / #36112 show that having every rank rename a shared temp directory is not the way to fix that).

Limits: tiny model, CPU, Windows, a process kill (not power loss), one run per kill point.

The harness and the exact commands are at https://github.com/rkb32/ckpt-chaos (MIT). I am happy to turn the get_last_checkpoint change into a PR with tests if that would help.

@qgallouedec

Copy link
Copy Markdown
Member

Closing: as rkb32 measured, this turns the crash into a silent restart from step 0. Skipping incomplete checkpoints in get_last_checkpoint (#35525) is the better fix.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants