Skip to content

Warn when the checkpoint random states cannot be restored - #4361

Open
CJstate wants to merge 1 commit into
huggingface:mainfrom
CJstate:rng-state-warning
Open

CJstate wants to merge 1 commit into
huggingface:mainfrom
CJstate:rng-state-warning

Conversation

@CJstate

@CJstate CJstate commented Oct 2, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes #4283.

load_accelerator_state() restores the RNG state inside a try/except, and reported a failure with logger.info("Could not load random states"). That message is below the effective log level of a default setup: accelerate.checkpointing has no level of its own, so logging inherits WARNING from the root logger and the record is dropped before it reaches any handler. A checkpoint whose RNG file is missing, truncated, or unreadable therefore resumed with no output at all, silently using different random numbers (and a step of 0, since step is stored in the same file) than the interrupted run.

This PR reports the failure at WARNING instead, on every process (each rank restores its own file), and includes the checkpoint path and the underlying error:

input_rng_file = input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl")
try:
    states = load(input_rng_file)
    ...
except Exception as e:
    logger.warning(
        f"Could not load the random states from {input_rng_file}: {e}. Training will resume, "
        "but the random number generators were not restored, so this run may not reproduce "
        "the interrupted one.",
        main_process_only=False,
    )

Resuming from a checkpoint without RNG state still works, it is just not hidden anymore.

Before / after, on the reproduction from #4283

Using the script attached to the issue, unmodified:

# before
RNG restored: False
RNG diverged: True
override_attributes: {}
effective log level: 30          # WARNING, so the INFO message was never emitted
                                 # ... no log line at all
# after
RNG restored: False
RNG diverged: True
override_attributes: {}
effective log level: 30
LOG[WARNING] accelerate.checkpointing: [RANK 0] Could not load the random states from
.../random_states_0.pkl: Weights only load failed. ... Training will resume, but the random
number generators were not restored, so this run may not reproduce the interrupted one.

Note about prior art: #4308 reached the same conclusion (this is the same fix, expressed independently). It was closed without review on 2026-09-25, in the same batch of PRs as #4285, #4239 and ~13 others, and the failure is still present on main. Since #4283 is still open, this PR restores the fix.

Tests

Three regression tests in tests/test_state_checkpointing.py::RandomStateRestoreTest:

  • test_unreadable_rng_state_warns: save a checkpoint, corrupt random_states_0.pkl, resume → exactly one WARNING naming the file, and resuming still returns.
  • test_missing_rng_state_warns: same for an absent RNG file.
  • test_readable_rng_state_restores_without_warning: a healthy checkpoint loads with no warning, restores Python/NumPy/PyTorch RNG state exactly, and restores step.

Verified on Windows 11 / Python 3.13 / PyTorch 2.9.1:

  • Before the fix (git stash of src/accelerate/checkpointing.py): 2 failed, 1 passed — both failures are AssertionError: no logs of level WARNING or higher triggered on accelerate.checkpointing.
  • After the fix: 3 passed.
  • The whole file: 21 passed (no other test relies on the INFO message).
  • ruff check and ruff format --check (0.13.1, the version pinned in setup.py) pass on both files.

Before submitting

Who can review?

@BenjaminBossan @SunMarc (core parts of the library / checkpointing)

`load_accelerator_state` reported a failure to restore the random states with
`logger.info`, which is below the effective log level of a default setup
(`logging` inherits `WARNING` from the root logger). A checkpoint whose RNG
state is missing or unreadable therefore resumed silently, using different
random numbers than the interrupted run.

Log it at `WARNING` instead, on every process since each rank restores its own
file, and include the checkpoint path and the underlying error. Resuming from a
checkpoint without RNG state still works.

Fixes huggingface#4283

Signed-off-by: CJstate <142857225+CJstate@users.noreply.github.com>
CJstate added a commit to CJstate/CJstate that referenced this pull request Oct 2, 2026
@CJstate

CJstate commented Oct 3, 2026

Copy link
Copy Markdown
Author

Small process note on CI: the five pull_request workflows for this branch (c66ff98d) all completed as action_required — they are queued for a maintainer to approve the run for a first-time contributor, so no check has actually reported yet.

Could someone click "Approve and run workflows" when convenient?

The description has the local before/after evidence (before: effective level 30, zero output; after: exactly one LOG[WARNING] accelerate.checkpointing: [RANK 0] Could not load the random states from ...: Weights only load failed.), including why main_process_only=False is required here.

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.

[Bug] RNG state restoration failure swallowed silently — resumed training diverges without warning

1 participant