[https://nvbugs/6163690][fix] Use PreTrainedTokenizerFast in trtllm-bench prepare_dataset to avoid NemotronH config parsing error

[pamelap-nvidia]

@coderabbitai summary

Description

Problem

trtllm-bench prepare-dataset (and the legacy benchmarks/cpp/prepare_dataset.py) fails with KeyError: '-' when the model is NVIDIA-Nemotron-Nano-9B-v2-NVFP4:

File ".../transformers/models/nemotron_h/configuration_nemotron_h.py", line 271, in _pattern_to_list
    return [pattern_mapping[char] for char in pattern]
KeyError: '-'

Root cause

AutoTokenizer.from_pretrained calls AutoConfig.from_pretrained internally to resolve the tokenizer class. This instantiates the built-in NemotronHConfig (added in transformers 5.5.3), whose _pattern_to_list maps M→mamba, E→moe, *→attention but is missing '-': 'mlp'. The NVFP4 model's config.json has hybrid_override_pattern: M-M-M-MM-M-M-M*-... where - encodes MLP layers, causing the KeyError.

This was not hit before because:

• rc13 shipped transformers 4.57.3 where nemotron_h was not yet a built-in type
• rc15's test ran with a pre-generated --dataset file, so tokenizer loading was never triggered

The upstream fix is in huggingface/transformers PR #44763 (merged 2026-04-22, first in transformers v5.10.1). TRTLLM currently ships transformers 5.5.4 and cannot upgrade to 5.10.1 without accuracy regressions in GPT-OSS tests (chat template changes, torch_dtype deprecation).

Fix

prepare_dataset.py only uses the tokenizer to encode text for synthetic dataset generation — it never needs model config. PreTrainedTokenizerFast.from_pretrained loads directly from tokenizer.json without invoking AutoConfig, bypassing the broken config parsing entirely.

The NVFP4 model ships a complete tokenizer.json, so the fast tokenizer loads correctly (131072-vocab, verified on rc19). An AutoTokenizer fallback is kept for models that do not provide a fast tokenizer.

Testing

• Verified PreTrainedTokenizerFast.from_pretrained succeeds on rc19 with NVIDIA-Nemotron-Nano-9B-v2-NVFP4 (131072-vocab, tokenizes correctly)
• AutoTokenizer fallback path is exercised by any model without tokenizer.json
• No change to the throughput command path (uses pre-generated dataset, never calls prepare_dataset)

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Tokenizer loading

Layer / File(s)	Summary
Fast tokenizer validation and fallback tensorrt_llm/bench/dataset/prepare_dataset.py	validate_tokenizer() first loads PreTrainedTokenizerFast from tokenizer.json, then falls back to AutoTokenizer while retaining padding and error-handling behavior.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#9250: Introduces the overlapping RootArgs.validate_tokenizer() flow in the same module.

Suggested reviewers: wanli-jiang

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title follows the required ticket/type format and clearly summarizes the tokenizer-loading fix.
Description check	✅ Passed	The description includes the required problem, root cause, fix, testing, and checklist sections with meaningful details.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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/bench/dataset/prepare_dataset.py`:
- Around line 42-49: Narrow the exception handler around
PreTrainedTokenizerFast.from_pretrained in the tokenizer-loading flow to catch
only OSError and ValueError, allowing unrelated loader errors to propagate while
preserving the AutoTokenizer fallback for intended fast-tokenizer load failures.

🪄 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: 0c84cc16-35ca-4270-b613-43b16afd8685

📥 Commits

Reviewing files that changed from the base of the PR and between 3c07ada and 2916bb8.

📒 Files selected for processing (1)

• tensorrt_llm/bench/dataset/prepare_dataset.py

[pamelap-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #61461 [ run ] triggered by Bot. Commit: 13080b7 Link to invocation

[tensorrt-cicd]

PR_Github #61461 [ run ] completed with state FAILURE. Commit: 13080b7
/LLM/main/L0_MergeRequest_PR pipeline #49682 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[pamelap-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61480 [ run ] triggered by Bot. Commit: 13080b7 Link to invocation

[tensorrt-cicd]

PR_Github #61480 [ run ] completed with state SUCCESS. Commit: 13080b7
/LLM/main/L0_MergeRequest_PR pipeline #49699 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[pamelap-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #61584 [ run ] triggered by Bot. Commit: 13080b7 Link to invocation

[tensorrt-cicd]

PR_Github #61584 [ run ] completed with state FAILURE. Commit: 13080b7
/LLM/main/L0_MergeRequest_PR pipeline #49795 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation