microsoft / microsoft/onnxruntime

LayerNorm fusion selects the Div output as scale, failing session creation even with keepdims=1

Open
#31,147 1 comment 0 reactions 0 assignees View on GitHub
stale
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

Follow-up to #30513. PR #31144 stops the fusion when ReduceMean has keepdims=0, which fixes the incorrect matching reported there. The session creation failure in that issue has a second, independent cause that also reproduces with keepdims=1, so it remains after #31144.

In LayerNormFusion and SimplifiedLayerNormFusion (onnxruntime/core/optimizer/layer_norm_fusion.cc), the scale selection loop picks the Mul operand whose rank matches the number of reduction axes. When the normalized input is rank 1, the Div output is also rank 1 and passes that test, so it is selected as scale. The Div node itself is removed by the fusion, so the fused node references a deleted output and session creation fails with:

```
Invalid model. Node input 'nm' is not a graph input, initializer, or output of a previous node.
```

The bias selection loop in LayerNormFusion can pick the Mul output the same way.

Repro on onnxruntime 1.26.0 with keepdims=1 everywhere. Fails at both opset 11 and opset 17:

```python
from onnx import helper, TensorProto
import onnxruntime as ort

def build(opset):
eps = helper.make_tensor("eps_v", TensorProto.FLOAT, [], [1.0e-5])
two = helper.make_tensor("two_v", TensorProto.FLOAT, [], [2.0])
nodes = [
helper.make_node("Constant", [], ["eps"], value=eps),
helper.make_node("Constant", [], ["two"], value=two),
helper.make_node("ReduceMean", ["x"], ["mn"], axes=[-1], keepdims=1),
helper.make_node("Sub", ["x", "mn"], ["ct"]),
helper.make_node("Pow", ["ct", "two"], ["sq"]),
helper.make_node("ReduceMean", ["sq"], ["vr"], axes=[-1], keepdims=1),
helper.make_node("Add", ["vr", "eps"], ["ve"]),
helper.make_node("Sqrt", ["ve"], ["sd"]),
helper.make_node("Div", ["ct", "sd"], ["nm"]),
helper.make_node("Mul", ["nm", "z"], ["mo"]),
helper.make_node("Add", ["mo", "b"], ["out"]),
]
graph = helper.make_graph(nodes, "g", [
helper.make_tensor_value_info("x", TensorProto.FLOAT, [2]),
helper.make_tensor_value_info("z", TensorProto.FLOAT, [2, 1]),
helper.make_tensor_value_info("b", TensorProto.FLOAT, [2]),
], [helper.make_tensor_value_info("out", TensorProto.FLOAT, [2, 2])])
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)])

for opset in (11, 17):
try:
ort.InferenceSession(build(opset).SerializeToString())
print(f"opset {opset}: loaded")
except Exception as e:
print(f"opset {opset}: {e}")
```

Output:

```
opset 11: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Invalid model. Node input 'nm' is not a graph input, initializer, or output of a previous node.
opset 17: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Invalid model. Node input 'nm' is not a graph input, initializer, or output of a previous node.
```

Expected: the model loads, with the fusion either applied correctly or skipped. Operands produced inside the matched subgraph should not be eligible as scale or bias. I will follow up with a PR.

Contributor guide

Open the contributing guide

Research direction

Start with onnxruntime/core/optimizer/layer_norm_fusion.cc and reproduce the failure using the provided Python model at opsets 11 and 17. Inspect the scale and bias selection loops in LayerNormFusion and SimplifiedLayerNormFusion, then verify that operands produced inside the matched subgraph are excluded and the session loads with fusion applied or skipped.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.