NVIDIA-NeMo / NVIDIA-NeMo/RL

SC checkpointing: persist partial rollout state so in-flight generations resume instead of restarting

Open
#3,594 2 comments 0 reactions 1 assignee Claimed by @macandro96 View on GitHub
bug Documentation Feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.