NVIDIA / NVIDIA/TensorRT-Edge-LLM

[BUG] Qwen3-TTS CP-FP8 calibration failed, it mixes text and codec embedding domains

Open
#175 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
563
Forks
135
Avg merge
14h 13m
Merged PRs (30d)
1

Description

Describe the bug

v0.10.0 documents FP8 quantization support for the Qwen3-TTS CodePredictor through:

tensorrt-edgellm-quantize llm ... --cp_quantization fp8

However, the Qwen3-TTS CodePredictor calibration path does not preserve the model's separate text-token and codec-token embedding domains.

This is a blocker for the documented Qwen3-TTS CodePredictor FP8 workflow. The supported precision contract is otherwise appropriate for this model:

  • Talker: FP16
  • CodePredictor: FP8 only
  • Code2Wav: FP16
  • KV cache: FP16

The issue is in tensorrt_edgellm/quantization/qwen3_cp_loader.py.

For Qwen3-TTS, _talker_inputs_from_text() sends tokenizer text IDs to talker.get_input_embeddings():

def _talker_inputs_from_text(model, input_ids):
    talker = model.talker
    talker_embed = talker.get_input_embeddings()
    return talker_embed(input_ids).to(talker.dtype)

The same talker_embed is then reused for the initial CodePredictor input:

talker_vocab = talker_cfg.text_config.vocab_size
# ...
random_token = torch.randint(0, talker_vocab, (bsz, 1), device=device)
last_token_embed = talker_embed(random_token).to(talker.dtype)

For the official Qwen3-TTS model, these are distinct domains:

  1. Text tokens must use the Talker text embedding followed by text_projection.
  2. Talker codec tokens use the Talker codec embedding.
  3. CodePredictor residual tokens use the CodePredictor per-codebook codec embeddings.

The exporter already reflects this separation by emitting both text_embedding.safetensors / text_projection.safetensors and codec embedding sidecars. The calibration path does not.

As a result, text-token IDs, or random IDs sampled from the text vocabulary, can be passed to a codec embedding table. This leads to CUDA embedding index-out-of-range failures. Mapping get_input_embeddings() to the text embedding is not a valid workaround either, because the later CP seed requires codec-embedding semantics.

There is also a loading issue in the documented generic quantization path. The official Qwen3-TTS checkpoint uses model_type: "qwen3_tts", while the generic quantizer loads via Transformers AutoModel factories and does not have a Qwen3-TTS-specific loader or qwen_tts registration path. With the stock Transformers package, model loading can fail before calibration with an unrecognized qwen3_tts configuration/model type.

This report concerns the calibration-interface mismatch after model loading. A valid fix likely needs a dedicated Qwen3-TTS loader and calibration adapter rather than changing the meaning of get_input_embeddings() globally.

Steps/Code to reproduce bug
  1. Start from TensorRT Edge-LLM v0.10.0.
  2. Install the tool dependencies and the official Qwen TTS package.
  3. Download the official Qwen3-TTS CustomVoice checkpoint.
  4. Run CP-only FP8 quantization as documented.

Installation method:

Built TensorRT Edge-LLM from source at tag v0.10.0 inside a Docker container.

git checkout v0.10.0
pip install ".[tools]"
pip install qwen-tts==0.1.1

Export command used:

export MODEL_DIR=/workspace/models/Qwen3-TTS-12Hz-0.6B-CustomVoice
export QUANTIZED_MODEL_DIR=/workspace/models/Qwen3-TTS-12Hz-0.6B-CustomVoice-CP-FP8

tensorrt-edgellm-quantize llm \
    --model_dir "$MODEL_DIR" \
    --output_dir "$QUANTIZED_MODEL_DIR" \
    --cp_quantization fp8 \
    --dtype fp16 \
    --device cuda \
    --num_samples 64

The same issue applies when using the official Hugging Face model ID instead of a local directory:

tensorrt-edgellm-quantize llm \
    --model_dir Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
    --output_dir /workspace/models/Qwen3-TTS-12Hz-0.6B-CustomVoice-CP-FP8 \
    --cp_quantization fp8 \
    --dtype fp16 \
    --device cuda \
    --num_samples 64

Observed failure modes:

  1. With stock Transformers AutoModel loading, the official checkpoint can fail to load because model_type="qwen3_tts" is not registered by Transformers alone.
  2. When the official qwen-tts model is loaded through a compatibility/registration path, calibration reaches the TTS CodePredictor loop and can fail with CUDA embedding index-out-of-range because text-vocabulary IDs are used with a codec embedding table.

Relevant source locations in v0.10.0:

  • tensorrt_edgellm/quantization/qwen3_cp_loader.py
    • _talker_inputs_from_text() uses talker.get_input_embeddings() for tokenizer text IDs.
    • random_token is sampled from talker_cfg.text_config.vocab_size and passed to that same embedding.
  • tensorrt_edgellm/quantization/quantize.py
    • The generic _load_model() path uses AutoModel factories.
    • No Qwen3-TTS-specific loader is present, unlike the dedicated Qwen3-ASR loader.
Expected behavior

The documented command below should successfully produce a CP-only FP8 quantized checkpoint for the official Qwen3-TTS checkpoint:

tensorrt-edgellm-quantize llm \
    --model_dir Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
    --output_dir ./qwen3-tts-cp-fp8 \
    --cp_quantization fp8

The implementation should:

  1. Load the official Qwen3-TTS checkpoint through an explicit Qwen3-TTS loader/registration path.
  2. Keep Talker, Code2Wav, and KV cache in FP16.
  3. Quantize only talker.code_predictor.* to FP8, while retaining down_proj, CP LM heads, and KV-cache BMM in FP16 as documented.
  4. Drive calibration with the model's real TTS input contract:
    • text token IDs -> text embedding -> text projection -> Talker;
    • codec token IDs -> codec embedding -> Talker/CodePredictor;
    • residual CP tokens -> the corresponding CodePredictor codec embedding.
  5. Avoid using a single get_input_embeddings() API for both text and codec token domains.

System information (x86 Host with GPU)

======================================================================

  • Container used (if applicable): Yes.
  • OS (e.g., Ubuntu 22.04, CentOS 7): Ubuntu 24.04.3 LTS
  • CPU architecture: x86_64
  • GPU name (e.g. H100, A100, RTX 4090): NVIDIA L4
  • GPU memory size: 24.0 GB
  • Number of GPUs: 2
  • Library versions:
    • Python: 3.12.3
    • TensorRT Edge-LLM version or commit hash: v0.10.0 (71dd1bae032e70771265917ec74d3ff4cad07a10)
    • CUDA: 13.1
    • PyTorch: 2.13.0+cu130
    • Transformers: 5.14.1
    • ModelOpt: 0.44.0
    • ONNX: 1.19.0
  • Any other details that may help:
    • Model: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
    • The same model was also reproduced from a local downloaded checkpoint directory.
    • qwen-tts==0.1.1 is installed.
    • The TensorRT Edge-LLM v0.10.0 pyproject.toml pins nvidia-modelopt==0.45.0, while this environment currently has ModelOpt 0.44.0. The embedding-domain issue is visible in the v0.10.0 calibration source independently of that version difference.
    • TTS FP16 export/build/inference works; this report is specifically about the documented CP-only FP8 quantization path.

======================================================================

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with tensorrt_edgellm/quantization/qwen3_cp_loader.py and quantize.py, then reproduce the documented CP-only FP8 command for the Qwen3-TTS checkpoint. Trace the text, codec, and residual embedding inputs and the generic model-loading path. Done means the official checkpoint loads and calibration produces the documented precision split without mixing token domains.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.