SC checkpointing: persist partial rollout state so in-flight generations resume instead of restarting
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
> 📊 **Visual explainer:** [how the dataloader cursor advances past in-flight prompts at save](https://terrykong.github.io/gh-pages-poc/terryk/pr3429-cursor-advance.html) — the three-pointers timeline and the saved/in-flight/never-dispatched split.
## Summary
SC checkpointing (#3429) saves the dataloader cursor at the **rollout pump's** position while persisting only **ready** replay-buffer groups. Prompts that were in flight (or yielded-but-undispatched) at save time are behind the restored cursor yet absent from the buffer.
**Scope of this issue: partial-rollout persistence.** Even after the prompts themselves are re-enqueued on resume (requested in the #3429 review — see below), the *partial generations* are still discarded and re-run from scratch. Full correctness — resuming in-flight rollouts rather than restarting them — requires checkpointing partial rollout state through TQ/DataPlane, which is the planned follow-up this issue tracks.
## Mechanism (at PR #3429 head `c039949`)
1. The cursor advances at **batch** granularity the moment `for prompt_batch in self._dataloader:` yields ([single_controller.py#L375](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/single_controller.py#L375)); dispatch inside the batch proceeds prompt-by-prompt, gated on buffer/in-flight permits ([#L398-L400](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/single_controller.py#L398-L400)).
2. `_save_checkpoint` snapshots that already-advanced cursor ([#L698](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/single_controller.py#L698)).
3. `TQReplayBuffer.state_dict` persists ready groups only; unready reservations are dropped ([replay_buffer.py#L861](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/async_utils/replay_buffer.py#L861)) — and their prompts are unrecoverable from the buffer, since `reserve()` stores `meta=None` until commit ([#L685](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/async_utils/replay_buffer.py#L685)).
4. The DataPlane is a fresh instance in the resumed job, so the checkpoint is the only carrier — nothing resumes the lost rollouts.
## Why this matters beyond throughput
The drop is **biased**: long-generation prompts (agentic, long-context) are disproportionately likely to be in flight at any save boundary, so they are systematically the ones lost. `in_order`'s top-up ([single_controller.py#L380-L390](https://github.com/NVIDIA-NeMo/RL/blob/c039949edb91e972cb0b40393a11f285458d5373/nemo_rl/algorithms/single_controller.py#L380-L390)) preserves the per-step group *count* but fills from fresh prompts; `windowed` preserves nothing.
## Status / split of work
- **In #3429 (requested in review):** persist the pending prompts' `DatumSpec`s (known driver-side at save time) and re-dispatch them through the normal `admit`/permit path on resume. This restores exact-once *prompt* coverage; the in-flight generation work is still redone from scratch.
- **This issue (follow-up):** checkpoint partial rollout state so in-flight generations *resume* instead of restarting — removing both the redundant recompute and the residual bias against long generations when a resumed-from-scratch rollout competes against fresh short ones.
Related: #3429 (SC checkpoint save/restore), where this was surfaced during review.
Contributor guide
Assessment
This issue has not been assessed yet.