Can a restore drop a prompt? Only if the sampler cursor comes back wrong

For NVIDIA-NeMo/RL #3925. All permalinks pin a4c10a78. The evidence is a round-trip test that runs the real code through a borrow, a checkpoint and a restore.

Three pages from the #3925 review — this one: whether restoring can lose a prompt.

Every group in the recovery ledger is in one of two phases:

phasewhat happenedsealed siblingson restore
RESERVED claimed from the dataloader, but the sampler admission never committed — so it was never dispatched always 0 re-admitted as a batch — the only phase the discard below can touch
ADMITTED admission committed; dispatch happens after, so generation may have started 0 up to group size − 1 relaunched unconditionally, never dropped

The phase is set when the group is first recorded — ADMITTED if admitted else RESERVED — and moves to ADMITTED when the admission commits. Re-admitting a RESERVED batch runs one branch that can release a prompt:

buffered = self._buffer.count_for_target_step(target_step) if buffered: dispatch_count = max(0, len(group_ids) - buffered) dispatch_group_ids = group_ids[:dispatch_count] for group_id in group_ids[dispatch_count:]: self._rollout_manager.discard_prompt_group(cut, group_id) # never re-served

single_controller.py:1149-1152

It fires when the replay buffer already holds a group stamped for the step this batch is being admitted to. The one mechanism that moves finished groups between steps is the borrow: it re-stamps a finished group into an earlier step and sends a spare to repay the later one, and it runs with rollout recovery on (:1983) as well as off (:2020). Even so, a borrow cannot make the discard fire — the test below runs one through a checkpoint and a restore to show it, and then shows the one thing that can.

The round trip, on the real code

Real: the rollout pump with its admission, spare pool, borrow and repayment; TQReplayBuffer; the recovery ledger; the in_order sampler; the checkpoint barrier; _capture_rollout_checkpoint_cut; and, in a fresh controller, _maybe_restore_replay_buffer, _maybe_restore_rollout_recovery and _redispatch_restored_rollouts. Faked: generation itself, and the tensor converter. Group size 2, 3 groups per step, lookahead 2, on_dropped_prompt: replace. Prompts are numbered in dataloader order; batch 3–5 became the spare pool.

1. Live: prompt 1 is given up on, so step 0 borrows from step 2 BEFORE step 0 012 prompt 1 failed step 1 678 step 2 91011 all finished; due last AFTER step 0 029 full — closes now, on step 2’s group 9 step 1 678 untouched step 2 10113 spare 3 dispatched FOR step 2 — still in flight 2. The checkpoint: 8 finished groups, 1 in flight, 3 never dispatched replay buffer — finished groups, with their step stamp step 0 prompts 0, 2, 9 step 1 prompts 6, 7, 8 step 2 prompts 10, 11 recovery ledger — unfinished groups prompt 3 · ADMITTED · step 2 · in flight prompts 12, 13, 14 · RESERVED · no step yet held at the gate: the cursor is at the lookahead limit Also saved: the sampler cursor, sampler_dispatch_index = 2 — the last step admitted. Nothing in the buffer is stamped past it. 3. Restore: the reserved batch is re-admitted — the cursor decides its step cursor restored from the checkpoint: 2 the batch becomes step 3 already stamped for step 3: none 12, 13, 14 dispatched nothing discarded — 14 of 14 prompts back cursor rebuilt from trainer_version: −1 the batch becomes step 0 already stamped for step 0: three (0, 2, 9) 12, 13, 14 discarded three prompts lost — the dataloader is already past them finished finished, from step 2 spare, in flight not generated yet given up on discarded

What the test prints, both ways (labels shortened; every value is verbatim):

saved cursor before the drop : {0: [0, 1, 2], 1: [6, 7, 8], 2: [9, 10, 11]} after borrow + repay : {0: [0, 2, 9], 1: [6, 7, 8], 2: [3, 10, 11]} ledger in checkpoint : (3, ADMITTED, step 2) (12, RESERVED) (13, RESERVED) (14, RESERVED) restored cursor : 2 after restore : {0: [0, 2, 9], 1: [6, 7, 8], 2: [3, 10, 11], 3: [12, 13, 14]} discarded on restore : [] telemetry : groups_considered 4 groups_redispatched 4 cursor rebuilt from trainer_version restored cursor : −1 (set_dispatch_index(trainer_version=0)) after restore : {0: [0, 2, 9], 1: [6, 7, 8], 2: [3, 10, 11]} discarded on restore : prompts 12, 13, 14 telemetry : groups_considered 4 groups_redispatched 1

What decides it, and what it means for the counters

A re-admitted batch is stamped dispatch_index + 1. Nothing in the buffer can carry that stamp: the borrow only takes from a step that was already admitted (promote_ready_group picks among slots that exist), and the repayment spare takes the lender’s stamp (:1987). So count_for_target_step is 0 for the new step, and the discard never runs. The cursor is written with every checkpoint — rollout snapshots at :3710, train-boundary saves at :4296 — and restored at :479. The rebuild from trainer_version at :477 only runs for a checkpoint that has no saved cursor.

What this means for the PR. On every checkpoint this code writes, groups_considered and groups_redispatched are equal — the discard is the only thing that could separate them, and it needs a cursor below stamps already in the buffer. One of the two counters is redundant. And the branch itself should not drop prompts: on a correct cursor it never runs, and on a wrong one it loses data with only a print. Better to raise there, so a bad cursor stops the run instead of quietly shrinking the dataset.

Nothing in CI runs a borrow through a checkpoint, and nothing re-admits a RESERVED batch against a real buffer — every scenario in the recovery matrix builds groups with admitted=True. The test above covers both, and is offered for the PR as is.