Six ways a run can be paused, three samplers, and what survives the restart

Test matrix for NVIDIA-NeMo/RL #3827, which tests #3480 · permalinks pin 2245298 · background: what a restart loses · who owns what.

One prompt group is one prompt with 2 rollouts. Training consumes 3 groups a step. The checkpoint lands mid-run, so the buffer holds a mix: some groups trained already, some finished and waiting, some still generating. The rule every case is checked against: a group the loader has handed out and the trainer has not used must come back — it sits behind the saved loader position, so nothing will ever produce it again.

Each row: the buffer at checkpoint time, then what each sampler gets back trained finished, waiting still generating evicted on purpose FAIL lost what it kept before XFAIL old gap PASS windowed in_order wt_fifo everything finished trained 9-11 · loader at 15 needs back 12,13,14 9 10 11 12 13 14 PASS FAIL FAIL gated samplers lose all three, and they were fully generated one very old group nothing trained · loader at 11 needs back 9,10 9 10 10 is far outside the staleness window PASS FAIL FAIL restore hands it back; evicting it is the sampler's call, not the checkpoint's one batch part-done trained 9-11 · loader at 15 needs back 12,13,14 9 10 11 12 13 14 1/2 0/2 XFAIL FAIL FAIL windowed keeps 13 only two batches at once trained 9-11 · loader at 18 needs back 12-17 9 10 11 12 13 14 15 16 17 1/2 0/2 1/2 0/2 XFAIL FAIL FAIL one dropped on purpose trained 9,11 · loader at 15 needs back 12,13,14 9 10 11 12 13 14 1/2 0/2 XFAIL FAIL FAIL 10 was thrown away on purpose, so it is not owed back trained out of order trained 9,10,11,13 · loader at 15 needs back 12,14 9 10 11 12 13 14 1/2 0/2 XFAIL PASS PASS 13 was trained while 12 was still going. The gated samplers pass here because nothing finished-and-waiting was left, so there was nothing for them to lose. 10 FAIL · 4 XFAIL · 4 PASS FAIL means this change lost a group the old code brought back. Those are plain failures, not expected ones. XFAIL means a gap that was there before too, so it is marked and tracked. Only in_order and weight_fifo ever FAIL, and only where something was finished and waiting.

Where the two failure lines come from

FAIL — in_order and weight_fifo say supports_buffer_checkpoint = False, so no index file is written (single_controller.py:981) and the restore returns straight away (:322). Before this change the buffer was saved every time with no such check (merge base :802) and restored for any sampler whose name matched (:270).

XFAIL — a group that has not committed is never written down. reserve() marks the slot not-ready (replay_buffer.py:910) and the save skips anything not ready (:1093). A group is either reserved or committed with nothing in between, so 1/2 and 0/2 are the same thing to a checkpoint. The old code did this too, so it is a gap, not a step backwards.

Does it work?

These were run, on CPU, in about 12 seconds. Real output:

10 failed, 19 passed, 4 xfailed in 11.69s AssertionError: in_order/lag1-next-step-complete: ['g12', 'g13', 'g14'] came back before this change and do not now. These groups finished generating and committed, so the old save wrote them and the old restore returned them on a same-sampler resume. recovered=[]

That is 33 tests: the 18 badges above, plus 15 smaller checks that pin why each line fails. Zero XPASS, which is the number that matters: every case marked "expected to fail" really does fail, so none of them is quietly passing and hiding a gap. The tests drive the real TQReplayBuffer, the real samplers and the real save and load calls — only the tensor converter is swapped out, so a case can use empty prompt records.

What to take from it. The two PASS badges on the bottom row are the check proving itself: there, the old code recovered nothing either, so there is nothing to lose and no failure is reported. Everywhere else a gated sampler is used, groups that finished generating no longer come back. The rollout tensors are still in TransferQueue after the restore — only the note of which rows form which group is missing — so closing this needs an index, not a re-run of the work.