fix: Add chat completion parity tests for TRTLLM v.s. vLLM and fix disparities - #3422
fix: Add chat completion parity tests for TRTLLM v.s. vLLM and fix disparities#3422hchings wants to merge 11 commits into
Conversation
hchings
changed the title
hchings
changed the title
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Reviewed at 5aa8e0a with an agent team (7 agents, two adversarial passes). Nice work — the parity harness is the right idea, and underneath the tests this fixes three real production bugs: the swallowed <think> tag, the missing prompt_token_ids on vLLM chat responses, and sampled non-EOS stop tokens being dropped during prefix splicing.
CONFLICTING with main in all three vLLM files. #3390 merged on 2026-08-10 and implements the prompt_token_ids work as a superset, so the vLLM half of this PR should mostly be dropped on rebase rather than merged. #3390 also added a top_logprobs precondition that will make the new vLLM parity test return HTTP 400 — details inline.
Three things block green CI:
test_tool_parser_contractscannot pass — thetext_before_and_after_tool_callexpectation is unreachable on TRT-LLM (two independent.strip()calls).- 2 × ruff
I001in the new parity test files. - After rebase: the vLLM parity requests need
top_logprobs.
Everything else is a suggestion. A few notes on how this was checked, so you know what to trust: I had no GPU, so all vLLM/TRT-LLM runtime claims are static analysis against pinned upstream source — but the two blockers above are verifiable by inspection and need no GPU. Numbers quoted from the golden fixture were extracted from the committed file, and Qwen3 token IDs were verified against Qwen/Qwen3-0.6B@c1899de.
If it helps, I wrote up the four disparities with a worked example built from real tokens in your golden (turn 1 → turn 2 splice, showing exactly which token gets skipped): https://terryk.gitlab-master-pages.nvidia.com/nemo-html/terryk/claude-summaries/pr-3422-trtllm-vllm-parity.html
Generated by Claude Code
| # Return the engine-recorded prompt after prefix replacement. | ||
| prompt_token_ids = list(final_res.prompt_token_ids or []) | ||
| for choice in response.choices: | ||
| choice.message.prompt_token_ids = prompt_token_ids |
There was a problem hiding this comment.
1 action item (rebase), 1 follow-up.
AI-1 — delete this hunk on rebase
The PR is CONFLICTING with main. #3390 "perf(vllm): return token metadata with chat completions" merged 2026-08-10 and supersedes this hunk specifically, as a superset — it attaches prompt_token_ids, generation_token_ids and generation_log_probs, and generalises the dump helper into a field loop. Its motivation was performance: it removes an awaited /tokenize round-trip Gym was making every turn (~1.17x end-to-end).
To be precise about what survives, since this PR touches four regions of the file:
| hunk | what | after rebase |
|---|---|---|
@@ -323 |
parse_chat_output |
survives |
@@ -393 |
compute model_stop_token_ids |
survives |
@@ -561 |
pass it to replace_prefix_tokens |
survives |
@@ -613 (this one) |
prompt_token_ids + control flow |
superseded — delete |
Only this hunk overlaps #3390 (which touched old lines 45, 594-628, 738), so it is the conflict. The utils.py and test_vllm_utils.py changes go with it. main also moved vLLM 0.20.0 -> 0.25.1 in the same window, so re-check anything validated against 0.20.0.
Also note #3390 gates the attach on request.logprobs and return_as_token_id and adds a top_logprobs precondition — which breaks the new parity test (see the comment on test_vllm_http_server_parity.py:149).
Follow-up — worth a sentence in the PR description
After rebase, the only change this PR still makes to vLLM runtime behaviour is hunks @@ -393 / @@ -561: vLLM now counts cfg["stop_token_ids"] and generation_config.eos_token_id as turn terminators, where before only tokenizer.eos_token_id counted. That is needed to keep both backends symmetric through the shared util — but it is a behaviour change to the backend that wasn't the subject of this PR, in a case that (as far as I can tell) nothing in-repo reaches. Calling that out explicitly would save the next reviewer the same trace.
There was a problem hiding this comment.
Re-opening for the follow-up only — AI-1 did land.
AI-1: the @@ -613 hunk is gone; git diff origin/main...HEAD on this file now returns three hunks and main's #3390 version of the attach path stands unmodified.
The follow-up did not: the PR description at 7671c1b still has no sentence about the widened stop-token set. vllm_worker_async.py:400 reads model_stop_token_ids = set(self.cfg.get("stop_token_ids") or ()) and folds in generation_config.eos_token_id at :401-407, where before only tokenizer.eos_token_id counted. The body's nearest line — "Verifies generated EOS/stop tokens are retained in response token IDs and logprobs while excluded from parsed response text" — describes a test assertion, not the vLLM runtime change.
| return golden | ||
|
|
||
|
|
||
| @pytest.fixture(scope="module", params=tuple(PARSER_SCENARIOS)) |
There was a problem hiding this comment.
test_vllm_http_server_parity.py:255
2 action items — both CI cost, neither blocking.
AI-1
pytest-shard assigns per item, not per module: sha256hash(item.nodeid) % num_shards. A module-scoped fixture can't hold its items together, so this file's 4 items (2 params × 2 tests) scatter across all three L0_Unit_Tests_Vllm_* shards — I hashed the real nodeids and got builds in all 3 shards (3-4 engine builds instead of 2). Folding test_parser_contracts into test_parity (one item per param) would cap it at 2.
AI-2
Neither parity file is listed in tests/unit/excluded_unit_tests.sh, whose stated principle is to exclude tests that "spin up WorkerGroups, Policy, VllmGeneration, ... load HF models" — which this file does. Scoped correctly this only affects the opt-in CI:Lfast lane (image_tag is non-empty only for Lfast/docs, and the PR default is L1 ⇒ FAST=0), but it's two lines to fix:
--ignore=unit/models/generation/test_vllm_http_server_parity.py
--ignore=unit/models/generation/trtllm/test_trtllm_http_server_parity.py
There was a problem hiding this comment.
Re-opening for AI-2 only — AI-1 did land.
AI-1: the two contract tests no longer touch the engine (_parse_with_vllm at test_vllm_http_server_parity.py:278 runs in-process now that parse_chat_output moved out of the worker), so only the two test_parity items build one. That is the cap this asked for, even though the suggested fold-into-test_parity was not done — the items still scatter across all three shards, they are just cheap now.
AI-2 did not: grep -n parity tests/unit/excluded_unit_tests.sh at 7671c1b returns nothing. The list is sourced only under FAST=1 (run_unit_shard_common.sh:28-29), so the CI:Lfast lane still collects both parity files and builds engines there.
bfaf824 to
ec02504
Compare
|
/ok to test ec02504 |
ec02504 to
c2ed2fc
Compare
|
/ok to test c2ed2fc |
hchings
added
the
CI:Lfast
c2ed2fc to
95bb5cf
Compare
|
/ok to test 95bb5cf |
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Re-reviewed at c2ed2fc after the rebase. All 11 findings from my previous review are genuinely fixed — I verified each against the code rather than the commit messages, since the SHAs in the replies predate the rebase. Several went further than asked: dropping parse_chat_output entirely rather than relocating it, and making the contract tests tokenizer-only so they no longer need a GPU engine. I also ran the new test_chat_template_parity_common.py locally — 19/19 pass. Flipping the golden so vLLM generates and TRT-LLM conforms is the right direction.
Two follow-ups below, both minor: one dead-code cleanup, and one re-open where the fix works mechanically but the test collapses into a duplicate.
Convergence is safe — I checked whether the content whitespace difference between backends (vLLM "before " vs TRT-LLM "before") could reach the training tokens. It cannot: replace_prefix_tokens keeps prior turns as verbatim token IDs and discards the template's re-render of that region, and final_prompt_token_ids is used both as the engine prompt and as the returned prompt_token_ids, so no train/inference skew is possible. The shared-splice refactor is also a strict no-op when model_stop_token_ids is omitted (0 differences over 108,569 randomised cases).
Not this PR, for the record: while tracing the above I found that top_k is handled inconsistently across backends — vLLM's HTTP path forces -1 and ignores cfg["top_k"], while TRT-LLM's HTTP path honours it since #3537, and the two families disagree on whether "disabled" is -1 or 0. That is all pre-existing on main and nothing here needs to change; I filed #3776 to document it.
One process note: the unit lanes still have not run (CI quality check: No tests run: Label CI:L*). Lint is green so the ruff fix is CI-verified, but the tool-parser contract fix deserves a real TRT-LLM lane run — that test was previously failing deterministically, and the fix has only been verified by inspection.
Generated by Claude Code
|
/ok to test aa903ec |
b7ffb77 to
673ef9a
Compare
|
/ok to test 673ef9a |
673ef9a to
ccb858d
Compare
|
/ok to test ccb858d |
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the NVIDIA-NeMo#3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…0rc21 Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
ccb858d to
7671c1b
Compare
|
/ok to test 7671c1b |
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
Scope: the three production fixes (stop-token threading in replace_prefix_tokens, TRT-LLM reasoning-parser state, stop-token retention), the new vLLM/TRT-LLM parity harness and its committed golden, and the new TRT-LLM gym nightly.
Checked:
- Splice semantics — hand-derived old vs new
replace_prefix_tokensacross the EOS / non-EOS-stop / max-tokens cases. The token the model side stops cutting and the one the template side now skips are the sameeos_token_id, so the EOS path is byte-identical to before; only the non-EOS-stop path changes, which is the fix. - Upstream API usage —
ReasoningParserFactory, both Qwen3 tool parsers andtry_get_generation_configread at the pinnedtensorrt-llm==1.3.0rc21and vLLM versions rather than from docstrings. - Golden provenance — the fixture is vLLM-generated (
source_backend: "vllm") and written only by the vLLM lane, so the cross-backend comparison lives in the TRT-LLM test while the vLLM one is a self-regression against the same file. - Coverage — each new test resolved against the shard that actually collects it, and the golden decoded turn by turn.
- Recipe layer — inheritance, node counts, suite bookkeeping, and the runnability of the new nightly end to end.
Remaining comments are on the new nightly (tool-parser choice, in_flight_weight_updates, an unread policy.optimizer block, and a request for curves), two production asks (the widened stop-token set and an aliased return), one coverage gap on the reasoning-state inference, the golden's size, two re-opened threads, and three stale-text nits.
| expose_http_server: true | ||
| default_chat_template_kwargs: | ||
| enable_thinking: false | ||
| tool_parser: qwen3_coder |
There was a problem hiding this comment.
qwen3_coder can't parse what Qwen3-0.6B emits, so tool calls in this nightly come back empty.
Qwen3's chat template emits a JSON body inside <tool_call> tags. Qwen3CoderToolParser expects <function=...> XML instead, but its has_tool_call only tests for the bare <tool_call> substring — so it enters parsing, _parse_block matches no <function=, and it returns zero calls. _make_parse_tool_calls then hits if not calls: return text, [] — the raw text survives as content, but the gym env sees no tool call, and there is no fallback parser.
Verified against the TRT-LLM commit this repo actually builds, not the PyPI release: 3rdparty/TensorRT-LLM-workspace/pyproject.toml [tool.trtllm] pins ref = bf2ef86f9a2652132b11773d4041e292c553c142.
tool_parser_factory.py:58-59— the YAML string routes through{"qwen3": Qwen3ToolParser, "qwen3_coder": Qwen3CoderToolParser}.tool_parser_factory.py:20-22— upstream's ownMODEL_TYPE_TO_TOOL_PARSERresolvesmodel_type: "qwen3"to"qwen3".Qwen/Qwen3-0.6B@c1899de2'sconfig.jsonismodel_type: "qwen3", soqwen3_coderis not what upstream would auto-pick for this model — it is not the target of any model_type.qwen3_coder_parser.py:57-58—has_tool_callisself.tool_call_start_token in textwith the token being the bare<tool_call>;:40-42—tool_call_function_regex = r"<function=(.*?)</function>|<function=(.*)$"is the only thing_parse_blocklooks for.qwen3_tool_parser.py:34— the JSON-format parser this model's output actually fits (bot_token = "<tool_call>\n").
In-repo, tool_parser always tracks the model's emitted format: the base config this recipe inherits uses nemotron_json for Nemotron Nano (grpo_workplace_assistant_nemotron_nano_v2_9b.yaml:257), and the other Qwen gym recipes use vLLM's JSON-format hermes. qwen3_coder's only other use in the repo is a Nemotron-Omni VLM recipe. This PR's own TRT-LLM parity test uses _build_tool_parser("qwen3") for the same model.
| tool_parser: qwen3_coder | |
| tool_parser: qwen3 |
There was a problem hiding this comment.
traced by claude. I'm not sure if change to qwen3 is correct. could you help double check?
| **(default_chat_template_kwargs or {}), | ||
| } | ||
|
|
||
| def _prompt_opens_reasoning(token_ids: list[int]) -> bool: |
There was a problem hiding this comment.
The True branch of _prompt_opens_reasoning has no coverage, and it's the branch a shipped recipe takes.
It's a closure inside create_app, so no test can import it — every other module-level helper in this file is imported and tested by test_trtllm_http_server.py. The CPU contract test passes reasoning_at_start straight into _build_reasoning_parser, and all three golden scenarios end their prompt in assistant\n or </think>, so the e2e path only ever reaches False. grpo_qwen3_30b_async_swe_trtllm.yaml pairs deepseek-r1 with Qwen3-30B-A3B-Thinking, whose template ends <|im_start|>assistant\n<think>\n.
Suggest fix:
- Hoist it to module scope taking
tokenizeras a parameter — that is the only name it captures. - Add a table test:
assistant\n-> False,<think>\n\n</think>\n\n-> False,...<think>\n-> True.
| gen_logprobs.pop() | ||
|
|
||
| gen_text = tokenizer.decode(gen_token_ids, skip_special_tokens=False) | ||
| # Parse without trailing stops; retain full response metadata. |
There was a problem hiding this comment.
nit: the comment on include_stop_str_in_output still describes the trimming this hunk removed.
trtllm_http_server.py:95 reads "so the adapter can trim tokens and logprobs together", but the lockstep pop() loop is gone and the flag is now load-bearing for the opposite reason — stop tokens must be retained so generation_token_ids matches vLLM and the next turn's replace_prefix_tokens can see model_ended_with_stop. _tokens_for_response_text is text-only and never touches logprobs.
Suggest fix: reword line 95 to "Include generated stop tokens so generation_token_ids stays contiguous with the next turn's prompt."
| count_seen += 1 | ||
| if count_seen == count_needed: | ||
| template_cut_start = pos | ||
| # Keep sampled stops; otherwise let the template add the missing EOS. |
There was a problem hiding this comment.
nit: the docstring still describes the pre-PR cut.
Lines 71-73 say the splice cuts "at the N-th EOS in template_token_ids", but this line cuts one past it whenever the model's own last token was a stop. The worked example below still evaluates correctly, so this is one stale sentence rather than a stale example — and model_stop_token_ids is not named in the prose at all.
Suggest fix: reword to "...cut at the N-th EOS in template_token_ids, or just past it when the model's own last token was a stop token", naming model_stop_token_ids as the set that decides.
| suffix_mismatches.append( | ||
| f"scenario={scenario_name!r}: transition {turn_index + 1}->" | ||
| f"{turn_index + 2} appended prompt suffix mismatch " | ||
| f"(vLLM={actual_suffix!r}, TRT-LLM={expected_suffix!r})" |
There was a problem hiding this comment.
nit: the expected side here is the vLLM golden, not TRT-LLM.
_load_golden asserts source_backend == "vllm" and this test is the only writer of the fixture, so a failure here sends the reader hunting in the wrong backend.
| f"(vLLM={actual_suffix!r}, TRT-LLM={expected_suffix!r})" | |
| f"(actual={actual_suffix!r}, golden={expected_suffix!r})" |
| uv run tests/check_metrics.py "$JSON_METRICS" \ | ||
| 'mean(data["train/reward"]) > 0.05' \ | ||
| 'median(data["train/token_mult_prob_error"]) < 1.1' \ | ||
| 'mean(data["train/gen_kl_error"]) < 0.02' \ | ||
| 'mean(data["train/grad_norm"], 2, 0) > 0.1' \ | ||
| 'mean(data["train/grad_norm"], 2, 0) < 30.0' |
There was a problem hiding this comment.
could you paste the wandb curves or metric results into the PR description to show that the new nightly test can run well?
| async_grpo: | ||
| in_flight_weight_updates: true |
There was a problem hiding this comment.
in_flight_weight_updates: true has no async loop to apply to — this recipe runs synchronous GRPO.
enabled isn't set here, and the base config has async_grpo.enabled: false, so the dispatch at run_grpo_nemo_gym.py:289 (elif config.grpo.async_grpo.enabled:) falls through to the synchronous branch at :334 and async_grpo_train never runs.
It isn't fully inert, though: line 54 interpolates this value into trtllm_cfg.in_flight_weight_updates, which trtllm_generation.py:474-479 turns into drain=not in_flight — so refit skips the drain under a sync loop.
Dropping the block is the smaller fix, and it matches the recipe's own name (no async); the base already supplies in_flight_weight_updates: false. If async was the intent instead, add enabled: true alongside it, the way grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml:3-5 does.
| async_grpo: | |
| in_flight_weight_updates: true |
| text_token_end = len(token_ids) | ||
| while text_token_end and token_ids[text_token_end - 1] in stop_token_ids: | ||
| text_token_end -= 1 | ||
| return token_ids if text_token_end == len(token_ids) else token_ids[:text_token_end] |
There was a problem hiding this comment.
The identity return aliases the caller's list, and gen_token_ids is what ships back as training data.
Both branches are behaviourally identical — the only consumer is tokenizer.decode at :302 — so the conditional buys one skipped copy of at most max_new_tokens ints and costs an alias: if anything downstream ever trims text_token_ids, it silently mutates the generation_token_ids sent at :331, which is the class of bug this hunk exists to prevent. test_response_text_excludes_stops_without_mutating_training_tokens asserts exactly that invariant, while test_response_text_reuses_unterminated_generation_tokens pins the alias with is and so blocks the simpler form.
Separately, text_token_end reaching 0 is reachable — a generation that is entirely stop tokens, e.g. [151645] — and neither new test covers it. It doesn't crash (decode([]) is ""), but it is the boundary the while text_token_end and ... guard exists for, so it is worth pinning.
Suggested companion edit, replacing test_trtllm_http_server.py:110-115:
def test_response_text_returns_all_tokens_when_unterminated():
generation_token_ids = [10, 20]
assert _tokens_for_response_text(generation_token_ids, {2, 3}) == [10, 20]
def test_response_text_is_empty_when_generation_is_all_stops():
assert _tokens_for_response_text([2, 3], {2, 3}) == []| return token_ids if text_token_end == len(token_ids) else token_ids[:text_token_end] | |
| return token_ids[:text_token_end] |
| model_stop_token_ids = set(self.cfg.get("stop_token_ids") or ()) | ||
| generation_eos_token_ids = model_config.try_get_generation_config().get( | ||
| "eos_token_id" | ||
| ) | ||
| if isinstance(generation_eos_token_ids, int): | ||
| model_stop_token_ids.add(generation_eos_token_ids) | ||
| elif generation_eos_token_ids is not None: | ||
| model_stop_token_ids.update(generation_eos_token_ids) |
There was a problem hiding this comment.
The generation-config eos should be a fallback, not a union — as written it silently widens an explicitly configured stop_token_ids.
configure_generation_config only fills this in when it is None (__init__.py:75-76), so a value that reaches the worker is either the user's explicit list or the auto-filled [tokenizer.eos_token_id] — in both cases it is the configured answer to "what terminates a turn", and unioning generation_config.json's eos list into it overrides that. A recipe pinning [151645] precisely to avoid terminating on <|endoftext|> would still get {151645, 151643}, model_ended_with_stop would fire on the excluded token, and the next turn's prompt would lose the template <|im_end|>.
Today the gym path nulls stop_token_ids outright (nemo_gym.py:952-954), so the fallback is what runs and behaviour is unchanged either way — this is about which precedence the code encodes for when that changes.
| model_stop_token_ids = set(self.cfg.get("stop_token_ids") or ()) | |
| generation_eos_token_ids = model_config.try_get_generation_config().get( | |
| "eos_token_id" | |
| ) | |
| if isinstance(generation_eos_token_ids, int): | |
| model_stop_token_ids.add(generation_eos_token_ids) | |
| elif generation_eos_token_ids is not None: | |
| model_stop_token_ids.update(generation_eos_token_ids) | |
| model_stop_token_ids = set(self.cfg.get("stop_token_ids") or ()) | |
| if not model_stop_token_ids: | |
| generation_eos_token_ids = model_config.try_get_generation_config().get( | |
| "eos_token_id" | |
| ) | |
| if isinstance(generation_eos_token_ids, int): | |
| generation_eos_token_ids = [generation_eos_token_ids] | |
| model_stop_token_ids = set(generation_eos_token_ids or ()) |
Note the isinstance normalisation still has to stay: eos_token_id is an int for many models and a list for Qwen3, and replace_prefix_tokens calls set() on whatever it receives (openai_server_utils.py:100), which would raise on a bare int.
| stop_token_ids: | ||
| - 151643 | ||
| - 151645 |
There was a problem hiding this comment.
Heads-up, pre-existing and no action needed in this PR: these lines cannot take effect on a gym run.
setup_nemo_gym_config clears both stop settings unconditionally, for every backend, after its backend branch:
# Stop strings or token ids are not supported
generation_config["stop_strings"] = None
generation_config["stop_token_ids"] = NoneIt is the only writer of that value, all three gym entrypoints go through it (run_grpo_nemo_gym.py:193, run_distillation_nemo_gym.py:107, run_grpo_single_controller.py:145), and nothing repopulates it afterwards — configure_generation_config runs earlier, at run_grpo_nemo_gym.py:180. So under gym the worker always sees None, whatever the recipe asks for.
The values are still covered by the server-side fallbacks: trtllm_http_server.py:154 adds tokenizer.eos_token_id (151645) and :161 adds the generation_config.json eos list, which for Qwen3-0.6B is [151645, 151643] — both of the ids here. So nothing breaks; the three lines just don't do anything. Worth deciding whether to keep them as documentation of intent or drop them, and worth knowing when reading the recipe.
4 participants
What does this PR do ?
Follow-up task for #3130. Context in #3342 (comment)
1. Unit tests for pre-merge -- focus on parity of prompt construction, tool parsing, reasoning parsing, and EOS-stripping.
preserving multi-turn token continuity.
2. Nightly -- e2e run comparing reward/logprob error/grad_norm, etc for SWE task.
WIP
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information