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

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

I maintainer di solito rispondono entro 2 giorni

@justinchuby ci sta già lavorando.

Dal 10/11/2026.

  • #3186 di @justinchuby — aperta

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
25/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Ferma
Stack tecnologico
python
Ambito
backend, machine-learning

Direzione di ricerca

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.

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

Descrizione

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
Lingua principale
Python
Stelle
456
Fork
159
Merge medio
3g 5h
PR unite (30g)
23

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 microsoft/onnxscript

Tutte le issue di microsoft/onnxscript

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.