feat(checkpoint): recover active vLLM generation prefixes by macandro96 · Pull Request #4506 · NVIDIA-NeMo/RL

macandro96 · 2026-10-05T16:03:21Z

What does this PR do ?

Adds rollout_recovery.target_level: prefix: at each coordinated Gym checkpoint, every active vLLM model call stages the tokens it has decoded so far, and after a restart the restored call resumes from that prefix and generates only its remaining tail.

Stacked on #5056 (deadline snapshots), which sits on #5049 (batch age), #5048 (rollout trace) and #5011 (rollout checkpointing refactor). The base branch is amahishi/rollout-deadline-snapshot, so this diff shows only the prefix-recovery commits. It is retargeted as that stack merges. The parity test that validates it end to end is in #4507, stacked on this one.

Issues

How it works

Save side. A Gym checkpoint already parks rollouts and commits participant state (turn level). With prefix, the controller also fences terminal token writes on every vLLM worker, and Gym asks each worker to cut its admitted calls:

sequenceDiagram
    autonumber
    participant SC as SingleController
    participant W as vLLM worker
    participant G as Gym
    participant TQ as TransferQueue

    SC->>W: begin fence (terminal writes wait, decoding continues)
    SC->>G: prepare + commit checkpoint
    G->>W: POST /ng-control/v1/generation-cut (frozen inventory)
    W->>W: read every inventory call's new tokens from vLLM, under its output lock
    W->>TQ: stage each call's new tokens as a chunk row
    W-->>G: receipt (durable_prefix | durable_failure per call)
    SC->>TQ: save data-plane snapshot
    SC->>W: end fence

Reading decode progress. Captured calls keep vLLM's own output kind (FINAL_ONLY for non-streaming chat); nothing is copied per decode step. At setup the worker wraps vLLM's OutputProcessor.process_outputs so each output chunk runs under one worker lock. A cut finds each inventory call's vLLM request through the map #5049 keeps from engine request ID to captured model call, takes that lock once per inventory, copies every call's tokens past what it has already staged, all at the same output boundary, and reads their logprobs outside the lock. Reading this way instead of queuing a read onto vLLM's engine loop matters at scale: the same pattern for router-replay route reads measured 52 ms median, 588 ms p95 and 40.68 s worst through the engine loop, against 1.36 ms / 4.24 ms / 15.3 ms read directly (10-node A/B runs). vLLM drops a request's state when it finishes, so _finish_request is also wrapped to keep that state, with its finish reason, until the request's consumer has taken the final output; a call vLLM just finished is still cut whole. With vllm_cfg.flat_logprobs, each sampled token's logprob is read straight from FlatLogprobs' arrays instead of building a dict of Logprob objects per position; for 16k generated tokens that is 0.9 ms instead of 6.6 ms, and the final-output read runs on the worker's HTTP event loop.

Restore side. The restored call's admission carries its chunk keys. The worker rebuilds the prefix, splices it into the prompt, shrinks max_tokens / min_tokens by the prefix length, and prepends the decoded prefix to vLLM's tail text before the response is parsed, so Gym, parsers, and verifiers see one uninterrupted completion. On success the chunk rows are replaced by the canonical terminal row.

Safety properties

Any call that cannot be cut regenerates from its previous boundary. A missing cut costs time, never correctness.
Fence epochs: a fence close that runs after the driver already timed out and released it is a no-op, so the gate cannot stick closed. Fence closes run on their own executor, never queued behind a slow cut.
Cut deadline: a cut is bounded by token_capture.control_timeout_s. If vLLM holds its output lock past it, each call is cut at the tokens it had already recorded; calls the cut does not reach in time are acknowledged as failures and regenerate.
Character boundaries: a cut that would stop inside a multi-byte character is refused; vLLM's tail detokenizer would otherwise echo the prompt into the response.
Sampling settings that a resumed tail cannot honor keep the call uncut: stop strings, multi-token bad_words (one-token bad words are banned at every step regardless of context, so they stay cuttable), presence/frequency penalties, and a per-request seed.
The direct reads touch private vLLM state, so setup refuses anything but GIL-enabled CPython and vLLM 0.29.x (the pinned release).

