Skip to content

fix(checkpoint): avoid overwriting checkpoint after resume - #22034

Open
wpyaoxueros wants to merge 2 commits into
Lightning-AI:masterfrom
wpyaoxueros:codex/fix-22033-checkpoint-overwrite
Open

wpyaoxueros wants to merge 2 commits into
Lightning-AI:masterfrom
wpyaoxueros:codex/fix-22033-checkpoint-overwrite

Conversation

@wpyaoxueros

@wpyaoxueros wpyaoxueros commented Oct 19, 2026 •

Copy link
Copy Markdown

What does this PR do?

When resuming with accumulate_grad_batches > 1, ModelCheckpoint starts with an empty last-saved-step marker even though the restored trainer is already at a checkpointed global_step. The first resumed batch can therefore save that same step again before an optimizer update, overwriting the checkpoint with later loop progress. Initialize the guard from the restored trainer step when a checkpoint was loaded.

Fixes #22033

Validation

  • The regression test fails on the base revision: the resumed checkpoint changes from (global_step=2, batches_completed=3) to (2, 4).
  • After the fix, the checkpoint remains (2, 3) and the original CPU-only reproduction shows weights resumed from the saved checkpoint matching uninterrupted training.
  • pytest --confcutdir=tests/tests_pytorch/callbacks -q tests/tests_pytorch/callbacks/test_model_checkpoint_edge_cases.py: 3 passed.
  • Ruff 0.17.1 checks and formatting passed for both changed code files.
  • The PyTorch changelog entry is included in this PR.
  • The full repository suite was not run.

AI disclosure: Codex assisted with reproducing the bug, implementing the fix, and writing the regression test.

This branch has not been deployed

No deployments
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.

ModelCheckpoint overwrites the checkpoint it resumed from when using gradient accumulation

1 participant