How the dataloader cursor advances past in-flight prompts at checkpoint save

Companion illustration for NVIDIA-NeMo/RL #3429 (SingleController checkpoint save/restore) and #3594. All permalinks pin the PR head c039949.

The SingleController path runs a rollout pump (generates completions) and a train pump (consumes them) concurrently, connected by a replay buffer. The cursor moves at batch granularity — the instant for prompt_batch in self._dataloader: yields a batch, the cursor is past it (single_controller.py:375) — while dispatch inside the batch proceeds one prompt at a time, each gated on a buffer/in-flight permit (:398-400). The checkpoint saves that already-advanced cursor (:698) but only ready buffer groups (replay_buffer.py:861).

The three pointers at save time

Save fires at train step T=10 — what each of the three pointers sees train pump step 9 ✓ step 10 ✓ → SAVE fires training_info.json records current_step=10 dataloader B9 done B10 done B11 = T+1: yielded, mid-dispatch B12 … saved cursor train_dataloader.pt says "resume from B12" — B11 is behind it batch B11 (30 prompts) 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 finished — [0:10] + [12] + [15:24] ready=True → SAVED to replay_buffer.pt, meta + payload, non-contiguous shape preserved state_dict :858-871 in flight — {10,11,13,14,24,25} dispatched, generation unfinished: unready → DROPPED skip-unready :861 never dispatched — [26:30) (always a contiguous tail) still sitting in the yielded prompt_batch, blocked on a permit — not in flight, not in the buffer, and already behind the cursor permit gates :398-400 On resume loader continues from B12; buffer holds the 20 green groups. The 10 red/gray prompts are gone: consumed from the loader, absent from the buffer, never re-drawn. Bound per save: ≤ one prompt batch + max_inflight_prompts No corruption, no duplication, no deadlock — a group committed between the two snapshots is behind the cursor AND in the buffer (saved once); permits stay balanced. The failure is pure silent data loss.

A concrete timeline (small numbers)

num_prompts_per_step = 2 (one batch = 2 prompts = 2 groups), max_buffered_rollouts = 4, windowed sampler, save_period = 2. Batch k feeds train step k:

batch 1 ─▶ generated ─▶ trained (step 1) ✓ gone from the buffer batch 2 ─▶ generated ─▶ trained (step 2) ✓ ◀━━ SAVE fires here (save_period=2) batch 3 ─▶ generated, not yet trained → buffer slots 0,1 (ready) batch 4 ─▶ dispatched, generation in flight → buffer slots 2,3 (unready) batch 5 ─▶ yielded, 0 of 2 prompts dispatched — parked on a full buffer batch 6 next read — this is where the saved cursor points Buffer at the instant of the save (capacity 4, full): slot: 0 1 2 3 from batch: 3 3 4 4 ready: True True False False status: SAVED SAVED DROPPED DROPPED
The consequence. The saved cursor says "resume from batch 6" — but batch 4 (slots 2,3, dropped as unready) and all of batch 5 (never dispatched) are behind it. On resume those 4 prompts are neither in the buffer nor re-drawn from the loader: silently never trained on. The bias is non-uniform — long-generation prompts are disproportionately likely to be in flight at a save boundary, so they are disproportionately the ones dropped. The interim fix (persist the pending prompts' DatumSpecs and re-dispatch on resume) is requested in the #3429 review; full partial-rollout persistence is tracked in #3594.