Avoid redundant vLLM /tokenize calls using native prompt token IDs - #2576
Closed
nanz-nv wants to merge 1 commit into
Closed
Avoid redundant vLLM /tokenize calls using native prompt token IDs#2576nanz-nv wants to merge 1 commit into
nanz-nv wants to merge 1 commit into
Conversation
…ned from vllm instead
nemo-automation-bot
Bot
added
the
community-request
Contributor
svcnvidia-nemo-ci
added
the
waiting-on-customer
Author
@ananthsub yes that PR looks great and solves my problem. Do you have an estimation on when that PR can get merged? Any blockers? |
svcnvidia-nemo-ci
removed
the
waiting-on-customer
github-actions
Bot
added
the
sla:triage-overdue
ananthsub
added a commit
that referenced
this pull request
Aug 18, 2026
#1558) ## Summary - Consume a complete token bundle from `choice.message` without making a second `/tokenize` request. - Support opt-in vLLM response token IDs for endpoints that return top-level `prompt_token_ids` and choice-level `token_ids`. - Retain `/tokenize` for integrations that return neither inline representation. - Reject partial or conflicting sources. Token capture supports one completion choice per request (`n` omitted or set to `1`). ## Why Gym needs the prompt token IDs, generated token IDs, and selected-token log probabilities used during generation. The previous adapter reconstructed prompt IDs through a second awaited `/tokenize` request. That request adds inter-turn latency and can reproduce a different prompt when prompt-affecting inputs are omitted. NeMo RL [#3390](NVIDIA-NeMo/RL#3390) attaches the token bundle to the assistant message. Gym reads that bundle directly and does not request duplicate response-level IDs by default. Gym cannot require that message shape from every training integration. Other integrations, including Verl, can request vLLM's response token IDs or continue using the compatibility fallback. ## Source precedence 1. A complete token bundle on `choice.message`. 2. Top-level `prompt_token_ids`, choice-level `token_ids`, and `choice.logprobs` from the vLLM response. 3. Generation metadata from `choice.logprobs` and prompt IDs from `/tokenize`. `request_prompt_and_generation_token_ids: true` enables the second source by sending `return_token_ids=true` to compatible vLLM endpoints. The option defaults to `false`. If both inline sources are present, their token IDs must agree. The selected training message is validated once to avoid repeated list scans and copies at long sequence lengths. The fallback forwards `model`, `messages`, `tools`, `chat_template_kwargs`, `mm_processor_kwargs`, and `required_prefix_token_ids`. This keeps fallback tokenization aligned with the generation prompt. `prompt_logprobs` is not requested. vLLM returns prompt IDs independently when `return_token_ids` is enabled, and Gym does not consume prompt-token log probabilities. ## Related work - Gym [#700](#700) is superseded by this change. - Gym [#2576](#2576) is superseded by this change. - NeMo RL [#3390](NVIDIA-NeMo/RL#3390) produces the message-level token bundle consumed by the first source. - NeMo RL [#3581](NVIDIA-NeMo/RL#3581) reuses replayed token metadata across turns. - Gym [#1784](#1784) can add the Dynamo `nvext.engine_data` source independently. - Gym [#2324](#2324) can retain its video-specific schema, conversion, and adapter behavior independently. --------- Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Contributor
|
closing out in favor of #1558 |
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
3 participants
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.
What does this PR do?
Summary
Use prompt token IDs returned by vLLM to avoid a redundant
/tokenizerequest after generation.If
prompt_token_idsare missing orNone, Gym falls back to/tokenizewhile preservingmodel,messages,tools, andchat_template_kwargs.Testing
Noneprompt token IDs.Checklist
pre-commit run --all-files) (so CI lint/format/copyright pass).git commit -s) (so the DCO check passes).