What happens to a training step when a prompt is given up on

Background for NVIDIA-NeMo/RL #3925. All permalinks pin a4c10a78. This is not a review finding — it explains a pre-existing, opt-in mechanism that came up while reading the recovery counters.

Three pages from the #3925 review — this one: what a spare prompt is for.

Everything on this page happens during a live run, not at restore. The code is in _dispatch_one_prompt inside the rollout pump, and it fires the moment a prompt exhausts its retry budgets mid-training — on a fresh run from scratch, with no checkpoint involved. Checkpointing does not change any of it: the same borrow runs with save_data_plane on or off, and the round-trip test exercises it with recovery on.

A training step consumes a fixed number of prompt groups. If a prompt is retried until its budget runs out and is finally given up on, that step is one group short. Something has to give, and on_dropped_prompt picks what:

on_dropped_prompt: Literal["shrink", "replace"] = "shrink" # ← the default

By default nothing described below happens. The step simply trains on fewer groups. Everything on this page is what you opt into with "replace", which exists to hold the batch size constant instead.

The hole, and the two ways to fill it

Step 7 needs 3 groups. One prompt was given up on. shrink (default) step 7 trains on 2 groups replace, no finished later group to take — dispatch a spare for the same step step 7 waits for the spare, then trains on 3 replace, a later step already finished — borrow now, repay later step 7 closes immediately, using a finished group taken from step 9 step 9 (the lender) step 9 is now one short …so the spare is dispatched for step 9 instead, which repays the loan borrowed group moves back to step 7

Where spares come from

A spare is a prompt the dataloader already produced that is deliberately not made into a training step. When the pool is below its low-water mark, an entire batch is diverted into it:

spare pool refilled with N prompt(s) (low-water mark replacement_reserve_prompts); this batch is not admitted as a training step

_divert_batch_to_reserve

The docstring is explicit that diverting must happen before admission, because “admitting a batch and then dispatching nothing for it would leave a target step that no group is ever generated for, which is exactly the hang the shortfall accounting exists to prevent.” Unused spares are not wasted — nothing was generated for them, and they stay in the pool.

knobdefaultmeaning
on_dropped_prompt"shrink"Whether any of this runs at all.
max_replacement_attempts1How many substitutions one step may make. Read only under "replace".
replacement_reserve_prompts1Low-water mark for the spare pool. Read only under "replace".

The borrow runs with checkpointing on or off

“Recovery on/off” is not a knob, and it is not about save/restore time. It is a flag computed once when the controller starts. It decides where in the code the borrow happens — not whether it happens:

self._rollout_recovery_enabled = bool( restoring_rollout_recovery # resuming a checkpoint that has a recovery payload or ( self._master_config.checkpointing["enabled"] and self._master_config.checkpointing.get("save_data_plane") and self._sampler.supports_buffer_checkpoint # True for all four shipped samplers ) )

single_controller.py:497-504

With the flag on, the borrow happens inside the checkpoint-atomic mutation cut (:1983), because the dropped prompt’s ledger record and the spare’s reservation have to change together. With it off, the same call sits a few lines lower (:2020). Either way _promote_into_step runs and the spare is sent to the lender. This is the config that gets you the borrow at all:

checkpointing: save_data_plane: true # or false — makes no difference to the borrow async_rl: sampler: name: in_order max_lookahead_versions: 1 # ≥ 1; at 0 no later step can have finished, so there is nothing to borrow rollout_failure: on_dropped_prompt: replace # required; at the default "shrink" none of this runs

The spare and the borrow are not a package. The link runs one way only:

needs the other?why
borrow → spareyes, always The borrow is a debt and the spare is the repayment. The code will not borrow without a spare in hand — “Borrowing without one would leave the lender short instead: the same hole, one step later.”
spare → borrowno The borrow needs a finished group in a later step to take. When there is none — lookahead 0, or nothing later has finished yet — the spare is dispatched for the same step and that step waits for it.

So every "replace" run uses the spare pool, and borrows whenever a later step has a finished group to lend. There is no configuration that borrows without a spare pool.

First, where a spare comes from. The pool is a separate holding area, filled by taking a whole dataloader batch instead of admitting it as a training step:

dataloader one batch = 3 prompts SPARE POOL not admitted as a training step — these 3 prompts skip the queue entirely Refilled the same way whenever the pool falls below replacement_reserve_prompts (default 1), so it is topped up on demand rather than filled once. Nothing is ever put back into it.

Now the two cases, drawn as the same five moments. Only step 3 differs — and what picks the case is whether a later step already has a finished group, not any setting.

no finished later group → the spare fills step 7 step 9 already finished → borrow 1 step 7 loses a prompt step 7 1 step 7 loses a prompt; step 9 already finished step 7 step 9 2 take one spare from the pool (3 → 2 left) 2 take one spare from the pool (3 → 2 left) 3 dispatch the spare FOR STEP 7 it starts generating; step 7 cannot close yet 3 move a FINISHED group from step 9 into step 7 the spare is not dispatched yet 4 step 7 waits for the spare to finish 4 step 7 full — trains immediately step 9 now one short — it lent one out 5 step 7 spare arrives — trains 5 dispatch the spare FOR STEP 9; when it lands: step 9 loan repaid step 7 work step 9 work spare step 7 pays the wait step 9 pays it instead — and step 9 is due later

Same number of rollouts either way, the same prompt is given up on, and the same one spare leaves the pool. The borrow only moves when the waiting happens — onto the step with the most slack, which is why promote_ready_group picks the furthest future step as lender.

Step 9’s groups are drawn in a second colour on purpose: when one lands in step 7 at moment 3, you can see that a single group moved, and which one. The two steps were not swapped, and step 7’s own two groups never change.

Dispatching a spare still means waiting for it to generate. The borrow removes that wait: take a group that is already finished from a later step, re-stamp it to the step with the hole so that step can close now, then dispatch the spare stamped for the step you took it from.

# with rollout recovery on — inside the checkpoint-atomic mutation cut if replacement is not None: lender_step = self._promote_into_step(target_step) # :1983 if lender_step is not None: target_step = lender_step # with it off — the same borrow, a few lines lower if not self._rollout_recovery_enabled: lender_step = self._promote_into_step(target_step) # :2020 if lender_step is not None: target_step = lender_step # the spare now belongs to the lender

single_controller.py:1983 and :2020 — _promote_into_step is the only caller of promote_ready_group in the repo.

The comment on promote_ready_group is where the debt idea is spelled out: the loan “must” be repaid, because “an unrepaid loan is the same hole one step later, carried forward until it reaches the last step, which has nobody to borrow from.” The furthest future step is chosen as lender because it is due last and has the most slack.

Does a borrow put a restore at risk? No. The borrow moves a group between steps that were both already admitted, and the spare takes the lender’s stamp, so nothing is ever stamped past the sampler cursor. After a restore, a re-admitted batch gets the next stamp after the saved cursor, finds nothing there, and dispatches in full. The third page runs exactly this through a checkpoint and back.

Is it worth keeping? An honest summary rather than a recommendation, since this is pre-existing and not part of #3925:

For: a constant batch size is a real property to want — under "shrink" a fleet wobble silently changes the number of groups a step trains on, which moves the effective batch size around without saying so.

Against: it is three knobs, a spare pool that costs a whole diverted batch, and a debt-and-repayment protocol whose ordering comments run longer than the code. The borrow only ever runs for in_order runs that opted into "replace" with lookahead ≥ 1 — a narrow corner to carry that much machinery for. If it were dropped, the spare substitution could stay and only the borrow/repay would go, which is where nearly all the complexity lives.