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

Inizia da tensorrt_edgellm/quantization/qwen3_cp_loader.py e quantize.py, quindi riproduci il comando CP-only FP8 documentato per il checkpoint Qwen3-TTS. Traccia gli input di testo, codec e residual embedding, oltre al percorso generico di caricamento del modello. Il lavoro è completato quando il checkpoint ufficiale viene caricato e la calibrazione produce la suddivisione di precisione documentata senza mescolare i domini dei token.

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.

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