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
Nobody has claimed this yet.
- 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
-
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 -
Export to ONNX with --mtp:
tensorrt-edgellm-export
/path/to/quant/091/Qwen3.5-4B
/path/to/onnx/091/Qwen3.5-4B/
--mtp -
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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