Skip to content

Fix stale checkpoint cleanup in save_pretrained - #14774

Open
kyo-zzz wants to merge 1 commit into
huggingface:mainfrom
kyo-zzz:fix/save-pretrained-stale-checkpoints
Open

kyo-zzz wants to merge 1 commit into
huggingface:mainfrom
kyo-zzz:fix/save-pretrained-stale-checkpoints

Conversation

@kyo-zzz

@kyo-zzz kyo-zzz commented Sep 15, 2026 •

Copy link
Copy Markdown

Fixes #14769.

ModelMixin.save_pretrained cleans a previous save by removing files that match the shard pattern of the current save, but it only recognized sharded files. Single-file checkpoints of the other container were left behind: re-saving a .safetensors checkpoint as .bin (or the reverse) kept the old file on disk, and because from_pretrained prefers sharded and safetensors checkpoints, the stale one silently took precedence over the freshly written weights — no error, no warning.

Relationship to #14844

#14844 fixed the sibling problem: cleanup matched another variant's shards through a bare prefix, and a formerly-sharded checkpoint left a stale index behind. It scopes the match to the container being written (filename.removesuffix(ext)), so a .safetensors → .bin re-save still leaves the old .safetensors on disk.

This PR is rebased on top of it and keeps that anchoring, then extends the cleanup to both containers.

Changes

  • Extract the cleanup into _get_superseded_checkpoint_files, which returns every artifact of the current save's variant that the new save replaces: single-file weights, shards, and sharded indexes, in either container.
  • Scope is the variant stem (diffusion_pytorch_model vs diffusion_pytorch_model.ema), so checkpoints of other variants that share the directory are preserved. Shards are matched through _REGEX_SHARD group 1 rather than a prefix, which keeps the fix: scope save_pretrained shard cleanup to the saved variant (avoid deleting another variant's shards) #14844 regression fixed.
  • Sharded indexes are cleaned as well, which fixes the stale-index FileNotFoundError case from the issue matrix.

Tests

tests/models/test_modeling_common.py::TestModelUtils

Verified red-then-green against the current main: reverting only src/diffusers/models/modeling_utils.py to main makes both new cases fail (assert 'diffusion_pytorch_model.ema.safetensors' not in [...]) while the #14844 tests still pass; with the fix applied all four pass. ruff check and ruff format --check are clean.

@github-actions github-actions Bot added fixes-issue models tests size/M PR with diff < 200 LOC and removed fixes-issue labels Sep 15, 2026
`ModelMixin.save_pretrained` cleaned a previous save by removing files whose
name matched the shard pattern of the current save, but it only recognized
sharded files. Single-file checkpoints of the other container were left
behind: re-saving a `.safetensors` checkpoint as `.bin` (or the reverse) kept
the old file on disk, and because `from_pretrained` prefers sharded and
safetensors checkpoints, the stale one silently took precedence over the
freshly written weights.

Extract the cleanup into `_get_superseded_checkpoint_files`, which returns
every artifact of the current save's variant that the new save replaces:
single-file weights, shards, and sharded indexes, in either container. Files
belonging to other variants (e.g. `ema`) share the directory and are left
untouched.

Add a regression test covering both directions of the container switch and
the variant coexistence case.

Fixes huggingface#14769
@kyo-zzz
kyo-zzz force-pushed the fix/save-pretrained-stale-checkpoints branch from 1c895ee to 055e979 Compare October 31, 2026 03:27

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

save_pretrained(safe_serialization=False) leaves the old safetensors checkpoint behind, and from_pretrained loads it instead of the new weights

1 participant