NVIDIA / NVIDIA/TensorRT-Edge-LLM

[BUG] ONNX Export fails for Qwen3.5 MTP draft model with 4-bit quantized checkpoint due to unpacked mtp.fc.weight shape mismatch

Open
#168 2 comments 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

Bug Description

When attempting to export an int4_awq quantized Qwen3.5-4B checkpoint to ONNX using tensorrt-edgellm-export with the --mtp flag, the export process crashes during the [MTP Draft] tracing phase with a matrix shape mismatch error.

Reproduction Steps
  1. Quantize the model:

    tensorrt-edgellm-quantize llm \
        --model_dir /path/to/Qwen3.5-4B \
        --output_dir /path/to/quant/091/Qwen3.5-4B/ \
        --quantization int4_awq
    
    
  2. Export to ONNX with --mtp:
    tensorrt-edgellm-export
    /path/to/quant/091/Qwen3.5-4B
    /path/to/onnx/091/Qwen3.5-4B/
    --mtp

  3. The export crashes during MTP draft ONNX export with the following RuntimeError:

File ".../tensorrt_edgellm/models/qwen3_5/modeling_qwen3_5_mtp.py", line 230, in forward
fused_hidden_states = self.fc(torch.cat((normed_embeds, normed_hidden_states), dim=-1))
File ".../tensorrt_edgellm/models/linear.py", line 171, in forward
return F.linear(hidden_states, self.weight, bias)

RuntimeError: a and b must have same reduction dim, but got [s23*s30, 5120] X [2560, 2560].
torch.onnx._internal.exporter._errors.TorchExportError: Failed to export the model with torch.export.

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

Reproduce the failure with tensorrt-edgellm-quantize and tensorrt-edgellm-export using the Qwen3.5-4B int4_awq checkpoint and --mtp flag. Inspect the MTP tracing path in tensorrt_edgellm/models/qwen3_5/modeling_qwen3_5_mtp.py around line 230 and the linear operation in tensorrt_edgellm/models/linear.py around line 171. Done means the quantized checkpoint exports to ONNX with MTP without the matrix-dimension or torch.export errors.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.