How a checkpoint can silently drop a batch of prompts on resume

NVIDIA-NeMo/RL PR #3480 — recovering the replay buffer from native TransferQueue checkpoints. All links pinned to 1c33a9a1.

This PR teaches a paused run to come back with its finished-but-not-yet-trained rollouts intact. To do that it now saves two new things: the replay buffer index (which groups are waiting, and which training step each is stamped for) and the sampler cursor (how many batches have been handed out so far). On resume both are restored exactly.

They have to agree. The cursor decides what step the next batch is stamped for, and the buffer says which steps are already covered. The problem is that the two are read about a hundred lines apart, with roughly fifteen await points between them — and the rollout pump keeps running through every one of those awaits.

Where the two reads happen, and what it costs

The cursor is read at t0 and the buffer at t3. Everything the rollout pump does in between (t1, t2) lands in the second read but not the first.

The cursor is snapshotted carefully — it sits in an await-free window with the dataloader position and the spare-prompt pool, and there is a comment saying exactly that. The buffer is snapshotted carefully too, inside the exclusive checkpoint barrier so no commit can tear it. Each half is correct on its own. The two were just never made atomic with each other, and admit() takes no barrier, so it is free to run in the gap.

What that costs on resume

Config for the walkthrough — every number distinct so you can trace each one:

quantity	value
`max_lookahead_versions` (the gate)	1	generation may run one step ahead
`num_prompts_per_step`	4	groups needed to close a step
trainer version at save	7	the step about to be trained
saved cursor	6	read at t0, before the pump moved
stamp the pump assigned at t1	7	collides with the trainer version — that is the bug

On restart the controller restores the cursor to 6 and the buffer with its 4 groups stamped for step 7. Then the rollout pump starts:

rollout pump                                     trainer
────────────────────────────────────────────     ──────────────────────────
reads a fresh batch of 4 prompts
  from the dataloader                            waiting for step 7
admit() → cursor 6 becomes 7
  stamps target_step = 7   ← already covered!
count_for_target_step(7) = 4                     select() finds the 4
  → dispatches 0 of 4 prompts                      RESTORED groups
  → "dropping the rest"                          step 7 closes normally
                                                 ✓ training looks healthy

The guard at single_controller.py:819-831 notices that step 7 is already covered and refuses to double-generate. That part is right. But the four fresh prompts it declined to dispatch were already pulled off the dataloader, and the dataloader does not rewind. They are gone — never generated, never trained, not in any checkpoint.

The failure is invisible. The step closes with a full batch of 4 groups, the loss curve is unremarkable, and no counter records a loss. The only trace is one print saying “dropping the rest” — which reads like the deliberate anti-duplication message it also is. One batch per affected resume, silently missing from the training stream.

The other direction is fine

If the pump admits at t1 but its rollouts are still generating at t3, the buffer saves nothing for step 7 while the cursor says 6. On restart the pump re-stamps step 7, sees count_for_target_step(7) == 0, and dispatches a full fresh batch. It self-heals. The in-flight rollouts are still lost, but that is the already-known, deliberately-deferred in-flight drop — not this bug.

The fix

The invariant to restore is one sentence: the cursor, the dataloader position and the buffer index must all describe the same instant. Two changes, both using the barrier that already exists.

1. Let `admit()` participate in the barrier

The gate poll must stay outside the barrier — it waits for the trainer to advance, which cannot happen mid-checkpoint, so holding a slot across it would deadlock. Only the increment and the stamp need protecting:

# staleness_sampler.py:381  _GatedSampler.admit
  async def admit(self, *, trainer_version_fn):
      while self._dispatch_index >= trainer_version_fn() + self._gate_window:
          await asyncio.sleep(_GATE_POLL_SECONDS)     # stays outside
+     async with self._buffer.checkpoint_barrier.mutation():
+         self._dispatch_index += 1
+         return self._stamp()
-     self._dispatch_index += 1
-     return self._stamp()

mutation() is the shared side of the same DataPlaneCheckpointBarrier that commit() already takes, so a save in progress blocks the increment instead of racing it.

2. Read the cursor inside the exclusive section

The dataloader and spare pool must move with it — the batch at the saved dataloader position is the one the next stamp applies to, so all three have to be captured together:

# single_controller.py:1872  _save_checkpoint
-     save_state.sampler_dispatch_index = self._sampler.dispatch_index
-     dataloader_state = self._dataloader.state_dict()
-     reserve_state = list(self._replacement_reserve)
      ...
# single_controller.py:1979  — the exclusive window that already exists
      async with self._data_plane_checkpoint_barrier.checkpoint():
+         # No admit, no commit, no dataloader read can interleave here,
+         # so these four describe one instant.
+         save_state.sampler_dispatch_index = self._sampler.dispatch_index
+         dataloader_state = self._dataloader.state_dict()
+         reserve_state = list(self._replacement_reserve)
          if self._sampler.supports_buffer_checkpoint:
              replay_metadata = self._buffer.metadata_state_dict(...)

A cheaper alternative, if reshaping the save is too invasive: reconcile on the restore side instead. After loading the buffer, advance the restored cursor to max(saved_cursor, max(target_step in restored groups)). That cannot fix a torn dataloader position, but it does close the drop shown here, and it is a three-line change in __init__.

Does it reproduce?

Not on hardware — nobody has run this. The functional test that would exercise it (grpo_dp_single_controller_tq_recovery.sh) checkpoints on a one-second timeout at step 1, where the pipeline is not yet warm, so it does not reach the state above.

What does exist is an executable model of the two functions involved. It reimplements _GatedSampler.admit and the ready-only filter in metadata_state_dict line-for-line from the source, and drives them through seven warm steps and a checkpoint. It is a model, not the real classes — importing those pulls in torch and Ray:

steady state: trainer_version=7 dispatch_index=6

[t0] save_state.sampler_dispatch_index = 6   (single_controller.py:1872)
[t1] rollout pump admits during the save -> dispatch_index=7, stamps target_step=7
[t2] its 4 rollouts COMMIT before the snapshot
[t3] metadata_state_dict() -> 4 group(s), target_step(s)=[7]

--- restart ---
restored: trainer_version=7, dispatch_index=6, buffer target_step(s)=[7]
rollout pump: admits, stamps target_step=7; count_for_target_step(7) = 4
  -> single_controller.py:822 drops 4 of 4 prompts (0 dispatched)
trainer at step 7: select() finds 4 of 4 group(s) needed

The honest ask before merging: a test that saves with a warm pipeline and a group already stamped for the upcoming step, restores, and asserts the resumed run dispatches a full batch rather than dropping one. No such test exists today — the drop branch of that guard is never exercised, and the sampler cursor and the buffer are never checked against each other.

So what. The whole point of this PR is that pausing a run and resuming it should train on the same data as never pausing. This is a path where it does not, and it fails quietly: a full batch of prompts leaves the dataloader and never reaches training. It only bites when the rollout pump admits during a save — likely, since a save follows a train step, which is exactly when the gate reopens and the pump moves. Worth fixing before merge, because once resumed runs are common the missing prompts are indistinguishable from noise.