[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 个 reaction 已指派 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 小时
30 天内合并 PR
285

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/Model-Optimizer 的其他 Issue

查看 NVIDIA/Model-Optimizer 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。