feat(sc): add checkpoint save/restore to SingleController - #3429
Conversation
…ting Port the replay-buffer checkpoint state capture from PR NVIDIA-NeMo#3138 onto the split single-controller path. state_dict snapshots ready slots on the event loop, then fetches each group's DataPlane rows; unready reservations (in-flight rollouts) are dropped. load_state_dict validates the envelope (partition, group size, sample_id uniqueness) before any DataPlane write, truncates to the current capacity keeping the freshest groups, and re-puts rows while rebuilding the parallel slot lists. Staleness filtering is intentionally left to the sampler's first evict. Covered by 9 new unit tests (round-trip, preflight rejection, capacity truncation); 20/20 pass in tests/unit/single_controller/ test_tq_replay_buffer.py. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Two additions the SC checkpointing path needs from the sampler layer: - resume_from_step: BaseSampler (and the three built-in policies + create_sampler) now accept the trainer step the run starts from — 0 for a fresh run, the restored current_step on resume. It seeds the dispatch cursor to preserve the fresh-start invariant _dispatch_index == trainer_version - 1. Without it a restored InOrderSampler stamps target_steps from 0 and every dispatched batch is instantly evicted (target < trainer_version), livelocking the train pump. create_sampler forwards the kwarg to custom samplers only on resume, so fresh starts don't constrain their constructors and an unsupported class fails loudly instead of silently running with an unseeded cursor. - supports_buffer_checkpoint: new PromptGroupSampler property gating replay-buffer save/restore. Only the ungated WindowedSampler returns True — gated policies dispatch a fixed quota per trainer step, so restored groups could never complete an already-consumed window. Covered by 8 new unit tests in test_sampler_interface.py (cursor seeding, gate behavior after resume, factory forwarding, custom fail-loud, checkpoint-support matrix). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Replace the checkpointing NotImplementedError guard with the actual driver-side resume wiring, following the grpo.py setup pattern: - Build a CheckpointManager unconditionally and resolve the latest checkpoint: load_training_info() populates save_state (default GRPOSaveState when starting fresh) and get_resume_paths() yields the weights/optimizer paths. - _build_trainer takes kw-only weights_path/optimizer_path and forwards them to TQPolicy (previously hardcoded to None) on both the colocated and non-colocated build paths. - Restore the dataloader position from train_dataloader.pt when present (load_dataloader_state, with its dataset-swap guard); warn and start fresh otherwise. Runs before _clamp_max_num_steps as before. - Forward checkpointing.pretrained_checkpoint into the policy config. - SingleControllerActorArgs carries two new fields, save_state and last_checkpoint_path, for the actor-side restore (next step). Saving itself is not wired yet — that lands in the SingleControllerActor train pump next. Existing tests updated for the new surface: the setup tests' hand-built checkpointing block now carries the keys CheckpointManager indexes, and the pump tests pass the two new ActorArgs fields. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Wire the actor side of SC checkpointing (the setup/resume half landed in the previous commit), with Megatron async_save supported end to end: - Restore: __init__ rebuilds counters (train_steps/trainer_version/ current_epoch/consumed_samples/total_valid_tokens) from the save_state loaded by setup, seeds the sampler with resume_from_step, and run() reloads replay-buffer groups (ungated samplers only, one capacity permit per restored group) before the pumps start. - Save: after each weight sync, _save_checkpoint mirrors async_grpo_train's block — finalize_pending flushes the previous background finalization, save_checkpoint returns after D2H staging under async_save, aux state (training info, dataloader position, replay buffer when the sampler supports it) is written synchronously, then begin_finalization defers the tmp->step rename until the async weight writes complete. run() flushes the last checkpoint via checkpointer.shutdown() on every exit path. - TimeoutChecker drives checkpoint_must_save_by: a timeout save also stops training early, matching the legacy loops. - latest_checkpoint_status.json is refreshed after each save for external watchdogs (reuses grpo's _write_latest_checkpoint_status). The pump tests' hand-built configs gain the checkpointing block the actor now reads (enabled=false keeps them write-free). Validated end to end on GB200 (Qwen3-0.6B, megatron async_save=true): 4-step run saves step_2/step_4 with no tmp_step_* leftovers; a checkpoint_must_save_by run stops early with a complete checkpoint; the resume run restores dataloader + 4 replay groups and continues from step 2 to step 4. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Port the checkpointing test suite from PR NVIDIA-NeMo#3138 onto the split SC architecture (actor_args, the PromptGroupSampler protocol, the split trainer step API) and extend it for the async-save path: - counter/sampler-cursor restore, save triggers (period boundary, last step, checkpoint_must_save_by timeout, disabled, save_optimizer), metric_name handling, dataloader state round-trip with the dataset-swap guard, and setup resume wiring (get_resume_paths forwarded to the trainer factory, training_info.json loaded). - replay-buffer persistence is asserted against sampler.supports_buffer_checkpoint (windowed saves/restores with one capacity permit per group; gated samplers skip both sides). - new async-save coverage: the tmp->step rename stays deferred until finalize_async_save completes and is flushed by shutdown; a failed background finalization re-raises at the next save; _save_checkpoint records val_metrics into val_reward and a val:* metric_name. 32 tests, in-process actor with fakes (ray.cluster_resources patched); 108 passed together with the existing single_controller suite. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
On exit, run() now propagates a failed checkpoint finalization only on the clean path; when an exception is already propagating the flush is best-effort (warning), so the original training failure stays the raised exception — matching async_grpo_train's guarded cleanup shutdown. logger.finish() moves into its own finally so it runs either way. Also drop the stale "SC does not support checkpointing yet." comment from the SC exemplar config. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
yuki-97
left a comment
There was a problem hiding this comment.
@haitian-nvidia thanks for supporting this! left some comments.
Three changes from the PR review: - A resume checkpoint missing train_dataloader.pt now raises instead of warning and starting from a fresh dataloader position. SC always writes the file on save, so absence means a corrupted checkpoint; this matches GRPO's load contract. - The sampler dispatch cursor is now seeded via a post-construction BaseSampler.set_dispatch_index(resume_from_step) hook (also on the PromptGroupSampler protocol) instead of threading a resume_from_step kwarg through five constructors and the factory — and it removes the fresh-vs-resume asymmetry for custom FQN samplers. - run()'s exit path flattens to logger.finish() + a bare checkpointer.shutdown(): Python's exception chaining already keeps an original training failure visible if the flush also fails, so the explicit guard read against the fail-loud policy. Tests updated accordingly (setter-based seeding cases, the missing-file test now expects FileNotFoundError, resume-wiring fixtures write dataloader state); 107 passed. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
- Round-trip now covers start_weight != end_weight and a non-None target_step (the InOrderSampler selection key), an empty buffer, and an unready slot sandwiched between ready ones (guards the by-index skip in state_dict). - One round-trip drives a real TensorDict payload (mixed dtypes and a non-contiguous view) through torch.save/torch.load, exercising the serialization path the production replay_buffer.pt actually uses instead of the opaque fake payload. - New composed resume test: an actor restored at current_step=2 running a live pump to max_num_steps=4 ends at exactly 4 steps and only checkpoints the post-resume boundary. 112 passed with the existing single_controller suite. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Replace the sampler-level supports_buffer_checkpoint capability with a recorded-identity check, per review: - Every save now writes replay_buffer.pt regardless of sampler, and records async_rl.sampler.name into training_info.json. - On resume, the buffer is restored only when the saved sampler name matches the current one; a mismatch warns and skips (checked before the torch.load, so a mismatched buffer file is never read). This also covers resuming an old checkpoint that predates the recorded name. - supports_buffer_checkpoint is dropped from PromptGroupSampler, BaseSampler and WindowedSampler; TQReplayBuffer is untouched (the sampler identity is run metadata, so it lives in training_info, not the buffer envelope). Also remove the val_metrics placeholder from _save_checkpoint (val lands with the future validation loop); the stale val_reward sentinel is still dropped from training_info. 108 passed with the existing single_controller suite. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Two runs of examples/run_grpo_single_controller.py sharing one checkpoint_dir and one config (same max_num_steps, so Megatron train_iters stays consistent across the restart): - Run 1 sets checkpointing.checkpoint_must_save_by=00:00:00:01 so the timeout save fires deterministically at the first step boundary and training stops early, leaving a complete step_1 checkpoint (weights, dataloader position, replay buffer; no tmp_step_* leftovers, i.e. the async finalization was flushed) with current_step=1 and the saving sampler_name recorded in training_info.json. - Run 2 drops the timeout and must restore the dataloader position and replay buffer, then continue from step 2 to step 4. Registered in L1_Functional_Tests_SingleController.sh (full tier; the fast tier keeps the two existing quick scripts). Verified end to end on 2 GPUs (Qwen3-0.6B, megatron async_save=true), ~8 minutes of runtime. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
# Conflicts: # nemo_rl/algorithms/single_controller.py # nemo_rl/algorithms/single_controller_utils/setup.py # tests/unit/single_controller/test_rollout_pump.py # tests/unit/single_controller/test_train_pump.py
Addresses two review threads on NVIDIA-NeMo#3429. GRPOSaveState now declares `sampler_name`, so the SC replay-buffer restore gate no longer rides on a runtime-only attribute. Both `# type: ignore[attr-defined]` sites and the setup-side re-attach are gone; `_get_grpo_save_state` carries the field through on its own, closing the two silent-drop paths that would have disabled the gate. Every algorithm's training_info.json now carries the field, so test_get_grpo_save_state_handles_legacy_checkpoint_and_filters_metrics expects it. checkpointing.metric_name is now validated in validate_single_controller_config instead of warning at every save. SC has no validation loop, so a "val:" metric is never collected and keep_top_k silently degrades to a no-op; the check rejects any non-"train:" prefix when checkpointing is enabled. Every SC config inherited "val:accuracy" from grpo_math_1B.yaml, so the SC exemplar and the one SC recipe that leaves checkpointing enabled now set metric_name: null. Also repairs test_logs_hyperparameters_and_concrete_weight_synchronizer, which has been broken since the actor started building a CheckpointManager in __init__, and applies pending ruff format fixes. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
…ffer Addresses the two replay-buffer review threads on NVIDIA-NeMo#3429. load_state_dict truncated an over-capacity checkpoint freshest-first, which is only safe when the sampler leaves target_step unset. For the InOrder family start_weight and target_step are both monotonic, so freshest-first keeps the far-future targets and drops the near-term ones; InOrderSampler.evict only drops target < current_train_weight, so those permits are never released and the rollout pump blocks on _buffer_capacity while the train pump waits for a target that can no longer be filled -- a silent deadlock, triggered whenever max_buffered_rollouts shrinks across save/resume. Keeping the smallest targets is no better: rollout would just re-produce them from the restored dataloader position. There is no meaningful subset to keep, so an over-capacity target-stamped checkpoint now raises, naming both ways out (raise max_buffered_rollouts, or delete replay_buffer.pt). Samplers that leave target_step unset keep the existing freshest-first truncation. In-order selection also assumes exactly num_prompts_per_step groups per target step, but a resumed run restored groups already stamped for the upcoming step and then dispatched a full batch on top of them. The surplus could never be selected and held its capacity permits until evict. The rollout pump now sizes each batch to the shortfall reported by TQReplayBuffer.count_for_target_step and drops the tail, matching the legacy async path; generation-level backpressure would let us keep it instead. Counting once after admit is both exact and necessary: each admitted batch gets a unique target_step, and reserve() runs inside the dispatched task, so a per-prompt check would not yet see the prompts it just dispatched. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Addresses the three test-coverage threads on NVIDIA-NeMo#3429. The async-save failure path was only exercised through a *subsequent* save's finalize_pending. The last save has no successor, so its background failure could only surface through run()'s exit path, and nothing asserted that shutdown() re-raises it. The new companion drives a single failing save and asserts both that the rename never happened (the checkpoint is still tmp_step_2) and that shutdown() raises. Every replay-buffer restore test ran with max_num_steps=0, so _train_pump's body never executed and the restore's permit accounting was only ever checked at rest. A regression that leaked the restored permits would starve the rollout pump without failing the suite. The new _run_restore_then_train_pump helper composes the two halves -- restore, then a live pump -- and asserts the semaphore returns to full: the restore drains it to zero and the pump releases one permit per selected group. test_run_restore_at_full_capacity_does_not_hang is deleted rather than kept: K acquisitions against a semaphore with K permits never block, so it asserted the same permit accounting as the K=3 test above it. The new live-pump test uses K == max_buffered_rollouts, so the full-capacity shape stays covered and now also proves the permits come back. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
terrykong
left a comment
There was a problem hiding this comment.
Agent-team review — 6 agents
Thanks for this — closing the NotImplementedError on the SC path is a real unblock, and the test work is genuinely strong. test_tq_replay_buffer.py's additions in particular (a real torch.save/torch.load cycle over a non-contiguous TensorDict view, preflight tests that also assert nothing reached the DataPlane) are the best part of the PR. Reusing CheckpointManager.begin_finalization rather than re-implementing a second flush protocol inside the actor was the right call.
Please rebase — and by hand. The PR is CONFLICTING with main. The conflicts come from a single commit, 0e687e6d0 (#3499, NeMo-Gym spinup), which touches none of the checkpointing logic — but it lands on the exact _build_trainer(weights_path=…, optimizer_path=…) call sites this PR adds, so please resolve manually rather than with -X ours/-X theirs.
Two findings worth acting on before merge, both about save triggers rather than the save path itself: the missing dataloader-position disjunct in is_last_step (a resumed run can end short and write no final checkpoint), and ft_save_period being silently ignored. The other four are a doc fix, a test-scope note, a docstring, and a one-line to_thread nit.
On the open threads: we verified the other seven unresolved threads against HEAD and they're genuinely addressed — the 8b8bfa847 truncation fix and the count_for_target_step top-up were each re-derived independently by two reviewers and are correct (the top-up is exactly permit-balanced, which is the non-obvious part). Those can be resolved. The one comment below on the functional test follows up on the diamond thread, where your position turns out to be right and the suggested determinism knob turns out to be a no-op.
Verification caveat, stated plainly: this review is static. The repo lockfile is Linux-only and the review host is macOS, so neither pre-commit nor any test was executed — there is no lint or test signal behind these findings, only source reading. Anything needing a runtime is flagged as unverified inline.
Generated by Claude Code
Signed-off-by: haitian-nvidia <haitianj@nvidia.com> # Conflicts: # nemo_rl/algorithms/async_utils/staleness_sampler.py # nemo_rl/algorithms/single_controller.py # nemo_rl/algorithms/single_controller_utils/setup.py # tests/unit/single_controller/test_rollout_pump.py # tests/unit/single_controller/test_sampler_interface.py
Addresses the four actionable threads from the second review round on NVIDIA-NeMo#3429. is_last_step now also fires when rollout is exhausted and the replay buffer is drained. A resumed run budgets max_num_steps against the full per-epoch batch count (_clamp_max_num_steps runs after the dataloader position is restored), so it can run out of data first; the pump then exited with status 0 right after a completed step whose checkpoint was never written. Anchoring on rollout exhaustion is the SC equivalent of grpo_train's dataloader-position disjunct. checkpointing.ft_save_period now ORs into the save trigger, matching grpo_sync and every other algorithm; it was accepted and silently ignored on the SC path. Also: _write_latest_checkpoint_status moves off the event loop like the other five blocking calls in _save_checkpoint; the three SC recipes narrow their stale "no validation and checkpointing" comment to the still-true validation half; and the functional test header now states it verifies save/restore mechanics only (prompts in flight at save are dropped while the dataloader cursor has advanced past them, so metric equivalence across resume is not defined). Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
|
/ok to test 6a3408d |
pyrefly's cycle-breaking narrows the asyncio.to_thread result at the init_tmp_checkpoint call to LiteralString and rejects the assignment; an explicit PathLike annotation pins the intended type. Also apply the two pre-commit hook diffs CI reported: drop a double blank line ahead of is_last_step and move the aliased grpo MasterConfig import below the parenthesized import group. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
Head branch was pushed to by a user without write access
|
/ok to test ab447cb |
|
/ok to test ab447cb |
The explicit PathLike annotation didn't help: pyrefly resolves the alias itself inconsistently while breaking the import cycle, so the annotated assignment still reports bad-assignment. Use the repo's inline ignore convention instead. Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
|
/ok to test 5a568a8 |
What does this PR do ?
Adds checkpoint save/resume to the SingleController path, closing the
NotImplementedErrorplaceholder left insetup_single_controller.Save (in the train pump, after each weight sync):
policy.save_checkpoint, and optimizer state whencheckpointing.save_optimizeris set. Megatronasync_saveissupported: the tmp→step rename is deferred with
begin_finalization(wait_fn=policy.finalize_async_save)and flushed atthe next save or on exit, so training/rollouts continue while the weight
write completes in the background.
training_info.json), dataloader position(
train_dataloader.pt), and — when the sampler supports it(windowed/over-sampled) — the TQ replay buffer's committed prompt groups
(
replay_buffer.pt, meta + DataPlane payloads).save_periodboundary, last step, andcheckpoint_must_save_by(timeout save + early stop).latest_checkpoint_status.jsonis refreshed after each save.Resume:
setup_single_controllerresolves the latest checkpoint, loadstraining_info.json, feedsweights_path/optimizer_pathto thetrainer, and restores the dataloader position (with the existing
dataset-swap guard).
SingleControllerActorrestores its counters, seeds the samplerdispatch cursor via the new
resume_from_stepargument, and reloadsreplay-buffer groups (re-acquiring one buffer-capacity permit per
group) before the pumps start.
create_sampler(..., resume_from_step=...)and asupports_buffer_checkpointproperty onPromptGroupSampler(True forthe windowed sampler; gated samplers skip buffer save/restore).
Usage
uv run examples/run_grpo_single_controller.py \ checkpointing.enabled=true \ checkpointing.checkpoint_dir=results/my-run \ checkpointing.save_period=10 # Re-running the same command resumes from the latest step_N checkpoint.Before your PR is "Ready for review"
Pre checks:
Additional Information
tests/unit/single_controller/— 108 passed, including 32new tests in
test_sc_checkpointing.py(save triggers, async-savefinalization ordering and failure propagation, metric_name handling,
dataloader round-trip, replay-buffer persistence, setup resume wiring).
async_save: true):a 4-step run saves complete
step_2/step_4checkpoints with notmp_step_*leftovers; acheckpoint_must_save_byrun saves and stopsearly; resuming restores the dataloader + replay buffer and continues
to step 4.
checkpointingblock).