How a checkpoint stops losing the prompts that were still generating

NVIDIA-NeMo/RL #3599 · the loss it fixes is drawn here · all permalinks pin the head 95844b6.

In async GRPO a loader hands out prompts a few at a time, the model answers them in the background, and training picks the answers up later from a holding area (the replay buffer) that only keeps finished ones. A checkpoint saves the loader's place — but that place is already past prompts that are still being answered, and those prompts are in no buffer either. Restart, and they are gone: on one long run, 675 of 736 skipped prompts came from this. The fix numbers every prompt as the loader hands it out (trajectory_collector.py:448), and saves the loader's place from before the unfinished prompts instead of its place now (:544). On restart the loader hands out that stretch again, and each prompt that is already trained or still sitting in the buffer is thrown away as it arrives (:486); the rest get answered again.

One checkpoint, two ways to restart

16 prompts, handed out 4 at a time. The checkpoint is written part-way through. ① When the checkpoint is written trained up to here loader has read this far 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 1st handout 2nd 3rd 4th — not handed out yet A training step takes a whole handout at once, so 4 is finished but has to wait in the buffer for 5, 6 and 7. ② Which place the checkpoint saves the loader's place is written down before every handout — the last few are kept: before 0 before 4 before 8 newest one at or before prompt 4 → before 4 — that one is saved ③ Old way: restart where the loader had read to loader restarts here 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 4 of 16 prompts are gone — the model never sees 5, 8, 9, 11. ④ New way: go back to the saved place loader goes back to here 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Prompts 4, 6, 7 and 10 come round again and are thrown away at once; 5, 8, 9 and 11 get answered again. trained already kept in the buffer, reused as is not handed out yet still being answered when the checkpoint was written lost for good answered again after the restart No prompt is used twice: each one is either already trained, or still in the buffer, or answered again — never two of those.

The fix, in code

1. Number each prompt, and write down where the loader was. Before every handout the code copies the loader's place, then writes a number on each prompt it hands out. The last 512 of those places are kept.

place = dataloader.state_dict() # where the loader is: 1 handout done so far batch = next(loader) # the next 4 prompts n = self._next_nemo_gym_task_index # 4 — the number the first of them gets stamp_task_indices(batch) # the 4 prompts become 4, 5, 6, 7 _dataloader_snapshots.append((n, place)) # "before prompt 4, the loader was here"

→ trajectory_collector.py:404-415, numbering at :448, the list of places at :176

2. When saving, pick the place that matches what training has finished. How far training got is read off the numbers on the prompts it just used. A step counter would not do: when a prompt is allowed to fail, the counter and the real position drift apart.

done_through = max(done_through, max(p["_ng_task_index"] for p in just_trained) + 1) # 4 saved = collector.get_checkpoint_state(done_through) # one call, so nothing moves in between └─ newest remembered place at or before prompt 4 → "before prompt 4" save(saved place, "train_dataloader.pt") # 1 handout done — not the 3 it has really done rollouts["frontier_ordinal"] = 4 # prompts under 4 are trained rollouts["resume_base_ordinal"] = 4 # this file starts handing out again at prompt 4

→ how-far-training-got at grpo.py:4728, save at :5283-5305, picking the place at trajectory_collector.py:544

3. On restart, hand that stretch out again and drop what you already have. The loader repeats prompts the run has already dealt with, so each one is thrown away the moment it arrives. Only the gap reaches the model — and once the repeated stretch is past, the checking stops.

done_through = 4 still_in_buffer = [4, 6, 7, 10] # groups that survived in replay_buffer.pt last_to_check = max(still_in_buffer + [done_through - 1]) # 10 # the loader starts at prompt 4 again: 4,5,6,7 | 8,9,10,11 | 12,13,14,15 for batch in loader: if lowest number in batch > 10: # 12,13,14,15 — past everything we know about stop checking; use the batch # normal from here to the end of the run keep = numbers that are >= 4 and not in still_in_buffer # 4,5,6,7 → keep [5] (4, 6 and 7 are in the buffer) # 8,9,10,11 → keep [8,9,11] (10 is in the buffer)

→ trajectory_collector.py:486-542, restart setup at grpo.py:4317-4340

Set checkpointing.load_replay_buffer: false (grpo_math_1B.yaml:124) and the buffer starts empty, so still_in_buffer is [] and 4, 6, 7 and 10 get answered again too. Same code either way.

Does it work?

Nobody has posted a run. There is no W&B link, no reward curve, no before-and-after prompt-coverage count on the PR — so there is no evidence yet that a resumed run learns better, or that the ~30% of prompts skipped on the long run that started this actually goes to zero. What exists is that first measurement (675 of 736 skipped prompts, traced to two causes, before the fix) and one test that runs a save like the one above both ways (with its own numbers, fed straight to the filter): the old way skips exactly [5, 8, 9, 11], the new way skips nothing and repeats nothing (test_async_utils.py:1530). The right thing to ask for before merging is one resumed long run with the skip count reported.

Two things to know before you restart a run. The prompts come back, but the answers do not: a prompt that is answered again gets a fresh answer from the newer weights, because a half-finished answer cannot be saved. You get the same prompts in the same order, not the same bytes. And if you restart from a checkpoint written before this PR, that run keeps the old, lossy behaviour for good — and so does every run started from the checkpoints it writes (grpo.py:4346 logs a warning). Only runs started fresh after this PR are protected from their first save.