Why forest training throws away every restored prefix

Explainer for NVIDIA-NeMo/RL #4506. All permalinks pin the PR head b6aff18f and the Gym version it pins, 5911d628. The review comment is in a pending review; its link works once the review is published.

Before
Turn-level recovery only. A model call still decoding at the checkpoint restarts from the start of its turn, so there is no saved prefix to lose.
This PR
Adds target_level: prefix: the tokens a running call has decoded are saved, and the restarted call continues from them. With token_capture.forest_training: true, the worker drops that saved prefix before using it.
Status
Open — not fixed in this PR. Nothing breaks loudly: every restored call quietly decodes again from scratch. Tracked in review comment C26.

Three pages on this PR, one question each — this one: why a restored prefix is lost when forest training is on.

Background. When a vLLM worker gets a model call from Gym, Gym also sends an admission for that call. After a restore, Gym puts the saved prefix into the admission as its generation_cut. Forest training changes how Gym links a call to the call before it, and that adds a request_replay record to every admission. Dotted words have a short explanation: hover or tap them.

Gym attaches the saved prefix to the call's admission and then deletes its own copy (Gym model.py:449-457, restored_cuts.take), so it is offered exactly once. With forest training, Gym links calls by source history (nemo_gym.py:1174-1176), so the admission carries request_replay. The worker sends any such admission through decide_capture_input (vllm_worker_async.py:2176-2203), which builds a new admission and leaves generation_cut out (capture_context.py:43-89). The worker's own comment says so: "A candidate decision rebuilds the admission without a cut".

Same restored call, two kinds of admission

Restored call: prompt 11 ids, 13-token budget, 5 tokens saved at the checkpoint prefix used prefix lost Normal admission (no request_replay) 1 · Gym admission mode token_in request_replay None generation_cut: 5 ids 2 · worker admission used as sent cut kept prefix checked 3 · request to vLLM prompt 11 + 5 = 16 max_tokens 13 − 5 = 8 4 · result decodes 8 new tokens counted as restored chunk row deleted Gym attaches cut :449 _resolve_usable_generation_cut Forest admission (forest_training: true) 1 · Gym admission request_replay set (source_history) generation_cut: 5 ids 2 · worker decide_capture_input builds a new admission cut dropped 3 · request to vLLM prompt 11 ids max_tokens 13 4 · result decodes all 13 again not counted anywhere chunk row left in TQ source_history :1174 worker branch :2176 decide_capture_input :43 Gym already deleted its copy of the cut when it attached it, so nothing retries it. The row cleanup in step 2 of the normal lane runs only when the worker sees the cut, so the forest lane never deletes the 5-token chunk row.

Nothing is wrong with the tokens the run trains on: the call just decodes its whole reply again, as turn recovery would. But the run pays for prefix mode (the extra snapshot work) and gets none of its benefit. The restore counters stay at zero, so it looks as if no call was ever cut. The chunk rows also stay in TQ until the next restore clears rows nothing points to.

What would fix it — proposed, not in the PR

Nothing below is implemented.

Where the fix lands: RL only, one function. validate_single_controller_config in config.py refuses the combination at setup, next to the existing top_logprobs refusal. No Gym change: Gym is right to hand the cut to every admission; the worker is the one that drops it.

@@ config.py L1799-L1804 · validate_single_controller_config
     if recovery_config.generation_prefix_cuts_enabled and top_logprobs > 0:
         # Sampler top-k is staged from a call's final response only, so a
         # restored generation prefix would have no candidates for its tokens.
         raise ValueError(
             "rollout_recovery.target_level='prefix' does not support "
             "policy.generation.top_logprobs > 0 (loss_fn.score_centering) yet"
         )
+    if (
+        recovery_config.generation_prefix_cuts_enabled
+        and token_capture_config.forest_training
+    ):
+        # Forest admissions go through decide_capture_input, which drops the
+        # restored cut, so every restored call would silently regenerate.
+        raise ValueError(
+            "rollout_recovery.target_level='prefix' does not support "
+            "token_capture.forest_training=true yet"
+        )

With this, a run that sets both fails at setup with a clear message, before any GPU work. The longer-term fix is to carry generation_cut through decide_capture_input so forest runs can use prefix recovery too; the setup refusal is the small, safe step for this PR and can be lifted when that lands.

C26 review comment — config.py:1805

So what. Prefix mode plus forest training is accepted at setup today and quietly behaves like turn mode, with extra snapshot cost and leftover rows in TQ. One setup check turns that into a clear error until the worker can keep the cut.