feat(vllm-model): on-policy sampling pin via sampling_overrides - #2127
Closed
ananthsub wants to merge 5 commits into
Closed
feat(vllm-model): on-policy sampling pin via sampling_overrides#2127ananthsub wants to merge 5 commits into
ananthsub wants to merge 5 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
ananthsub
changed the title
This was referenced Jul 23, 2026
ananthsub
marked this pull request as ready for review
July 23, 2026 16:18
github-actions
Bot
added
the
sla:triage-overdue
ananthsub
force-pushed
the
ananthsub/tokidcap/delivery
branch
from
July 29, 2026 06:55
74a48d6 to
71de7c5
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 06:55
4a07f8d to
774339d
Compare
Claude Code makes model calls that are not part of the rollout: it generates a conversation title and probes quota. They reach the model server on the same rollout-prefixed URL and get captured, and because they are genuine policy output -- real token ids, real log probs -- nothing downstream can tell they do not belong. Training on them optimizes the policy to write conversation titles under the rollout's reward. The record now keeps what the *harness* asked for (requested_model, has_tools), read off the parsed request body at the handler rather than by touching the body again in middleware. That is the signal, because a harness asks for a small model for these calls even though the server serves one model. Classification uses two signals and needs no harness-specific code in the core: an optional explicit pattern list for deployments that know their harness, and self-calibration -- whichever model generated the most tokens in a rollout is the policy model, and calls asking for a different one are side calls. Records written before this field existed carry an empty requested_model and are all kept, so nothing changes for them. Excluded calls are reported (side_calls_excluded) rather than silently dropped, and a rollout whose calls were *all* side calls is masked instead of yielding an empty trajectory. This is the second half of the title-call problem. The first was structural: a short side call became the main chain and the real rollout was dropped, fixed by selecting on generated-token mass. Even with the right chain selected, the side call would still have been stitched in. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The builder infers lineage after the fact, by finding the earlier call whose token sequence is the longest prefix of this call's prompt. That cannot do two things. It cannot resolve a retry: capture records a call once the response is assembled, including one the client never received, so a retry leaves two records with the same prompt and different generations, and both are equally valid children. And it cannot run before the call, which is where supplying the engine an exact prefix has to happen. This resolves the parent at request time using only what the harness already sends. A harness must echo the conversation to continue it, so the assistant turns in a request are the ones we produced; hashing them in order identifies the call that produced the last one. Nothing is added to the wire and nothing depends on the harness preserving a field we invented. Tool-call arguments are canonicalized (sorted-key JSON) for comparison only, because harnesses re-serialize them between turns -- compact one turn, pretty the next. Without that every tool-using turn would miss. The model's original argument string is what stays in the record. The index is a map keyed by call, not a running cursor, so it is a tree: two sub-agents branching from one parent both resolve to that parent and both would get the same prefix. A cursor would hand the second branch a prefix containing the first branch's generation, which the splice would apply without complaint. Entries are added and never mutated, so concurrent sub-agents cannot corrupt each other's lineage. Three cases deliberately return no parent rather than a guess: a new or rewritten conversation, an ambiguous match (two recorded calls with byte-identical output), and an evicted rollout. Each falls back to prefix inference, which is what happens today. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Without this the engine builds every prompt by re-rendering the conversation through the chat template. Re-tokenizing an assistant turn can produce a different split than the one sampled, and for a reasoning model the template drops earlier thinking entirely. Either way the new prompt does not extend the previous prompt-plus-generation, the builder cannot chain, and only the first call is trained on. Supplying the parent's cumulative ids as required_prefix_token_ids makes NeMo-RL's splice keep them verbatim and append only the newly rendered tail. The harness does not send this field and could not. Gym constructs the outbound engine request itself, so it is injected in _preprocess_chat_completion_create_params next to the sampling pin. The inbound request supplies identity; the outbound request supplies payload. The splice applies whatever it is given without checking that it belongs to this conversation, so supply fires only on a unique, verified parent and otherwise forwards the request untouched. Each record carries prefix_supplied and the server tracks a supplied/eligible ratio, so a run can be audited afterwards rather than inferred from whether chains happen to be contiguous. Off by default; it needs a backend that honours the field. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The route that serves a rollout's raw training tokens is registered on the same app the harness calls to generate. That is acceptable inside a trusted cluster. It is not acceptable once the harness runs in a sandbox whose only egress is this server, because the harness could read its own training data -- or another rollout's. token_id_capture_read_token requires a bearer token on the route, compared in constant time. When it is unset the route stays open and warns once, so existing deployments keep working and the gap is visible rather than silent. This is the one piece of sandboxing that has to land before the sandbox work, not with it. Everything else about capture is unaffected by sandboxing: capture happens in the model server, outside the sandbox, and the sandbox only ever sees text. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Add a framework-agnostic sampling pin to vllm_model: sampling_overrides forces temperature/top_p on every chat request, read from generic policy_generation_* keys with on-policy defaults. This keeps an external harness's rollouts on-policy for training without Gym holding any framework-specific knowledge — the training framework publishes its sampling into the generic keys. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/delivery
branch
from
July 29, 2026 07:12
71de7c5 to
5c317c1
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 07:12
774339d to
7342d4e
Compare
ananthsub
marked this pull request as draft
July 29, 2026 07:12
github-actions
Bot
removed
the
sla:triage-overdue
Contributor
Author
|
Superseded: recreating this PR with the correct stack base ( |
This was referenced Jul 29, 2026
This was referenced Jul 29, 2026
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
1 participant
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.
On-policy sampling pin in the vLLM model server
Fourth PR in the stack (base: #2126).
Problem
An external agent picks its own sampling settings when it calls the model. During training the generation worker checks that requests use the trainer's sampling settings, so a
temperatureortop_pthe agent chose that doesn't match will either fail the run or bias the recorded rollout.What it adds
sampling_overridesinvllm_model: settemperatureandtop_pon every chat request before generation, replacing whatever the agent sent.policy_generation_temperature,policy_generation_top_p) with on-policy defaults. Gym doesn't know anything about the trainer; the trainer writes its settings into those keys (see the matching NeMo-RL change) and Gym just reads them.Notes for review
vllm_modelnever references NeMo-RL. Any training framework can drive this by filling in those config keys.Tests
No isolated unit test — the
vllm_modelserver test runs in its own environment. Verified in the 2-GPU Megatron GRPO run: sampling matched the trainer and the run stayed on-policy (generation KL ~0.04).Stack