[None][fix] Prevent special tokens leaking into detokenized output on slow tokenizers#16331
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #58990 [ run ] triggered by Bot. Commit: |
|
PR_Github #58990 [ run ] completed with state
|
4068de7 to
fcfa02f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59639 [ run ] triggered by Bot. Commit: |
|
PR_Github #59639 [ run ] completed with state
|
fcfa02f to
00925b2
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59912 [ run ] triggered by Bot. Commit: |
|
PR_Github #59912 [ run ] completed with state
|
00925b2 to
47049c2
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60652 [ run ] triggered by Bot. Commit: |
|
PR_Github #60652 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60943 [ run ] triggered by Bot. Commit: |
|
PR_Github #60943 [ run ] completed with state |
|
/bot run --disable-fail-fast --only-multi-gpu-test |
|
PR_Github #61152 [ run ] triggered by Bot. Commit: |
|
PR_Github #61152 [ run ] completed with state
|
Walkthrough
ChangesTokenizer special-token caching
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/tokenizer/tokenizer.py`:
- Around line 344-376: The fingerprint in _special_token_ids currently uses
len(inner._added_tokens_decoder) as a proxy for additional_special_tokens
changes, which can remain unchanged after reassignment. Replace that proxy with
inner.additional_special_tokens_ids in the fingerprint while retaining the named
special-token IDs and existing fallback behavior.
In `@tests/unittest/llmapi/test_tokenizer_decode.py`:
- Around line 74-89: Reorder the test around tokenizer.convert_ids_to_tokens so
the cache is primed before registering tok42, then add tok42 and perform the
second conversion to verify invalidation and filtering of both special tokens.
Remove the duplicate no-op add_special_tokens call, and update the
setup/comments to reflect the genuine pre-mutation and post-mutation sequence.
Do not add the optional removal/reassignment regression test unless the
underlying fingerprint fix is part of this change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 97aec9da-0e25-4138-a30a-ba66971494e6
📒 Files selected for processing (2)
tensorrt_llm/tokenizer/tokenizer.pytests/unittest/llmapi/test_tokenizer_decode.py
47049c2 to
d870ca4
Compare
8af372e to
8f481f6
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61225 [ run ] triggered by Bot. Commit: |
longlee0622
left a comment
longlee0622
left a comment
There was a problem hiding this comment.
fine to merge. If we care about slow-tokenizer + large-special-token-set models under streaming, it'd be worth a quick before/after profile of detok host overhead on such a model to confirm the added per-token cost is acceptable.
|
Is there a real-world use case where special tokens are modified after detokenization has already started?
|
|
PR_Github #61225 [ run ] completed with state
|
|
Thank you both for the suggestions. After verifying the issue, I confirmed that this bug is deterministic: the cache is generated before Gemma4's mutation step, which can cause |
8f481f6 to
12267f9
Compare
TransformersTokenizer snapshotted the tokenizer's special tokens in __init__ (all_special_ids for slow convert_ids_to_tokens, all_special_tokens for convert_tokens_to_string) and reused them forever. That snapshot is taken too early: a model's input processor can register special tokens after the wrapper is built but before serving (e.g. gemma4 adds <|video|>), so the affected tokens stopped being skipped and leaked into detokenized output when skip_special_tokens=True -- an accuracy regression. Compute the (ids, strings) sets once, lazily, on the first detokenization call. By request time all build-time special-token registration is done and nothing mutates it further, so a first-use cache is correct; it also keeps the slow-tokenizer detok path O(1) per streamed token (avoiding both the original per-id all_special_ids rebuild and any per-call refresh). Add regression tests on a genuine slow tokenizer: special tokens registered after wrapping must be reflected by both convert_ids_to_tokens and convert_tokens_to_string. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
12267f9 to
ebfc905
Compare
|
/bot run --disable-fail-fast |
Shixiaowei02
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/tokenizer/tokenizer.py (1)
200-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the repository’s Python typing requirements across the changed code.
tensorrt_llm/tokenizer/tokenizer.py#L200-L201: annotate the lazy cache astuple[set[int], set[str]] | None.tensorrt_llm/tokenizer/tokenizer.py#L334-L349: usetuple[set[int], set[str]]for_special_tokens().tests/unittest/llmapi/test_tokenizer_decode.py#L25-L54: annotate all new tokenizer methods and their return values.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/tokenizer/tokenizer.py` around lines 200 - 201, The lazy cache at tensorrt_llm/tokenizer/tokenizer.py lines 200-201 must be annotated as tuple[set[int], set[str]] | None; update _special_tokens() at lines 334-349 to use tuple[set[int], set[str]]; and annotate all newly added tokenizer methods and return values in tests/unittest/llmapi/test_tokenizer_decode.py lines 25-54 according to the repository’s Python typing requirements.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/llmapi/test_tokenizer_decode.py`:
- Around line 56-94: Add tests/unittest/llmapi/test_tokenizer_decode.py to an
appropriate llmapi unit-test list under tests/integration/test_lists/test-db/ or
tests/integration/test_lists/qa/, using the existing list-file conventions so
both newly added tests are included in the test suite.
---
Nitpick comments:
In `@tensorrt_llm/tokenizer/tokenizer.py`:
- Around line 200-201: The lazy cache at tensorrt_llm/tokenizer/tokenizer.py
lines 200-201 must be annotated as tuple[set[int], set[str]] | None; update
_special_tokens() at lines 334-349 to use tuple[set[int], set[str]]; and
annotate all newly added tokenizer methods and return values in
tests/unittest/llmapi/test_tokenizer_decode.py lines 25-54 according to the
repository’s Python typing requirements.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cd4de593-eac4-4fce-a50d-5d3dc9333c44
📒 Files selected for processing (2)
tensorrt_llm/tokenizer/tokenizer.pytests/unittest/llmapi/test_tokenizer_decode.py
| def test_convert_ids_to_tokens_reflects_special_tokens_registered_after_wrap() -> None: | ||
| # Specials registered after wrapping (before first detok) must be reflected. | ||
| inner = _ToySlowTokenizer() | ||
| inner.add_special_tokens( | ||
| { | ||
| "bos_token": "tok1", | ||
| "eos_token": "tok2", | ||
| "additional_special_tokens": ["tok5"], | ||
| } | ||
| ) | ||
| tokenizer = TransformersTokenizer(inner) | ||
| assert inner.is_fast is False | ||
|
|
||
| inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) | ||
|
|
||
| ids = [10, 42, 11, 5, 12] | ||
| got = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) | ||
| assert got == inner.convert_ids_to_tokens(ids, skip_special_tokens=True) | ||
| assert got == ["tok10", "tok11", "tok12"] | ||
|
|
||
|
|
||
| def test_convert_tokens_to_string_reflects_special_tokens_registered_after_wrap() -> None: | ||
| # The same must hold for convert_tokens_to_string's skip path. | ||
| inner = _ToySlowTokenizer() | ||
| inner.add_special_tokens( | ||
| { | ||
| "bos_token": "tok1", | ||
| "eos_token": "tok2", | ||
| "additional_special_tokens": ["tok5"], | ||
| } | ||
| ) | ||
| tokenizer = TransformersTokenizer(inner) | ||
| assert inner.is_fast is False | ||
|
|
||
| inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) | ||
|
|
||
| tokens = ["tok10", "tok42", "tok11"] | ||
| got = tokenizer.convert_tokens_to_string(tokens, skip_special_tokens=True) | ||
| assert got == "tok10 tok11" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
test_file='tests/unittest/llmapi/test_tokenizer_decode.py'
rg -n --fixed-strings "$test_file" \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa || true
pytest tests/unittest/ \
-k 'test_convert_ids_to_tokens_reflects_special_tokens_registered_after_wrap or test_convert_tokens_to_string_reflects_special_tokens_registered_after_wrap'Repository: NVIDIA/TensorRT-LLM
Length of output: 202
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## tracked test-related files\n'
git ls-files 'tests/unittest/llmapi/test_tokenizer_decode.py' \
'tests/integration/test_lists/test-db/*' \
'tests/integration/test_lists/qa/*' | sed 's#^`#-` #'
printf '\n## locate references to the unittest file in test lists\n'
rg -n --fixed-strings 'tests/unittest/llmapi/test_tokenizer_decode.py' tests/integration/test_lists || true
printf '\n## inspect the unittest file around the added tests\n'
cat -n tests/unittest/llmapi/test_tokenizer_decode.py | sed -n '1,180p'Repository: NVIDIA/TensorRT-LLM
Length of output: 9995
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## read list-file docs\n'
sed -n '1,220p' tests/integration/test_lists/test-db/README.md
printf '\n---\n'
sed -n '1,220p' tests/integration/test_lists/qa/README.md
printf '\n## any list entries referencing unittest paths\n'
rg -n 'tests/unittest/|test_tokenizer_decode|llmapi' tests/integration/test_lists || true
printf '\n## sample list file structure\n'
sed -n '1,120p' tests/integration/test_lists/test-db/l0_sanity_check.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 38410
Test coverage summary — insufficient.
- Added tests:
test_convert_ids_to_tokens_reflects_special_tokens_registered_after_wrap,test_convert_tokens_to_string_reflects_special_tokens_registered_after_wrap. - Neither test is listed under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/; addtests/unittest/llmapi/test_tokenizer_decode.pyto the relevant list file (for example, one of the llmapi unit-testtests/integration/test_lists/test-db/l0_*.ymllists).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/llmapi/test_tokenizer_decode.py` around lines 56 - 94, Add
tests/unittest/llmapi/test_tokenizer_decode.py to an appropriate llmapi
unit-test list under tests/integration/test_lists/test-db/ or
tests/integration/test_lists/qa/, using the existing list-file conventions so
both newly added tests are included in the test suite.
Source: Path instructions
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/tokenizer/tokenizer.py (1)
200-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse precise Python 3.10 type annotations for the lazy cache.
The new cache and helper return type currently lose the distinction between integer IDs and string tokens. Use
tuple[set[int], set[str]] | Nonefor the field andtuple[set[int], set[str]]for_special_tokens().As per coding guidelines, use precise annotations, built-in generics, and
|types.Proposed typing update
- self._special_tokens_sets = None + self._special_tokens_sets: tuple[set[int], set[str]] | None = None - def _special_tokens(self) -> Tuple[set, set]: + def _special_tokens(self) -> tuple[set[int], set[str]]:Also applies to: 334-349
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/tokenizer/tokenizer.py` around lines 200 - 201, Update the _special_tokens_sets field to use the precise Python 3.10 annotation tuple[set[int], set[str]] | None, and annotate _special_tokens() to return tuple[set[int], set[str]]. Preserve the lazy cache behavior while distinguishing integer IDs from string tokens.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/tokenizer/tokenizer.py`:
- Around line 200-201: Update the _special_tokens_sets field to use the precise
Python 3.10 annotation tuple[set[int], set[str]] | None, and annotate
_special_tokens() to return tuple[set[int], set[str]]. Preserve the lazy cache
behavior while distinguishing integer IDs from string tokens.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5cac293a-ebc9-4997-82f7-d9be38b91d09
📒 Files selected for processing (2)
tensorrt_llm/tokenizer/tokenizer.pytests/unittest/llmapi/test_tokenizer_decode.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unittest/llmapi/test_tokenizer_decode.py
4 participants
Description
This pull request refactors how special tokens are handled in the
TransformersTokenizerclass to ensure that any special tokens registered after the tokenizer is wrapped are correctly recognized during detokenization. It also adds new tests to verify this behavior, especially for slow tokenizers. The changes improve correctness and efficiency, particularly when handling streamed detokenization and special tokens.Special Token Handling Improvements
TransformersTokenizerto be lazy and cached on first detokenization, instead of computed at initialization. This ensures that special tokens registered after wrapping (but before serving) are recognized, and improves performance by avoiding repeated computation. (tensorrt_llm/tokenizer/tokenizer.py) [1] [2]convert_ids_to_tokensandconvert_tokens_to_stringmethods to use the new cached special-token sets, ensuring correctness when skipping special tokens. (tensorrt_llm/tokenizer/tokenizer.py) [1] [2]Testing Enhancements
_ToySlowTokenizer) and two new tests to verify that special tokens registered after wrapping are properly recognized and skipped during detokenization and string conversion. (tests/unittest/llmapi/test_tokenizer_decode.py)These changes ensure that the tokenizer wrapper behaves correctly with dynamic special token registration and that performance is maintained for streamed detokenization.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
TransformersTokenizerspecial-token handling to avoid stale, init-time snapshots when special tokens are added/changed after the wrapper is constructed. Removed eager cache construction and replaced it with a lazy, first-use cache (self._special_tokens_sets) computed frominner.all_special_ids/inner.all_special_tokens.convert_ids_to_tokens(ids, skip_special_tokens=True)now uses a slow path only whenidsis a list and the wrapped tokenizer is not fast; it filters by the memoized special-id set from_special_tokens()[0].convert_tokens_to_string(..., skip_special_tokens=True)uses the memoized special-token strings set from_special_tokens()[1]when the wrapper’s custom detokenization path is active.inner.all_special_ids/inner.all_special_tokensare unavailable, catchingAttributeErrorandNotImplementedError.QA Engineer Review
tests/unittest/llmapi/test_tokenizer_decode.py_ToySlowTokenizer(PreTrainedTokenizer)to forceTransformersTokenizerthrough the slow/non-fast filtering logic deterministically.test_convert_ids_to_tokens_reflects_special_tokens_registered_after_wraptest_convert_tokens_to_string_reflects_special_tokens_registered_after_wraptests/integration/test_lists/.