[ONNX] graph_surgery / llm_export compute head_dim as hidden_size // num_attention_heads, ignoring config.head_dim (breaks Qwen3, Gemma, Mistral-Small)

Đang mở
  • Có gì không ổn?
#3,245 3 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

@sylvesterkaczmarek đang làm issue này rồi.

Từ ngày 6/11/2026.

  • #3246 của @sylvesterkaczmarek — đang mở

Đánh giá

Độ khó
3/5
Thời gian dự kiến
1-2 ngày
Mức phù hợp với người mới
71/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python
Lĩnh vực
machine-learning

Hướng nghiên cứu

Read modelopt/onnx/graph_surgery/utils/rope_cache.py, modelopt/onnx/graph_surgery/gqa_replacement.py, and modelopt/onnx/llm_export_utils/export_utils.py, comparing their head-size calculation with onnx/quantization/ort_utils.py and torch/quantization/kv_cache_auto_quant.py. Start by examining tests/unit/onnx/test_gqa_graph_surgery.py and its config fixtures. Done when all three ONNX paths use config.head_dim when present, retain the fallback, and regression coverage checks a config where the values differ.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

bug

Describe the bug

Three ONNX code paths derive the attention head size as hidden_size // num_attention_heads and never read config.head_dim:

  • modelopt/onnx/graph_surgery/utils/rope_cache.py:65 (get_rope_caches): cos/sin cache width
  • modelopt/onnx/graph_surgery/gqa_replacement.py:907 (replace_attention_with_gqa): KV-cache I/O shapes, packed QKV dim, and the GQA scale=1/sqrt(head_dim)
  • modelopt/onnx/llm_export_utils/export_utils.py:109 (llm_to_onnx): dummy past_key_values used for export tracing

Many current models set head_dim independently of hidden_size / num_attention_heads, so these paths produce wrong shapes and/or a wrong attention scale:

Model hidden_size // heads config.head_dim
Qwen/Qwen3-0.6B 1024 // 16 = 64 128
Qwen/Qwen3-4B 2560 // 32 = 80 128
mistralai/Mistral-Small-24B-Instruct-2501 5120 // 32 = 160 128

Gemma-2/3 (head_dim 256) are affected too, even though gqa_replacement.py has Gemma-specific handling.

This is the same root cause that was already fixed for Puzzletron in #2482 (see also #3016 / #3112). Other code in this repo already handles it correctly: onnx/quantization/ort_utils.py and torch/quantization/kv_cache_auto_quant.py prefer config.head_dim and only fall back to the division.

Impact: should have. ONNX GQA export/surgery for Qwen3-class models produces wrong cache shapes / attention scale.

Steps/Code to reproduce bug
from modelopt.onnx.graph_surgery.utils.rope_cache import get_rope_caches

cos, sin, cfg = get_rope_caches("Qwen/Qwen3-0.6B", max_seq_len=16)
print(cos.shape, cfg.head_dim)

Output: (16, 32) 128

Control (Qwen/Qwen2.5-0.5B, no separate head_dim): (16, 32), which is correct, since 896 // 14 = 64 = head_dim.

Expected behavior

cos.shape == (16, 64) (head_dim 128 // 2). Actual is (16, 32): the cache is built from 1024 // 16 = 64 instead of config.head_dim = 128.

Suggested fix: use the pattern already used elsewhere in the repo, ideally via one shared helper:

head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads

tests/unit/onnx/test_gqa_graph_surgery.py currently uses a Qwen2.5-0.5B-shaped config (896 / 14 = 64, no separate head_dim), where both formulas agree, so CI cannot catch this. A regression test should use a config where head_dim != hidden_size // num_attention_heads.

I'm happy to send a PR for this.

Who can help?
  • ?

System information

  • Container used (if applicable): none
  • OS: Windows 11 Pro 10.0.26200
  • CPU architecture: x86_64 (AMD64)
  • GPU: none used (CPU-only repro; Intel Iris Xe integrated graphics)
  • Library versions:
    • Python: 3.11.1
    • ModelOpt version or commit hash: bc00ef3 (main, checked 2026-11-06)
    • PyTorch: 2.15.0+cpu
    • Transformers: 5.22.1
  • Any other details that may help: found by code reading, reproduced on CPU, and checked against the published HF config.json files.
Ngôn ngữ chính
Python
Star
5.5k
Fork
800
Merge trung bình
3 ngày 8 giờ
Pull request đã merge (30 ngày)
285

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/Model-Optimizer

Tất cả issue của NVIDIA/Model-Optimizer

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.