Changes

Area	What it does
Config / setup	`target_level: prefix` requires the turn-level prerequisites and the vLLM async engine. Megatron inference, multimodal capture, router replay and sampler top-k capture (`policy.generation.top_logprobs > 0`, used by `loss_fn.score_centering`) are rejected at setup; top-k is only staged from a call's final response, so a restored prefix would have none. Workers receive the control bearer token and timeout.
vLLM worker (`vllm_worker_async.py`, new `generation_cut_capture.py`)	Per-call token list with a staged-token count, snapshot gate, generation-cut endpoint (route pinned to Gym's `GENERATION_CUT_ROUTE`), direct progress reads under vLLM's output lock with finished-request retention, prefix rebuild/resume, the restored prefix replayed through vLLM's own detokenizer, and restored-response parsing. Every uncut call is acknowledged with one fixed failure reason and counted.
Controller (`single_controller.py`, `rollout_checkpointing/`)	Fence around the Gym checkpoint; Gym-referenced chunk rows count as expected staging rows; snapshot cadence keeps cutting while a dispatched prefix grows; cut metrics per checkpoint and prefix metrics per step.
Gym environment checkpoint (`environments/checkpointing/nemo_gym.py`)	Holds the generation fence across the snapshot and reports Gym's cut counts.
TQ sink / ledger	Chunk-row staging keys; `preserves_generation_prefix` recovery level.
Tests	Worker cut/resume/fence/progress-read unit tests (fence races, deadline, mid-character prefixes, finished-request retention, uncuttable settings) and `tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh` (full tier in L1).

Commits, in order: the feature, Gym's generation-cut wire names (Ananth), two review rounds, reading progress from vLLM at cut time, single-token bad words, failure-reason counters, direct reads under vLLM's output lock, a cleanup of checks that Gym, setup or a sole caller already make, finding a cut call's vLLM request through #5049's link instead of a second one of our own, and reading sampled logprobs from vLLM's flat arrays.

Lines: source +2789 / −151, unit tests +3121 / −97, functional tests +615 / −9, docs/config +24 / −12.

Usage

token_capture:
  enabled: true
  control_timeout_s: 60.0
rollout_recovery:
  target_level: prefix
rollout_checkpointing:
  restore_mode: latest
  # As with turn level (#5056): a snapshot follows every trainer checkpoint and
  # the checkpoint_must_save_by deadline; set an interval for more.
  snapshot_attempt_interval_s: 5.0
checkpointing:
  save_data_plane: true

Testing

Local, in the CI container image, with no packages installed, on this head:

Unit, --nemo-gym-only: test_vllm_token_capture_hosting.py and test_captured_media.py, 195 passed.
Unit, default lane: test_vllm_captured_call_progress.py, test_rollout_trace.py, test_vllm_utils.py, test_vllm_chat_template_wiring.py, test_vllm_restored_prefix_replay.py, test_gym_environment_checkpoint.py, test_gym_prefix_recovery_parity.py, test_timer.py (237 passed, 1 skipped); test_checkpointing.py (180 passed); test_setup.py prefix/turn/top-k/score-centering/deadline cases (62 passed); test_vllm_generation.py capture and generation-cut route cases (18 passed); test_single_controller_actor.py train-pump cases (37 passed).
Unit, --vllm-only --nemo-gym-only: the bad-words check against real vLLM SamplingParams passes.
Unit, --vllm-only: the flat-array logprob read matches vLLM's own per-position mappings on a real FlatLogprobs; the default-lane test_vllm_utils.py (77 passed) covers repeated tokens, the floor and unpublished positions.
Functional: grpo_async_gym_single_controller_prefix_recovery.sh with test(checkpoint): add Workplace prefix-recovery parity #4507's harness passes on this head (Qwen3-1.7B, 1 node): with flat_logprobs: true (452 s), 8 calls cut and restored, token/mask/reward parity exact, generation logprobs 99.11% bit-identical, max |diff| 0.212; on the commit before the flat read with the default flat_logprobs: false (461 s), 99.10%, max |diff| 0.162. Every regular snapshot cut 8 of 8 calls with 0 failed and 0 skipped, and every failed_cuts_* count stayed 0.
Not re-run on this base: the turn-recovery functional test and full unit lanes.

CI status

No CI has run on #4506 itself yet (bot checks only).

Metrics

Per rollout checkpoint (timing/rollout_checkpoint/*): generation_cut_fence_hold_seconds, and Gym's cut outcome as generation_cut_calls, generation_cut_failed_calls, generation_cut_skipped_calls.

Per training step (generation_prefix/*, totals since the vLLM workers started):

Metric	Meaning
`restored_calls`, `restored_prefix_tokens`	Calls resumed from a restored prefix, and their prefix tokens
`failed_cuts`, `failed_cuts_<reason>`	Uncut calls, by reason: `no_capture_or_completed_state`, `completed_not_staged`, `completed_missing_output_limit`, `completed_missing_finish_reason`, `uncuttable_request`, `observation_error`, `missing_output_limit`, `zero_generated_tokens`, `incomplete_utf8_boundary`, `cut_error`, `control_deadline`
`completed_cuts_kept`, `cut_lifecycle_races`	Finished calls a cut kept whole; calls that finished while being cut
`progress_reads`, `progress_read_calls`, `progress_read_wait_us`, `progress_read_hold_us`	Direct progress reads, the calls they covered, and the total time spent waiting for and holding vLLM's output lock

Fence hold time in the latest functional run (Qwen3-1.7B, 4 prompts × 2 generations, 5 rollout checkpoints): 0.13–0.21 s; total rollout-checkpoint save at most 0.31 s.

Why the snapshot fence

The fence exists only to keep the terminal path from deleting chunk rows that the in-progress Gym commit references. From fence begin (before Gym prepares and cuts) until the TQ snapshot is saved, terminal staging waits, so no terminal PUT or chunk-row delete reaches TQ between the cut and the snapshot. A call that finishes decoding inside the window is cut as an already-terminal prefix, and its response is withheld until the fence is released.

Of the two writes the fence holds back, only the delete is unsafe:

Delete (unsafe): without the fence, a call cut in this checkpoint that finishes before the TQ save would delete the chunk rows Gym's commit references, and the save fails with "Gym checkpoint references staging rows missing from live TQ state".
Terminal PUT (safe): Gym does not reference that row, restore clears unreferenced staging rows, and turn-level checkpoints already accept terminal PUTs during the save.

Moving the delete to the controller does not work: a still-running call's next cut references its earlier chunks, and only the worker knows them. Deferring deletes on the worker and pausing only that queue during the checkpoint window does work. #4655 implements that queue; the follow-up there makes deferred deletion the only prefix path and removes the PUT gate, so this fence goes away and completions no longer wait on a checkpoint.

Known follow-ups

After phase 2 restores the TQ checkpoint, TQ's controller logs dtype mismatch: existing=torch.int64, incoming=torch.float32 once for rollout_staging. It appears in every functional run on refactor(checkpoint): move rollout checkpointing out of the single controller #5011 and later and does not fail the run; not yet investigated.
Measure the direct progress reads on a large run (generation_prefix/progress_read_*); the functional test kills phase 1 before its first training step, so it never logs them.
Prefix recovery with score centering: stage sampler top-k with each prefix chunk and prepend it at the final write.
With async_rl.generation_fleet_health, the fence fans out to every worker, so a dead shard fails the snapshot (and the run after max_consecutive_failures).
Chunk rows from aborted calls stay in TQ until the next restore clears them.
Replace this fence with perf(checkpoint): batch obsolete prefix cleanup in background #4655's deferred delete queue (see "Why the snapshot fence").
Move GenerationCutCaptureMixin into a standalone GenerationCutCapture object built in setup_token_capture, like Gym's RolloutTokenCapture.

🤖 Generated with Claude Code

copy-pr-bot · 2026-10-05T16:05:27Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-10-06T02:48:34Z