Qwen2-VL SFT with extended vocab generates only vocab_size-1 (168063) tokens in TensorRT-LLM 0.17, while HF works fine

[lhh0916]

System Info

Environment
TensorRT-LLM: 0.17.0.post1
Python: 3.12
CUDA / GPU: (e.g. RTX 6000 Ada / A100 – please fill your actual GPU)
Transformers: 4.48.x (HF version used by Qwen2-VL)
Model type: qwen2_vl
Base model: Qwen2-VL-2B-Instruct
SFT checkpoint: custom SFT on top of Qwen2-VL-2B-Instruct, with extended vocab via new_special_tokens (visual tokens like <|0|>, <|1|>, ..., <|N|>)

Summary

• The official Qwen2-VL-2B-Instruct model, converted and run via examples/multimodal/run.py, works correctly under TensorRT-LLM 0.17.

• A SFT checkpoint of Qwen2-VL-2B-Instruct with extended vocab (large new_special_tokens list like <|0|> ~ <|10000|>) works correctly under PyTorch (HF model.generate()), but when converted to TensorRT-LLM:

1. examples/multimodal/run.py + this SFT checkpoint + a single image + a simple English prompt
2. always generates almost exclusively token id = vocab_size-1 (168063) in the generated segment.

• With a custom decode function that preserves all token ids, the output is a long sequence of <|168063|> repeated, i.e. the model in TRT seems to be “collapsed” to always choose vocab_size-1.

Request
Could you please:

1. Confirm whether extended-vocab SFT checkpoints for Qwen2-VL (with many <|n|> style special tokens) are officially supported in TensorRT-LLM 0.17?
2. If so, help investigate why generation under TRT collapses to vocab_size-1 for such checkpoints, even though HF model.generate() works correctly and configs (HF/TRT ckpt/engine) match.
3. If not yet supported, provide guidance on:

• How to safely extend Qwen2-VL vocab and still be compatible with TensorRT-LLM,
• Or what constraints we should respect when adding new_special_tokens for visual tokens.

