How an unfinished prompt now survives a checkpoint

NVIDIA-NeMo/RL PR #3480, new since the last review round — the rollout recovery ledger. Links pinned to bca9a721. Part 1 of 3 · the atomic cut · the proof

The last review round ended with a known gap: a rollout still generating when a checkpoint lands was dropped, and its prompt with it — the dataloader had already moved past it and never rewinds. The author's answer is a new file, rollout_recovery.py, built around one rule: at any instant, every prompt has exactly one durable owner. Either the dataloader has not handed it out yet, or the new ledger owns it (handed out, rollouts not finished), or the replay buffer owns it (finished, waiting to be trained on). A checkpoint saves all three owners in one atomic cut, so a restart can rebuild exactly the in-flight work it interrupted.

One prompt, three owners

dataloader prompt 41 not yet taken saved: train_dataloader.pt recovery ledger (NEW) taken, rollouts not finished RESERVED → ADMITTED saved: rollout_recovery.pt replay buffer all rollouts finished rows live in TransferQueue; saved: replay_buffer_metadata.pt reserve_group TQ commit then discard_group spare pool whole batches set aside to replace failed rollouts divert substitute → saved: replacement_reserve.pt inside the ledger — two phases RESERVED dataloader advanced past it; the sampler has not stamped a step yet ADMITTED step assigned (target_step); rollouts dispatched, still generating On restart: ADMITTED groups redispatch as-is; RESERVED batches go through sampler admission again. Each record keeps its group_id, so the same TQ slot is reused. what one checkpoint contains (all captured in a single atomic cut — see part 2) train_dataloader.pt — where the loader stopped replacement_reserve.pt — the spare pool (only when non-empty) rollout_recovery.pt — unfinished groups: locators only, no prompt text, no tensors replay_buffer_metadata.pt + data_plane/ — finished groups: index + the TQ rows themselves sha256 of rollout_recovery.pt, its schema version, and its group count are written into the TQ checkpoint metadata — a torn or stale pair is rejected at load, never silently redispatched.
Ownership moves left to right. The ledger covers exactly the window that used to lose prompts: after the dataloader moved on, before the rollouts finished. The spare pool is a side store for batches deliberately set aside to replace failed rollouts.

Worked example — all numbers distinct. Config: 4 rollouts per prompt group, gate lets generation run 1 step ahead, trainer at step 7:

prompt groupstate when the checkpoint landsowner savedon restart
g20 (step 7)all 4 rollouts finished, in TQreplay buffertrained on directly — nothing regenerates
g21 (step 8)2 of 4 rollouts still generatingledger, ADMITTED, target_step=8redispatched for step 8; all 4 regenerate
g22 (no step yet)taken from the loader, waiting at the gateledger, RESERVEDgoes through admission again, then dispatches
g23still in the datasetdataloader positionserved normally when the loader reaches it

Before this change, g21 and g22 were simply gone: the loader was saved past them and no file held them. The earlier review round got the finished case (g20) recovered; this round closes the unfinished ones.

What the ledger saves — and how restore trusts it

The in-memory record keeps the live prompt by reference, but state_dict() persists only a locator: the dataset row number (idx) and task_name. No prompt text, no tensors — the file stays tiny regardless of prompt length. Restore rebuilds each prompt by reading dataset[idx] again and re-running the same one-row collate_fn the original dispatch used (_rehydrate_rollout_recovery_prompts). This is safe under shuffling: shuffle changes the order rows are visited, not which row an idx names.

Restore refuses to guess. The chain, each step a hard error on mismatch:

1  TQ metadata advertises the file    → file missing?               error
2  file exists, metadata silent?      → orphan from another run     error
3  sha256(file) vs metadata           → torn/stale pair             error   single_controller.py:624
4  schema_version == 1                → written by other code       error
5  group count vs metadata            → truncated file              error
6  per-group field validation         → corrupt record              error   rollout_recovery.py:347
7  drop groups the replay buffer owns → never two owners            discard_canonical_groups
8  rebuild each prompt, check idx + task_name match                 error on mismatch
The one thing restore cannot detect: a dataset whose row content changed under the same idx. The author removed content fingerprints deliberately (commit fd1e424e5) and pinned the behavior with a test — a reordered or shrunk dataset is caught, an edited row is not. Resuming across a dataset edit silently trains the recovered groups on the new content.
So what. With this in place, stopping a run mid-step and resuming it trains on the same prompts as never stopping: finished work is reused, unfinished work is regenerated from the same dataset rows, and nothing the loader handed out can fall between the files. The regenerated rollouts are new samples (different tokens, same prompt) — the data stream is preserved, not the exact trajectories.