Repository navigation
fix(token-capture): use next-turn prefill routes at R3 turn boundaries - #4579
Closed
pjin-nvidia wants to merge 1 commit into
Closed
pjin-nvidia wants to merge 1 commit into
pjin-nvidia wants to merge 1 commit into
Conversation
With token capture, each turn of a multi-turn rollout is a separate vLLM call. A call never runs its last generated token back through the model, so its aligned routes carry only a placeholder for that token (arange(topk) at every layer, not the -1 missing-route sentinel), and Router Replay forced those experts in the Megatron forward passes. Every later position attends to that token, so each intermediate turn boundary was trained with a route that inference never used. The next call's prefill computes the real route. The non-TQ NeMo-Gym path already splices it in (nemo_rl/environments/nemo_gym.py); this does the same for TQ token capture, for both direct route assembly and token_capture.defer_routed_experts_to_policy. - Capture: a token_in call that generated tokens also returns the prefill route of the previous turn's last token (routed_experts_boundary, [1, layers, topk]) and its absolute index (routed_experts_boundary_index = prev_len - 1). The Gym adapter copies both into the call's extras, so the extras digest commits them. Token IDs, masks, logprobs, and delta lengths are unchanged. - Staging: TQTokenSink checks the sidecar (token_in child, index == prev_len - 1, same dtype and per-token shape as the delta routes) and stores it in a routed_experts_boundary column registered on the R3 staging partition; the index stays in the extras metadata. Only rows with a boundary write the column, and only those rows read it. - Assembly: route plans move to schema v3 with per-span boundary_token_index and repair_previous_token. The reassembler allows a repair only along the receipt's selected chain and only between two spans that both generated tokens, so abandoned retries and response order cannot change a parent. The repair edits the assembled tensor, never the parent's staged routes. A missing or corrupt sidecar follows the existing failure policy: direct assembly rejects the rollout; deferred assembly uses native routing for it. The trajectory's final token keeps its placeholder. - Compatibility: v2 plans and rows staged without a boundary remain readable; the latter keep the placeholder and log a warning. With R3 and token capture, the vLLM worker fails at startup on a Gym adapter without boundary support, data-plane checkpoints record routed_experts_boundary_schema_version, and resuming a native TQ checkpoint without it fails at startup. Staging-schema migration is not implemented: a restored Mooncake controller would register the new column lazily while producers are live. - docs/guides/router-replay.md describes the behavior and upgrade path. - Tests cover capture -> TQ -> direct/deferred replay for 1-3 turns (int8/int16 routes, all -1 boundaries), retry selection, tampered or missing sidecars, a data-plane checkpoint round trip, legacy rows, v2/v3 plan decoding, the Gym adapter gate, and checkpoint metadata/restore checks. Bumps Gym to ad46d9f6b: main's pin 17abf2ddc plus "feat(capture): commit the vLLM turn-boundary route in call extras". Rebased onto VLM media capture (#4191): the staging partition registers the boundary column next to the media columns, stage() derives the media checksum and the boundary index from the same metadata JSON, and fetch_for_finalization() verifies media metadata before parsing the boundary index from it. Signed-off-by: Peter Jin <pjin@nvidia.com>
github-actions
Bot
added
the
Documentation
pjin-nvidia
force-pushed
the
pjin/routing-replay-v2-fix-20260928
branch
from
October 8, 2026 00:14
9dce0ff to
e39475d
Compare
Contributor
Author
|
closing, merged #4610 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
With token capture, each turn of a multi-turn rollout is a separate vLLM call. A call never runs its last generated token back through the model, so its aligned routes carry only a placeholder for that token (arange(topk) at every layer, not the -1 missing-route sentinel), and Router Replay forced those experts in the Megatron forward passes. Every later position attends to that token, so each intermediate turn boundary was trained with a route that inference never used. The next call's prefill computes the real route. The non-TQ NeMo-Gym path already splices it in (nemo_rl/environments/nemo_gym.py); this does the same for TQ token capture, for both direct route assembly and token_capture.defer_routed_experts_to_policy.
Bumps Gym to ad46d9f6b: main's pin 17abf2ddc plus "feat(capture): commit the vLLM turn-boundary route in call extras".
Rebased onto VLM media capture (#4191): the staging partition registers the boundary column next to the media columns, stage() derives the media checksum and the boundary index from the same metadata JSON, and fetch_for_finalization() verifies media metadata before parsing the boundary index from it.
Also see: NVIDIA-NeMo/Gym#4201