[Thanatos9404]

What does this PR do ?

Raise a failed background checkpoint finalization at the start of the next training step instead of at the next save.

Issues

Fixes #4707

Usage

N/A (internal). No config or API changes for users.

Problem

With background finalization, a checkpoint whose async write or tmp_step_N to step_N rename fails does not stop the run. Training keeps going until the next save, which is up to save_period steps later, and only then raises. By that point the newest visible step_N is older than the step the run has reached, and the error is reported from a later step's save than the one that broke.

Root cause

CheckpointManager.begin_finalization() runs wait_fn and _rename_checkpoint() on a daemon thread. On failure the thread only stores the exception in self._finalize_error (checkpoint.py#L439). The only code that raises it is finalize_pending() (#L446-L458), and the training loops reach that only:

• at the next save, through the finalize_pending() call at the top of begin_finalization() (#L422) or the explicit checkpointer.finalize_pending() the GRPO loops make before init_tmp_checkpoint() (grpo.py#L4118, #L6010), or
• at exit, through shutdown().

Nothing looks at _finalize_error between saves.

Change

• CheckpointManager.raise_if_finalization_failed(): if the background thread has already recorded an error, it calls finalize_pending(), which joins the finished thread, clears the state and raises the same RuntimeError("Background checkpoint finalization failed") with the original exception as its cause. It never waits for a finalization that is still running.
• Every training loop that uses begin_finalization() calls it at the top of each step: GRPO (grpo_train and async_grpo_train), GRPO over TransferQueue (grpo_sync.py), SFT, SFT v2, DPO, RM, PPO (ppo_train and async_ppo_train), distillation, cross-tokenizer off-policy distillation, and the single controller train pump. In the async loops it sits right after the existing per-step check_health call; in the single controller it runs before the maybe_gpu_profile_step calls.
• Single controller: the rollout checkpoint pump can be inside finalize_pending() on a worker thread (via asyncio.to_thread) while it holds _checkpoint_save_lock. The train pump skips its poll while that lock is held, so two threads never flush at once; a skipped poll is caught at the next step. The call stays synchronous (no to_thread), so the locked() check and the poll are atomic on the event loop.
• Docstrings and the comments next to the three begin_finalization() call sites now say that a failure that has already happened is raised at the start of the next step.

Why a non-blocking check

The issue lists a few options. I went with a check that does not block because of what async save is for:

• Joining the finalize thread at every step boundary would make each step wait for the previous checkpoint's disk write, which removes the overlap between the write and training that async_save was added for. With save_period > 1 the write currently has several steps to finish.
• Raising from the finalize thread or a done-callback cannot reach the training thread safely; it would have to interrupt it at an arbitrary point, possibly inside a collective or a Ray call.
• Reading _finalize_error once per step costs nothing when saves succeed, keeps the async write overlapping with training, and reuses the existing error path, so the exception type and message do not change. finalize_pending() before the next save and shutdown() at exit stay the hard barriers they are today.

The trade-off: a write that is still running when the next step starts is not waited for. If it fails during that step, the failure surfaces at the following step boundary (or at the next save or exit, whichever comes first). So the delay is bounded by the step in flight when the write fails, instead of by the save period.

Tests

New tests:

• tests/unit/utils/test_checkpoint.py::TestRaiseIfFinalizationFailed: raises a failure that already happened (original exception kept as __cause__, tmp_step_1 left in place, raised only once); does not block while a finalization is in flight; no-op when nothing failed.
• test_grpo.py::test_grpo_failed_checkpoint_finalization_stops_before_next_step[grpo_train] and [async_grpo_train]: a real CheckpointManager, save_period=2, max_num_steps=6, and policy.finalize_async_save raising OSError. Asserts the run raises after exactly 2 training steps and that tmp_step_2 was never renamed.
• test_checkpointing.py::TestAsyncSaveFinalization::test_failed_background_finalization_stops_before_next_step: the same scenario through the single controller train pump (save_period=2, max_num_steps=4), asserting it stops at _train_steps == 2.
• test_single_controller_actor.py::test_train_pump_polls_finalization_only_while_no_save_holds_the_lock: no poll while the lock is held, one poll once it is free.
• test_failed_checkpoint_finalization_stops_before_training in test_sft.py, test_dpo.py, test_rm.py and test_distillation.py: with the mocked checkpointer's raise_if_finalization_failed raising, the loop raises and train is never called.

The GRPO and single controller loop tests join the finalize thread before the next step starts, so the failure has already happened when training reaches the step boundary, as in the issue's trace. On main, all 11 new tests fail: both GRPO loops and the train pump run steps 3 and 4 and only fail at the step-4 save (assert 4 == 2), the four wiring tests fail with DID NOT RAISE, the lock test sees no call, and the three CheckpointManager tests hit AttributeError.

Edits to existing tests:

• test_failed_background_finalization_raises_at_next_save is now ..._raises_by_next_save, with its comment saying the failure surfaces no later than the step-4 save; the assertion is unchanged.
• _train_pump_controller in test_single_controller_actor.py builds the actor with object.__new__ and sets each attribute the pump reads. It now also sets _checkpointer = MagicMock(), since the pump reads it every step.

Commands and results (CPU only, Python 3.13.14, torch 2.13.0+cpu, `HF_HOME` set; not the CI container)

# the 11 new tests: all pass here; all fail with nemo_rl/ reverted to main
python -m pytest tests/unit/utils/test_checkpoint.py::TestRaiseIfFinalizationFailed \
  tests/unit/algorithms/test_grpo.py::test_grpo_failed_checkpoint_finalization_stops_before_next_step \
  tests/unit/single_controller/test_checkpointing.py::TestAsyncSaveFinalization::test_failed_background_finalization_stops_before_next_step \
  tests/unit/single_controller/test_single_controller_actor.py::test_train_pump_polls_finalization_only_while_no_save_holds_the_lock \
  tests/unit/algorithms/{test_sft,test_dpo,test_rm,test_distillation}.py::test_failed_checkpoint_finalization_stops_before_training

# related suites
python -m pytest tests/unit/utils/test_checkpoint.py tests/unit/single_controller tests/unit/algorithms/x_token \
  tests/unit/telemetry/test_sft_v2.py tests/unit/data_plane/test_sync_one_hop.py \
  tests/unit/data_plane/test_architecture_invariants.py \
  tests/unit/algorithms/{test_grpo,test_sft,test_sft_v2,test_dpo,test_mpo,test_rm,test_ppo,test_distillation}.py \
  tests/unit/algorithms/{test_grpo_checkpoint_engine,test_grpo_group_identity,test_grpo_router_replay_async}.py
# 1956 passed, 9 skipped

ruff check / ruff check --select I / ruff format --check <changed files>   # ruff 0.9.9: clean, 19 files
pyrefly check                                                              # pyrefly 0.24.2: 14 errors, at the same locations as on main, none in changed files

Compatibility

• Additive API: one new public method on CheckpointManager. No config changes and no change to the checkpoint layout.
• No behavior change when finalization succeeds. When it fails, the same exception is raised, only earlier, and the failed tmp_step_N stays in place as before (resume still ignores it).
• Code that drives a training loop with its own fake checkpointer (not a MagicMock or CheckpointManager) needs the new method; there are none in the repo.
• Merge order with refactor(utils): migrate checkpointing configs to BaseModel #4152: that PR changes CheckpointManager.__init__ to read an attribute config (callers pass CheckpointingConfig.model_validate(...)). The new GRPO test builds a CheckpointManager from a plain dict (test_grpo.py around L4270), so whichever PR lands second needs a one-line adapt there. The other new tests go through existing fixtures and config builders.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Unit tests only (listed above); I don't have GPUs for the functional tests.
• Did you add or update any necessary documentation? Not needed: no user-facing config or docs describe when finalization errors surface.

Additional Information

• Suggested CI label: CI:L0 (bug fix).
• Out of scope, noticed while reading: the GRPO loops and the single controller write latest_checkpoint_status.json (last_successful_ckpt_save_completion) right after begin_finalization(), before the rename has happened, so a save whose finalization later fails is still recorded there as successful. Happy to send a follow-up if that file is meant to track finalized checkpoints only.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.