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

オープン
  • 何か問題がありますか？
#3,245 コメント 3 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

@sylvesterkaczmarek がすでに取り組んでいます。

2026年11月6日 から。

  • #3246 @sylvesterkaczmarek による — オープン

評価

難易度
3/5
見積もり時間
1〜2日
初心者へのやさしさ
71/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
machine-learning

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

説明

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.
主要言語
Python
スター
5.5k
フォーク
800
平均マージ
3日 8時間
マージ済み PR（30日）
285

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/Model-Optimizer のほかの issue

NVIDIA/Model-Optimizer の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。