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

Abierto
  • ¿Algo no está bien?
#3,245 3 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

@sylvesterkaczmarek ya está trabajando en esto.

Desde el 6/11/2026.

  • #3246 de @sylvesterkaczmarek — abierto

Evaluación

Dificultad
3/5
Tiempo estimado
1-2 días
Aptitud para principiantes
71/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.
Lenguaje dominante
Python
Estrellas
5.5k
Forks
800
Merge medio
3 d 8 h
PR fusionados (30 d)
285

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/Model-Optimizer

Todos los issues de NVIDIA/Model-Optimizer

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.