feat(token-id-capture): keep harness side calls out of the trajectory - #2179
Closed
ananthsub wants to merge 4 commits into
Closed
feat(token-id-capture): keep harness side calls out of the trajectory#2179ananthsub wants to merge 4 commits into
ananthsub wants to merge 4 commits into
Conversation
An external agent harness returns no token ids, and the ids do not survive to the client for a streamed response or one translated to Anthropic Messages. The model server holds the assembled response with the ids on it for a moment before either of those happens, which is the one point that covers every dialect. nemo_gym/token_id_capture/ records one TokenEntry per correlated model call: prompt ids, generated ids, one log prob per generated token, and the assistant text and tool calls, since a trainer reads the text for its own penalties. Records go to <rollout_id>.tokens.jsonl, per-file flock and fsync, awaited so a record is durable before the call returns. Nothing is added to the client response. TokenSink and TokenSource are the write and read seams. Gym owns the record shape and the capture code; a framework supplies the implementation and runs it where its tokens are produced, so sink placement is a deployment choice rather than a fork in the design. The package is a leaf -- no fastapi, ray, uvicorn, aiohttp or torch -- so a framework's inference worker can import it; a subprocess test enforces that. TokenEntry also carries optional parent_call_id, cum_len and digest. cum_len and digest are stamped at capture; parent_call_id stays null until the model server can resolve a parent. A capture failure is logged and also writes a <rollout_id>.tokens.incomplete marker, so a rollout that lost a call is distinguishable from a complete one. Capture stays best-effort: a bad payload must not break the harness run. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The builder is a pure function over a rollout's TokenEntry records. per_request emits one sequence per call. prefix_merging chains calls by the token-prefix relationship, parenting each call to the earlier call whose prompt-plus- generation is the longest prefix of this call's prompt, so an append-only multi-turn rollout becomes one chain. A prompt that extends nothing starts a new root, which is what a compacted or rewritten context looks like. Both are order-independent. Loss masks follow provenance: generated tokens are trained with their captured log probs, and anything re-fed into a prompt is not. The projection re-emits contiguous Responses items carrying both content and token ids, which is what NeMo-RL's postprocess already consumes. Main-chain selection is by generated-token mass across all roots, not by the first root. Entries are processed in increasing prompt length, so the first root is whichever root has the shortest prompt; a rollout's own first call is large because of the harness system prompt, so an auxiliary short-prompt call would be selected instead and the rollout dropped at delivery without an error. Recorded parent links are used when present and verified by digest rather than trusted, falling back to prefix inference on mismatch. A retry of the final call is reported unresolved rather than tie-broken, since nothing can say which generation the client received. Chain count, quarantined fraction and delivered fraction are returned rather than discarded, and a malformed capture returns an unbuilt result instead of raising into the caller's loop. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…med records Rollout collection replaces response.output with the merged, contiguous items for agents that opted into capture, so NeMo-RL reads response.output the same way for native and external-harness rollouts. Native agents are excluded by the per-agent opt-in: they already return exact ids inline, and a rebuild could differ from what the model server returned. Retention runs in both directions. Consumed records are deleted once folded into response.output (NG_KEEP_TOKCAP retains them), and stale records are cleared before dispatch. Both are needed because rollout ids are deterministic and the store appends, so a rerun would otherwise stitch a previous attempt's calls together with this one's. The build's counts ride the record under _ng_token_capture. Without them a rollout that trained on one of five calls is indistinguishable from one that trained on all five. A rollout captured incompletely, or whose final call was retried ambiguously, is marked for masking and warns. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
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>
|
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. |
This was referenced Jul 29, 2026
ananthsub
force-pushed
the
ananthsub/tokidcap/delivery
branch
from
July 29, 2026 12:38
5c317c1 to
58cde9d
Compare
Contributor
Author
|
closing out to restack the side call filtering as an extension on the main stack |
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.
Keep harness side calls out of the trajectory
Part of the token-id capture stack for #1824.
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.
What it adds
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: a harness asks for a small model for these calls even though the server serves one model.side_calls_excluded) rather than silently dropped.Notes for review
requested_modeland are all kept, so nothing changes for them.Tests
tests/unit_tests/test_token_id_capture.py,tests/unit_tests/test_trajectory_builder.pyStack
One commit per PR, each based on the previous branch (bottom of the stack targets
main):All nine are drafts. #2183 supersedes #2127, which could not be retargeted after
the stack grew (GitHub forbids changing the base of a PR that is part of a stack).