feat(sc): add checkpoint save/restore to SingleController

[haitian-nvidia]

What does this PR do ?

Adds checkpoint save/resume to the SingleController path, closing the
NotImplementedError placeholder left in setup_single_controller.

Save (in the train pump, after each weight sync):

• Policy weights via policy.save_checkpoint, and optimizer state when
  checkpointing.save_optimizer is set. Megatron async_save is
  supported: the tmp→step rename is deferred with
  begin_finalization(wait_fn=policy.finalize_async_save) and flushed at
  the next save or on exit, so training/rollouts continue while the weight
  write completes in the background.
• Training state (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).
• Triggers: save_period boundary, last step, and
  checkpoint_must_save_by (timeout save + early stop).
  latest_checkpoint_status.json is refreshed after each save.

Resume:

• setup_single_controller resolves the latest checkpoint, loads
  training_info.json, feeds weights_path/optimizer_path to the
  trainer, and restores the dataloader position (with the existing
  dataset-swap guard).
• SingleControllerActor restores its counters, seeds the sampler
  dispatch cursor via the new resume_from_step argument, and reloads
  replay-buffer groups (re-acquiring one buffer-capacity permit per
  group) before the pumps start.
• New sampler surface: create_sampler(..., resume_from_step=...) and a
  supports_buffer_checkpoint property on PromptGroupSampler (True for
  the 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:

• 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? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• Unit tests: tests/unit/single_controller/ — 108 passed, including 32
  new tests in test_sc_checkpointing.py (save triggers, async-save
  finalization ordering and failure propagation, metric_name handling,
  dataloader round-trip, replay-buffer persistence, setup resume wiring).
• E2E on GB200 (Qwen3-0.6B, megatron backend, async_save: true):
  a 4-step run saves complete step_2/step_4 checkpoints with no
  tmp_step_* leftovers; a checkpoint_must_save_by run saves and stops
  early; resuming restores the dataloader + replay buffer and continues
  to step 4.
• No documentation changes needed (config surface reuses the existing
  checkpointing block).

[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.

[yuki-97]

@haitian-nvidia thanks for supporting this! left some comments.

[yuki-97]

Round-trip only exercises start_weight == end_weight and target_step is None. _add_group(weight=w) defaults end_weight=weight, and commit never sets target_step — so a load bug that swapped start_weight / end_weight at replay_buffer.py:1002-1003 or silently dropped target_step (the InOrderSampler's primary key at staleness_sampler.py:431) would still pass.

Suggest fix: parameterize with start=1, end=2 on one group and target_step=7 on another.

[haitian-nvidia]

Done in cbbbe2b — test_round_trip_preserves_end_weight_and_target_step parameterizes exactly as suggested: one group with start=1, end=2 and one with target_step=7, asserting all three lists survive the round-trip. Built via the real reserve(target_step=...) → commit path (rather than a hand-built envelope) so the save side of target_step serialization is exercised too.

[yuki-97]

FakeDataPlaneClient.get_samples returns a hand-crafted {"payload_for": sample_ids} dict and load_state_dict re-puts it verbatim — but the production replay-buffer file goes through torch.save(buffer_state, replay_buffer.pt) on TensorDict-valued fields_data (see single_controller.py:290,766). Any bug in that real serialization path (dtype/device promotion, non-contiguous strides, custom-Tensor pickle) is uncaught here.

Suggest fix: one round-trip that stores an actual TensorDict payload through an in-memory BytesIO torch.save/torch.load cycle.

[haitian-nvidia]

Done in cbbbe2b — test_round_trip_tensordict_payload_through_torch_save stores a real TensorDict as fields_data (mixed int64/float32 dtypes + a non-contiguous [:, ::2] view) through an in-memory BytesIO torch.save/torch.load(weights_only=False) cycle (matching the production load), then asserts per-key torch.equal after load_state_dict re-puts.

[yuki-97]

Three round-trip shapes aren't covered: (a) empty buffer (common resume shape when no group was committed before the checkpoint); (b) middle-unready [ready, unready, ready] (state_dict skips unready by index; the trailing-unready test wouldn't catch an off-by-one in that skip); (c) target_step != None — the InOrderSampler resume path, which needs _make_group_entry(target_step=7) to construct.

[haitian-nvidia]

Done in cbbbe2b — (a) test_round_trip_empty_buffer; (b) test_state_dict_skips_middle_unready ([ready, unready, ready], asserting the neighbouring groups' fields don't shift); (c) covered by test_round_trip_preserves_end_weight_and_target_step (see the thread above) via the real reserve(target_step=7) path.

[yuki-97]

All four _run_actor_run call sites use max_num_steps=0, so _train_pump never runs post-restore. TestCounterRestore covers constructor seeding and TestSaveTrigger covers the fresh pump, but nothing composes the two to prove that resuming with save_state["current_step"]=N and running the pump yields _train_steps = N + k (not k, not 2N + k).

Suggest fix: one test that primes save_state["current_step"]=2, runs with max_num_steps=4, and asserts final _train_steps == 4.

[haitian-nvidia]

Done in cbbbe2b — test_resumed_pump_continues_to_max_steps primes save_state["current_step"]=2, runs the live pump with max_num_steps=4, and asserts _train_steps == 4 (not 2, not 6), plus that only the post-resume step_4 boundary is checkpointed and consumed_samples accumulates on top of the restored value.

[yuki-97]

SC guards load_dataloader_state with if os.path.exists(dataloader_state_path) and prints a warning on miss, while GRPO's init_train_dataloader (grpo.py:451-452) calls it unconditionally and lets a missing file raise. SC always writes train_dataloader.pt on save (line 759), so absence means externally-tampered / partial ckpt — silent-fresh-start then replays already-seen samples with no signal to the user.

Suggest fix: drop the exists guard and let load_dataloader_state raise, matching GRPO's contract.

[haitian-nvidia]

Done in 26c0f92 — dropped the os.path.exists guard; a resume checkpoint missing train_dataloader.pt now raises (matching GRPO's contract). The unit test now expects FileNotFoundError.

[yuki-97]

resume_from_step is threaded through 6 signatures + a special-case != 0 gate for CustomSamplerConfig at line 561-562 that silently overrides YAML model_extra on resume but silently loses to it on fresh runs.

Since only _dispatch_index needs to be restored, wdyt replacing the kwarg with a BaseSampler.set_dispatch_index(resume_from_step) post-construction hook, called once from SingleControllerActor.__init__ after create_sampler? Removes the 5 sub-constructor signature changes, the factory kwarg, and the CustomSampler asymmetry.

[haitian-nvidia]

Done in 26c0f92 — replaced the resume_from_step kwarg threading with a BaseSampler.set_dispatch_index(resume_from_step) post-construction hook (also on the PromptGroupSampler protocol), called once from SingleControllerActor.__init__. The factory kwarg and the custom-sampler fresh-vs-resume asymmetry are gone.

[yuki-97]

sys.exc_info()[0] is not None + try/except buys "shutdown errors don't mask training errors" — but Python's built-in exception chaining already preserves the original via __context__, and the shape reads against skills/error-handling's fail-loud policy. The whole block flattens to:

Suggested change

	# Flush the last checkpoint's background finalization; a failure
	# raises on a clean exit but never masks a propagating exception.
	propagating = sys.exc_info()[0] is not None
	try:
	await asyncio.to_thread(self._checkpointer.shutdown)
	except Exception:
	if not propagating:
	raise
	warnings.warn(
	"Checkpoint finalization failed while handling an "
	"exception; the original exception will be re-raised.",
	stacklevel=2,
	)
	finally:
	self._logger.finish()
	self._logger.finish()
	await asyncio.to_thread(self._checkpointer.shutdown)

[haitian-nvidia]

Done in 26c0f92 — flattened to self._logger.finish() + a bare await asyncio.to_thread(self._checkpointer.shutdown), relying on exception chaining.

[yuki-97]

for the val part, is it make sense to you to remove it and add back when supporting val?
I'm also ok with keeping it here since val should be supported soon, so just a notice to add some related tests when supporting val.

[haitian-nvidia]

Done in 60de5aa — removed the val placeholder from _save_checkpoint (the val_metrics parameter and its branch); we'll add it back with the validation loop, along with tests. The stale val_reward sentinel from _default_grpo_save_state is still dropped so it never lands in training_info.json.

[yuki-97]

instead of supports_buffer_checkpoint, I think maybe better to

1. save sampler_name in ckpt.
2. if sampler_name equal with current, load buffer
3. if sampler_name not equal, warning and return.

something like below:

Suggested change

	if not (
	self._sampler.supports_buffer_checkpoint
	and self._last_checkpoint_path is not None
	):
	return
	if not self._last_checkpoint_path is not None:
	return
	if sampler_name != current_sampler_name:
	print(...)
	return

[haitian-nvidia]

Done in 60de5aa, with one refinement on where the name lives: supports_buffer_checkpoint is gone — every save writes replay_buffer.pt and records async_rl.sampler.name into training_info.json (run metadata, alongside current_step), rather than into the buffer envelope. On resume the saved name is compared against the current sampler before the torch.load, so a mismatched (or pre-this-change) checkpoint warns and skips without reading the buffer file. TQReplayBuffer stays sampler-agnostic.

[yuki-97]

can we add functional tests like tests/functional/grpo_async_replay_buffer_checkpoint.sh and tests/functional/sft_resume_diamond.sh?

[haitian-nvidia]

Done in 7c1d6da — added tests/functional/grpo_checkpoint_single_controller.sh (registered in L1_Functional_Tests_SingleController.sh, full tier). Same two-run shape as grpo_async_replay_buffer_checkpoint.sh: run 1 uses a 1-second checkpoint_must_save_by so the timeout save fires deterministically at the first step boundary and stops early (also covering the timeout-save path); run 2 resumes from step_1, asserting the dataloader + replay-buffer restore lines and reaching step_4. Verified end to end on 2 GPUs (~8 min runtime).

[haitian-nvidia]

@mehraakash