Why resuming an in_order run throws away the batch it was about to train on

A review finding on NVIDIA-NeMo/RL #3480 (recover replay buffer from native TransferQueue checkpoints), open at the time of writing. All permalinks pin the PR head 608e9bd; "before" links pin its merge base 7cba25a. For how the pieces are laid out across processes, see which process owns which piece of the checkpoint.

Training here runs two loops at once. A rollout pump takes prompts from the dataloader and generates answers; a train pump consumes finished answers a batch at a time. Finished-but-not-yet-trained answers wait in a replay buffer, and the pump is allowed to run a set number of batches ahead of the trainer — max_lookahead_versions for in_order, max_staleness_versions for weight_fifo (staleness_sampler.py:353). Both default to 1, which is what the shipped SingleController config uses, and the example below. A checkpoint has to save that buffer, or the answers in it are lost. Before this PR every save wrote the buffer's own contents to replay_buffer.pt, whichever sampler was in use (before: self._buffer.state_dict(...)). Samplers themselves are never saved — they have no state_dict, before or after this PR. This PR only saves the buffer when the sampler declares it can restore it (now: :981), and only windowed declares that (staleness_sampler.py:257).

What the restart loses

3 prompts per training step, generation runs 1 step ahead. Save fires after step 4. at save step 4 done 9 10 11 12 13 14 15 16 17 loader cursor = 15 saved to train_dataloader.pt already trained finished, waiting for step 5 not read from the loader yet The 3 orange prompts are done generating and are stamped for step 5. The loader has already moved past them. resume before 9 10 11 12 13 14 15 16 17 buffer restored; step 5 trains 12,13,14 The saved buffer comes back. Its prompts are stamped for step 5, the step about to run, so nothing is skipped. resume with #3480 9 10 11 12 13 14 15 16 17 buffer empty; step 5 trains 15,16,17 generated, paid for, never trained on — and never generated again No buffer file was written, so there is nothing to restore. The cursor still says 15, so prompts 12, 13 and 14 sit behind it and are never handed out again. The run reads 15, 16, 17 instead and keeps going. The dataset is read once (max_num_epochs: 1 in the shipped config), so there is no second pass to pick them up. Nothing is printed at any point. No error, no warning, no metric.

The lost prompts are not spare capacity that happened to be sitting around. The step that just finished consumed everything stamped for itself, so what is left in the buffer is the work queued for the steps after it — and before this PR those groups came back selectable, because InOrderSampler.select picks groups whose stamp equals the step being run, and a resume restores that step number unchanged.

How much that is depends on your config, so the 3 prompts above is the smallest interesting case, not a fixed cost. The gate lets generation run max_lookahead_versions batches ahead, and the buffer is sized num_prompts_per_step × (max_lookahead_versions + 1) — one live batch plus one per lookahead version (staleness_sampler.py:321). So at the shipped default of 1 you lose one batch per restart; raise the lookahead and you lose that many batches. Set it to 0 and nothing is buffered, so nothing is lost. A separate hard cap, max_buffered_rollouts, can hold the number below what the gate would otherwise allow.

The fix, and why it goes at startup

There are three places the code could tell you. Only the first one is any use.

# 1. actor startup — both facts are known, and nothing is lost yet checkpointing.enabled = True sampler = in_order # supports_buffer_checkpoint = False # the last moment the user can still change the config

put the warning here, next to the check that already lives at single_controller.py:186

# 2. save, after step 4 — too late, and it happens every save if self._sampler.supports_buffer_checkpoint: # False replay_metadata = ... # skipped # no buffer file written; prompts 12, 13, 14 are now unrecoverable

single_controller.py:981

# 3. resume — nothing left to report if not self._sampler.supports_buffer_checkpoint: return # buffer empty, cursor = 15, count = 0 # a warning here cannot say how much was dropped — the file was never written

single_controller.py:322 — this is the line that returns in silence today

So the suggested change is one print at startup, when checkpointing.enabled is on and the chosen sampler cannot restore a buffer. It costs nothing, it fires once, and it fires while the user can still switch to async_rl.sampler.name=windowed or turn checkpointing off.

What about rollouts that were still running?

They are lost too, and this one hits every sampler — windowed included. It is not a hidden bug: the PR says so itself, listing "in-flight reservations that have not committed canonical TQ rows" under what it does not yet recover, with #3456 named as the work that would close it. Worth drawing anyway, because it stacks on top of the loss above.

The first figure kept things simple by showing 15, 16 and 17 as unread. In a real run the pump has usually handed some of them out already — that is what max_inflight_prompts allows — so the cursor sits further along. Same save, same step, just drawn with the in-flight prompts included.

The same run, with the in-flight prompts drawn in. Each prompt needs 2 generations. at save step 4 done 12 13 14 15 16 17 18 19 20 1/2 2/2 0/2 cursor = 18 committed, waiting for step 5 still generating — never saved not read yet The three red prompts were handed out and are part-way through. 16 has even finished both of its generations, but has not committed yet. None of them is in the saved index, and the cursor is already past them, so a restart never hands them out again. Every sampler behaves this way.

The reason is a single flag. When the pump hands a prompt out it books a slot marked not-ready (ready_list.append(False)), and the slot only flips to ready once every generation is done and written (ready_list[idx] = True). The save walks that list and skips anything not ready (metadata_state_dict); its docstring puts it plainly — "In-flight reservations are intentionally omitted." So prompt 16 above, fully generated but not yet committed, is skipped exactly like prompt 17 that has barely started.

Nothing compensates for that on the loader side. The cursor is written down exactly where it stands (self._dataloader.state_dict()), and _save_checkpoint never looks at the in-flight bookkeeping it already keeps in _inflight_by_group_id. It does not wait for those prompts, does not record them, and does not move the cursor back before them.

How this differs from the same question on #3599. That one is about async GRPO, which does try to protect in-flight prompts: it saves the loader's place from before the oldest unfinished one. The gap there is narrow — a tolerated failure lets a step be refilled from later prompts and pushes the saved place past prompts still being answered, and it needs several settings lined up to happen. SingleController has no such protection to defeat. There is no rewind, so every in-flight prompt at save time is lost, every time, on any config. Simpler to reason about, and much easier to fix later: the information needed is already sitting in _inflight_by_group_id.

Does it work?

Nobody has run it. There is no resumed-run link, no before-and-after count of skipped prompts, and no curve. That is not an oversight in the review — the feature's own end-to-end test never executed either. grpo_dp_single_controller_tq_recovery.sh is registered as a full-mode test, the PR carries the CI:Lfast label, and the job log for the SingleController lane reads:

+ echo 'FAST: Skipping: uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh'

The workflow has no merge-queue trigger, so the next time that script would run is after the change has already landed on main. Every check on the pull request is green, and none of them touched this code path.

Two things are worth asking for before merge. Relabel the pull request CI:L1 and link the run, so the recovery test actually executes once. And for this finding specifically, resume an in_order run from a checkpoint and report how many prompts were skipped — that number is the whole argument, and right now nobody has it.

What this means for you. If you run SingleController with in_order or weight_fifo and checkpointing on, every restart silently throws away whatever generation had run ahead — one batch at the default lookahead of 1, more if you raised it — along with the prompts behind them, and those prompts never come back. It is quiet: the run resumes, the loss curve looks fine, and the only sign is that your model saw slightly less of the dataset than you think. Before this PR the same restart kept them. Until a warning lands, windowed is the only sampler where a restart is lossless.