SFT checkpoint config:
(https://github.com/MIV-XJTU/FSDrive/blob/main/configs/sft.yaml)

Who can help?

@Tracin @juney-nvidia @ncomly-nvidia

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Model & Config Details
HF SFT config (config.json of SFT checkpoint) (key fields):
{ "model_type": "qwen2_vl", "vocab_size": 168064, "hidden_size": 1536, "num_hidden_layers": 28, "num_attention_heads": 12, "intermediate_size": 8960, "max_position_embeddings": 32768, "rope_scaling": { "type": "mrope", "mrope_section": [16, 24, 24] }, "vision_start_token_id": 151652, "vision_end_token_id": 151653, "vision_token_id": 151654, "image_token_id": 151655, "video_token_id": 151656, "bos_token_id": 151643, "eos_token_id": 151645 }

Tokenizer / vocab extension

• configs/sft.yaml (training config) has:
  new_special_tokens: "<|0|>,<|1|>,...,<|10000|>,..."
• In the SFT checkpoint:
  added_tokens.json and tokenizer.json contain entries like:
  { "<|0|>": 151657, "<|10000|>": 161657, ...}
• So vocab is extended from the original 151936 to 168064, and HF model.generate() correctly produces visual code sequences like <|4509|><|786|>... plus text.

TRT ckpt config (trt_ckpt/fp16/1-gpu/config.json) and
Engine config (trt_engines/fp16/1-gpu/config.json, pretrained_config):

• Both have the same core fields as HF config:
  { "architecture": "Qwen2VLForConditionalGeneration", "vocab_size": 168064, "hidden_size": 1536, "num_hidden_layers": 28, "num_attention_heads": 12, "intermediate_size": 8960, "position_embedding_type": "mrope", "qwen_type": "qwen2_vl", "max_position_embeddings": 32768, ...}
  We also verified with a small script that HF config, TRT ckpt config and engine pretrained_config fully agree on these fields.

Conversion & Engine Build Steps
For the SFT checkpoint, we use the Qwen2-VL-aware wrapper script (similar to FSDrive’s 02_build_engines.sh):

export HF_MODEL_PATH=/path/to/sft/checkpoint  # extended-vocab SFT ckptexport TRTLLM_BACKEND_ROOT=/path/to/tensorrtllm_backendexport OUTPUT_ROOT=/path/to/output_dir  # contains trt_ckpt and trt_engines
bash 02_build_engines.sh

Internally this does roughly:

1. convert_checkpoint_qwen2vl.py (wrapper) → calls examples/qwen/convert_checkpoint.py
2. trtllm-build --max_batch_size 4 --max_input_len 2048 --max_seq_len 3072 --max_multimodal_len 3888
3. build_visual_engine.py --model_type qwen2_vl --model_path $HF_MODEL_PATH

So we have:

• LLM engine: $OUTPUT_ROOT/trt_engines/fp16/1-gpu
• Visual encoder engine: $OUTPUT_ROOT/trt_engines/multimodal_encoder

Minimal Repro Script
Using official examples/multimodal/run.py with this SFT ckpt:

cd /path/to/tensorrtllm_backend/tensorrt_llm/examples/multimodal
export HF_MODEL_PATH=/path/to/sft/checkpoint
export ENGINE_PATH=/path/to/output_dir/trt_engines/fp16/1-gpu
export MM_ENGINE_PATH=/path/to/output_dir/trt_engines/multimodal_encoder
export IMG=/path/to/local/image.jpg
# Enable debug
export TRT_MM_DEBUG_PTUNING=1
export TRT_MM_DEBUG_DECODE=1
python run.py \  
  --hf_model_dir "$HF_MODEL_PATH" \  
  --visual_engine_dir "$MM_ENGINE_PATH" \  
  --llm_engine_dir "$ENGINE_PATH" \  --image_path "$IMG" \  
  --input_text "Describe this driving scene."

Key log lines:
[TRT_MM_DEBUG_PTUNING] num_fake_tokens=324 num_visual=324 prompt_table.shape=(324, 1536)[TRT_MM_DEBUG] decode start output_ids.shape=(1, 1, 3072) input_lengths=[350] vocab_size=168064 pad_id=151643 eos_id=151645 (gen_pad=151643) [TRT_MM_DEBUG] batch=0 beam=0 output_ids.shape=(1, 1, 3072) input_len=350 start=350 seq_len=2722 first5=[168063, 168063, 168063, 168063, 168063] [Q] ['Describe this driving scene.'] [A]: ['<|168063|><|168063|><|168063|><|168063|>... (repeated 168063) ...']Generated 1153 tokens
So:

• input_lengths[0] = 350
• The generated segment [350:] is almost entirely token id 168063, which is vocab_size-1 (168064-1).
• Prompt tuning looks correct: 324 fake tokens == 324 visual tokens, prompt_table (324, 1536).
  For comparison, with official Qwen2-VL-2B-Instruct HF checkpoint + appropriately built engines using the same run.py in the same environment, the output is a normal English description (no 168063 collapse).

Additional Debug: Custom Decode
To rule out tokenizer issues, we implemented a custom decode helper in runtime/multimodal_model_runner.py that:

• Uses convert_ids_to_tokens(ids, skip_special_tokens=False)
• For any None tokens, falls back to _convert_id_to_token or finally to an explicit f"<|{tid}|>" representation
• Then calls convert_tokens_to_string(tokens) or "".join(tokens) as fallback

This ensures we can see all token ids that TRT generates, and avoids NoneType errors in decode.
With this in place, the output for the SFT engine is exactly:
[A]: ['<|168063|><|168063|><|168063|>...']
i.e. the model in TRT is genuinely generating 168063 on almost every step, not just a decode artifact.

Expected behavior

For the same SFT checkpoint, PyTorch HF model.generate() produces:

• Visual tokens: <|4509|><|786|>... (384 code tokens)
• Plus textual description, waypoints, decisions, etc.

TensorRT-LLM should produce a similar distribution over the extended vocab:

• Visual tokens mostly in the extended region (<|0|> ~ <|16383|>),
• Normal text tokens, not concentrated on vocab_size-1.

actual behavior

Under TensorRT-LLM 0.17.0.post1:

• For this SFT checkpoint with extended vocab (vocab_size=168064),
• Generation phase is almost completely collapsed to token id 168063 = vocab_size-1,
• Even when prompt tuning and engine/HF configs are consistent.

additional notes

This issue is specific to the SFT checkpoint with extended vocab:

• Official Qwen2-VL-2B-Instruct + TensorRT-LLM 0.17 + examples/multimodal/run.py works fine (no 168063 collapse).
• Only the extended-vocab SFT ckpt exhibits this behavior.

My suspicion is that the current Qwen/Qwen2-VL conversion and/or runtime code path in TensorRT-LLM does not correctly handle:

• vocab_size extension,
• added_tokens / additional_special_tokens for large ranges like <|0|> ~ <|10000|>,
• And thus ends up with an inconsistent embedding / lm_head region near vocab_size-1.