microsoft / microsoft/onnxruntime
LayerNorm fusion selects the Div output as scale, failing session creation even with keepdims=1
- 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
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