Partial rotary fusion sets `rotary_embedding_dim` without `num_heads`, which ONNX Runtime requires

未关闭
  • 哪里不对？
#3,185 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 2 天内回复

@justinchuby 已经在做这个了。

开始于 2026年11月10日。

  • #3186 来自 @justinchuby —— 未关闭

评估

难度
2/5
预计耗时
1-3 小时
新手友好度
25/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
停滞
技术栈
python
领域
backend, machine-learning

调研方向

Read PartialRotaryEmbeddingFusion.rewrite in the rotary embedding fusion implementation and inspect the linked pull request #3186. The repro in this issue shows that fusion creates a RotaryEmbedding with rotary_embedding_dim but no num_heads; use it to check the failure. Done means the fused model runs in ONNX Runtime with the expected head count.

由索引模型根据 Issue 内容生成。

描述

PartialRotaryEmbeddingFusion.rewrite (fuse_partial_rotary_embedding, run by optimize_for_ort) rebuilds ONNX Runtime's RotaryEmbedding over the whole head with rotary_embedding_dim set, but keeps the original node's attributes, which have no num_heads. ONNX Runtime requires num_heads whenever rotary_embedding_dim is set, so a model rotating only part of each head (GPT-NeoX's rotary_pct) runs before the fusion and fails after it. We hit it on GPT-NeoX. Setting num_heads in the rewrite (axis 1 of a 4D input) would fix it.

Repro (onnxscript 0.7.2, onnx-ir 1.1.0, onnxruntime 1.31.1):

"""onnxscript's partial rotary fusion (`fuse_partial_rotary_embedding`, part of `optimize_for_ort`) rebuilds
ONNX Runtime's `RotaryEmbedding` with `rotary_embedding_dim` but without `num_heads`, which ONNX Runtime requires
alongside it: a model rotating part of each head (GPT-NeoX's) runs before the fusion and fails after it."""

import numpy as np
import onnx_ir as ir
import onnxruntime as ort
import onnxscript
from onnx import TensorProto, helper, numpy_helper
from onnxscript.rewriter.ort_fusions.rotary_embedding import fuse_partial_rotary_embedding

B, H, S, D, ROTARY = 2, 4, 5, 16, 4
constants = {
    "zero": [0], "rotary": [ROTARY], "end": [np.iinfo(np.int64).max], "axis": [3], "step": [1],
}
initializers = [numpy_helper.from_array(np.array(value, np.int64), name) for name, value in constants.items()]
rng = np.random.default_rng(0)
initializers += [numpy_helper.from_array(rng.standard_normal((S, ROTARY // 2)).astype(np.float32), name) for name in ("cos", "sin")]
graph = helper.make_graph(
    [
        helper.make_node("Slice", ["x", "zero", "rotary", "axis", "step"], ["rotated_part"]),
        helper.make_node("Slice", ["x", "rotary", "end", "axis", "step"], ["passed_part"]),
        helper.make_node("RotaryEmbedding", ["rotated_part", "positions", "cos", "sin"], ["rotated"], domain="com.microsoft"),
        helper.make_node("Concat", ["rotated", "passed_part"], ["y"], axis=-1),
    ],
    "g",
    [helper.make_tensor_value_info("x", TensorProto.FLOAT, [B, H, S, D]), helper.make_tensor_value_info("positions", TensorProto.INT64, [B, S])],
    [helper.make_tensor_value_info("y", TensorProto.FLOAT, [B, H, S, D])],
    initializers,
)
proto = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.microsoft", 1)], ir_version=10)
feeds = {"x": rng.standard_normal((B, H, S, D)).astype(np.float32), "positions": np.tile(np.arange(S), (B, 1))}


def run(model) -> str:
    try:
        session = ort.InferenceSession(ir.to_proto(model).SerializeToString(), providers=["CPUExecutionProvider"])
        return f"runs, y[0, 0, 1, :4] = {session.run(None, feeds)[0][0, 0, 1, :4]}"
    except Exception as error:
        return f"fails: {str(error).splitlines()[0].split('was false. ')[-1]}"


ort.set_default_logger_severity(3)
print("onnxscript", onnxscript.__version__, "| onnx_ir", ir.__version__, "| onnxruntime", ort.__version__)
model = ir.from_proto(proto)
print("original:", run(model))
print("fuse_partial_rotary_embedding:", fuse_partial_rotary_embedding(model), "fusion(s)")
rotary = next(node for node in model.graph if node.op_type == "RotaryEmbedding")
print("fused attributes:", {name: attribute.value for name, attribute in rotary.attributes.items()})
print("fused:", run(model))

Output:


onnxscript 0.7.2 | onnx_ir 1.1.0 | onnxruntime 1.31.1
original: runs, y[0, 0, 1, :4] = [1.4040405  0.3131538  2.0222473  0.18501948]
fuse_partial_rotary_embedding: 1 fusion(s)
fused attributes: {'rotary_embedding_dim': 4}
fused: fails: num_heads must be provided if rotary_embedding_dim is specified
主要语言
Python
星标
456
派生
159
平均合并
3 天 5 小时
30 天内合并 PR
23

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

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

microsoft/onnxscript 的其他 Issue

查看 microsoft/onnxscript 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

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