fix(trtllm): report response truncation so overlong filtering can see it - #4042
Open
yupengtang wants to merge 2 commits into
Open
yupengtang wants to merge 2 commits into
yupengtang wants to merge 2 commits into
Conversation
`stop_strings` is part of the shared generation contract -- it is declared on `GenerationConfig` and again per sample on `GenerationDatumSpec` -- and vLLM, SGLang and Dynamo all read it. TRT-LLM built its `SamplingParams` with `stop_token_ids` and never `stop`, on the direct path and over HTTP, so the same config stopped generation on the other three backends and ran on to `max_new_tokens` here. The two controls are not interchangeable: `stop_token_ids` cannot express a multi-token boundary like `</answer>`, which is exactly the shape a chat or agentic rollout stops on. Nothing errors, so a run only shows it as rollouts that overrun their boundary. `TrtSamplingParams` already accepts `stop` alongside `stop_token_ids`, so both paths now pass it. The direct path merges the configured list with the per-sample ones the way `BaseVllmGenerationWorker._merge_stop_strings` does -- one `SamplingParams` is built per batch, so a sample's stop strings apply to the batch, matching vLLM rather than inventing a different rule. An empty result stays `None` so TRT-LLM keeps its own default instead of receiving an empty list. The HTTP helper gains the argument next to `stop_token_ids`, mirroring how NVIDIA-NeMo#3537 threaded the sampling config through after the same drift was found in `top_k`. Signed-off-by: Yupeng Tang <85978465+yupengtang@users.noreply.github.com>
`truncated` is the column rollouts use to mark a sample whose response hit `max_new_tokens` without a stop token. vLLM, SGLang and Dynamo all return it; TRT-LLM returned no such column, and `run_multi_turn_rollout` starts from an all-False tensor and only overwrites it when the backend supplied one. So on TRT-LLM every sample looked untruncated. That is not only a metric: `grpo.overlong_filtering` reads this column to zero the loss multiplier for samples that ran out of budget. With the column stuck at False those samples train as if they had finished on their own, which is the opposite of what the flag asks for, and nothing reports it. `truncation_rate` under-reports for the same reason. TRT-LLM already says this on each output -- `finish_reason == "length"` is its spelling for hitting the token cap without a stop token, the same signal `BaseVllmGenerationWorker` reads. Fill the column from it, and give the empty-batch shortcut the matching empty tensor so both branches return the same shape. Observation truncation, set in `rollouts.py` when an env observation is cut to fit the context, was never affected -- only the response side was missing. Scoped to TRT-LLM. Megatron has the same gap but reports no finish reason; its mcore path would need truncation inferred from the termination id, which is a judgement call worth making separately. Signed-off-by: Yupeng Tang <85978465+yupengtang@users.noreply.github.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
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
2 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 ?
Makes TRT-LLM report response truncation, so
grpo.overlong_filteringstops silently doing nothing on that backend.truncatedmarks a sample whose response hitmax_new_tokenswithout a stop token. vLLM, SGLang and Dynamo all return it inGenerationOutputSpec; TRT-LLM returned no such column.That is not simply a missing field.
run_multi_turn_rolloutstarts from an all-False tensor and only overwrites it when the backend supplied one:So the column exists and reads all-False on TRT-LLM.
grpo.py:3374then does:which zeroes nothing. Samples that ran out of budget train as if they had finished on their own, the opposite of what the flag asks for, with no error and no warning.
truncation_rateunder-reports for the same reason.Observation truncation (an env observation cut to fit the context,
rollouts.py:1116) was never affected, since it is computed rollout-side. Only the response side was missing.Issues
No issue filed for this. Distinct from #2163 / #2395, which are about
truncatedbeing absent and raisingKeyError; here the column is present and simply always False.Usage
No config or API change.
overlong_filtering: truenow behaves on TRT-LLM the way it already does on the other backends.Design & Code Changes
TRT-LLM already reports this per output:
finish_reason == "length"is its spelling for hitting the token cap without a stop token (executor/result.py:finish_reason: Optional[Literal['stop', 'length', 'timeout', ...]]). That is the same signalBaseVllmGenerationWorkerreads, so this is a plumbing fix rather than new inference.generate_asynccollectsgen.finish_reason == "length"per output and returns it astruncated.torch.zeros(0, dtype=torch.bool), so both branches return the same columns.Scoped to TRT-LLM. Megatron has the same gap but reports no finish reason, and its mcore path would need truncation inferred from
termination_idagainst the token cap. That is a judgement call about mcore semantics rather than a port of an existing signal, so I left it out rather than guess; happy to follow up if you want it in the same shape.Before your PR is "Ready for review"
Pre checks:
Additional Information
Tests
Two added in
test_trtllm_worker_async.py, plus the existingtest_generate_async_converts_padded_batch_and_logprobsextended so its fake outputs now carryfinish_reason, which the real TRT-LLM object always has, and it asserts the new column.test_generate_async_reports_truncation_from_finish_reason:"stop"and"length"in one batch produce[False, True]test_generate_async_empty_batch_returns_empty_truncated: the shortcut returns an empty bool tensor rather than omitting the columnReverting only
trtllm_worker_async.pytomain, the first fails on the missing column rather than on a signature change:With the change (both commits on this branch):
No failures on either version. The 5 skips are local: my stub
tensorrt_llmhas notensorrt_llm.servesubmodule.CPU only. I have no GPU box, so the tests needing a real TRT-LLM runtime did not run and I could not watch a rollout actually get filtered. The change is confined to how the output batch is assembled, which the tests above cover.
ruff 0.9.9check, import sort and format are clean on all four files.