[ely2ba]

Addresses the unresolved counter behavior first reported by @mshvartsman in a 24 January 2024 comment on #1601.

What does this PR do?

A fresh process starts with project_configuration.iteration=0. load_state() restores a checkpoint but leaves that counter at zero, including when the caller supplies an explicit path.

With total_limit=2, the next save can delete old checkpoint_0 and reuse its name for the newest state. A later restart loads older checkpoint_1 because its number is higher. In the reproduction, a newly saved weight of 0.7 is replaced on resume by the older 0.8.

Change

After a successful load with automatic checkpoint naming enabled, scan the project's checkpoints directory and advance the counter to:

iteration = max(iteration, newest_existing_checkpoint_number + 1)

• Cover both automatic selection and explicit checkpoint paths, including older checkpoints.
• Use the newest checkpoint already in the project, not just the checkpoint that was loaded.
• Preserve a larger caller-supplied counter.
• Update the counter only after the full load succeeds, including custom objects.

Regression coverage

Test	Checkpoint loaded after restart
test_automatic_loading_after_restart	Latest checkpoint via load_state()
test_explicit_loading_after_restart	Older checkpoint_0 supplied by path

Both tests save twice with total_limit=2, rebuild the accelerator, load, train, and save. They check that the next save is checkpoint_2, then rebuild again and verify that the newly saved parameters are restored.

Both tests cover PyTorch and safetensors model serialization.

Validation

• Both new tests fail without the production fix in both serialization formats: 4 failing regression cases.
• After the fix, tests/test_state_checkpointing.py: 22 passed locally, including both new tests in both serialization formats.
• Ruff lint, Ruff formatting, and git diff --check pass.

Tested on Accelerate main at f47439c, Python 3.12.13, PyTorch 2.15.0, macOS arm64, CPU.

Before submitting

• Read the contributor guidelines.
• Added necessary regression tests.
• Maintainer approval: related discussion is linked above; no new approval is claimed.

Reproduction

Runnable example and observed output before the fix

• Source: Accelerate main at f47439c (1.16.0.dev0).
• Runtime: Python 3.12.13; PyTorch 2.15.0; NumPy 2.5.4; macOS 26.1 arm64.
• Execution: Single process, CPU, no mixed precision. Configuration is supplied in the script; no launcher configuration is required.

import tempfile

import torch
from accelerate import Accelerator
from accelerate.utils import ProjectConfiguration


def start(project_dir):
    """What a (re)started training process does: build everything, then resume if a checkpoint exists."""
    config = ProjectConfiguration(
        project_dir=project_dir, automatic_checkpoint_naming=True, total_limit=2
    )
    accelerator = Accelerator(cpu=True, project_config=config)
    model = torch.nn.Linear(1, 1, bias=False)
    torch.nn.init.ones_(model.weight)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    model, optimizer = accelerator.prepare(model, optimizer)
    return accelerator, model, optimizer


def train_step(accelerator, model, optimizer):
    accelerator.backward(model(torch.ones(1, 1)).sum())
    optimizer.step()
    optimizer.zero_grad()


with tempfile.TemporaryDirectory() as project_dir:
    accelerator, model, optimizer = start(project_dir)
    for _ in range(2):  # checkpoint_0, checkpoint_1
        train_step(accelerator, model, optimizer)
        accelerator.save_state()

    accelerator, model, optimizer = start(project_dir)  # restart 1
    accelerator.load_state()  # "will pick up from the latest checkpoint"
    train_step(accelerator, model, optimizer)
    accelerator.save_state()
    print("saved after restart 1: weight", model.weight.item())

    accelerator, model, optimizer = start(project_dir)  # restart 2
    accelerator.load_state()
    print("resumed after restart 2: weight", model.weight.item())

Observed output before the fix:

[RANK 0] Deleting 1 checkpoints to make room for new checkpoint.
saved after restart 1: weight 0.6999999284744263
resumed after restart 2: weight 0.7999999523162842