What prefix recovery keeps that turn recovery throws away

Explainer for NVIDIA-NeMo/RL #4506, part 5 of a stacked series on #5056 (turn-level recovery and deadline snapshots). All permalinks pin the PR head b6aff18. Review comments are in a pending review; their links work once it is published.

Before
Turn-level recovery. A model call that was still decoding at the checkpoint is thrown away; after a restart it runs again from scratch, from the start of its turn.
This PR
Adds rollout_recovery.target_level: prefix. Each checkpoint also saves the tokens every running call has decoded so far, and the restarted call continues from them.
Status
New feature, opt-in. Last round's findings (stop strings, penalties, seeds across a cut) are fixed. Open: with forest training every restored prefix is dropped, and prefix mode cannot start next to restart-only Gym servers. Both are in the review comments. No CI has run on this head yet.

C26 review comment — config.py:1805 C27 review comment — config.py:794

Three pages on this PR, one question each — this one: what prefix recovery adds next to turn recovery.

Background. In a Gym rollout, an agent server (Gym) calls the model many times. NeMo-RL's vLLM workers answer, and store each finished call's token ids as a row in TQ. Every few seconds the single controller takes a rollout snapshot. With turn recovery, Gym saves its agents in that snapshot, so a restart goes back to the last finished turn (turn recovery, explained). Dotted words have a short explanation: hover or tap them.

At each Gym checkpoint the controller first holds back terminal row writes on every vLLM worker; decoding keeps running (nemo_gym.py:86). Gym then asks each worker to cut its running calls. The worker reads each call's tokens straight from vLLM, under a lock it wraps around vLLM's output step (vllm_worker_async.py:616), and saves the new ones as a TQ chunk row (generation_cut_capture.py:677). After a restart the worker reads those rows back, adds them to the end of the prompt, and lowers max_tokens and min_tokens by the same count (vllm_worker_async.py:2225). It then puts the decoded text in front of the reply (:1538). A call that cannot be cut — the deadline passed, the cut lands inside a multi-byte character, or it uses a setting a resumed call cannot honor, such as stop strings or a seed — falls back to the turn path.

Where each path does the work — and where they join

One model call: 13-token budget, 5 tokens decoded at the checkpoint new in this PR same code, both paths work thrown away Turn recovery (before) 1 · checkpoint call decoding: t1 t2 t3 t4 t5 not saved 2 · saved finished turns + Gym agent state cut: None 3 · restart call prompt: 11 ids max_tokens 13 min_tokens 7 4 · vLLM decodes 13 new tokens t1..t5 redone TURN :204 splice :2225 Prefix recovery (this PR) 1 · hold + cut hold row writes Gym POSTs a cut freeze t1..t5 2 · TQ chunk row ids [t1..t5] prefix_count 5 output limit 13 3 · restart + splice prompt 11+5 = 16 max 13−5 = 8 min 7−5 = 2 4 · vLLM decodes 8 new tokens text = t1..t5 + new text hold :86 Gym POST :1031 stage chunk :677 cut endpoint :2483 splice :2225 limits :195 prepend text :1538 5 · terminal row both paths run the same code: finish capture :1813 complete_call :1910 13 output tokens turn: 13 new prefix: 5 + 8 Gym continues the agent's turn What still differs on the prefix path (everything else, including the splice, is shared) • the reply text starts with the decoded t1..t5, and the reasoning / tool parser is fed those ids first • the logged output count adds the 5 restored tokens to the 8 new ones • once the terminal row is saved, the chunk rows are deleted from TQ parser :432 count :258 delete chunks :1859

Same call, two restarts

The numbers in the figure, side by side. They are picked by hand to show the mechanism; no GPU run backs this page.

model call: prompt 11 ids, max_tokens 13, min_tokens 7 — 5 tokens decoded when the checkpoint runs turn (before) prefix (this PR) saved for this call nothing chunk row t1..t5 prompt after restart 11 ids 16 ids (11 + 5) max / min tokens 13 / 7 8 / 2 tokens decoded again 13 (t1..t5 redone) 8 (nothing redone) final output 13 tokens 13 tokens (5 + 8)

Why keep both. Turn recovery needs nothing from the vLLM worker beyond normal token capture. Prefix recovery needs the vLLM async engine and a Gym build with the cut endpoint, and setup rejects it with multimodal capture or router replay (setup.py:1380). Any call it cannot cut simply restarts the turn way, so turn recovery stays the fallback inside prefix mode.

What it costs, and what is still open

The gain is the decode work up to the last cut: a long call keeps everything it decoded before the checkpoint instead of nothing. The price: with prefix on, the "nothing changed, skip this snapshot" shortcut is off whenever a model call is running (snapshots.py:619), so every snapshot interval holds finished replies on all workers and pauses Gym. The only timing so far is from a 4-prompt, 1-node run (fence held 0.13–0.21 s); the review asks for one at scale.

C39 review comment — snapshots.py:628

Open review findings (tracked in the review comments). (1) Forest training drops every restored prefix. With token_capture.forest_training on, the worker rebuilds the admission without the saved cut, so each restored call decodes everything again and its chunk rows stay in TQ (explained here). (2) Prefix mode cannot start next to restart-only Gym servers. Setup tells you to set allow_restart_only_gym_servers, and the config check then rejects that flag because it still only allows turn (config.py:759).

(1) C26 review comment — config.py:1805 (2) C27 review comment — config.py:794