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

Aperta
  • Qualcosa non va?
#3,245 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

@sylvesterkaczmarek ci sta già lavorando.

Dal 6/11/2026.

  • #3246 di @sylvesterkaczmarek — aperta

Valutazione

Difficoltà
3/5
Tempo stimato
1-2 giorni
Idoneità per principianti
71/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
machine-learning

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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.
Lingua principale
Python
Stelle
5.5k
Fork
800
Merge medio
3g 8h
PR unite (30g)
285

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/Model-Optimizer

Tutte le issue di NVIDIA/Model-Optimizer

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.