microsoft / microsoft/onnxruntime
SimplifiedLayerNormFusion Incorrectly Fuses Downstream Mul with Shape-Incompatible Scale, Causing Runtime Shape Error
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
ONNX Runtime's **SimplifiedLayerNormFusion** optimizer incorrectly fuses a downstream `Mul` node into a LayerNorm pattern when the `Mul`'s scale operand has a shape that doesn't match the input's last dimension.
### To reproduce
```python
from onnx import helper, TensorProto
import onnxruntime as ort
import numpy as np
# Constants
eps_t = helper.make_tensor("eps_v", TensorProto.FLOAT, [], [1.0e-5])
two_t = helper.make_tensor("two_v", TensorProto.FLOAT, [], [2.0])
scale_t = helper.make_tensor("scale_v", TensorProto.FLOAT, [2], [0.0, 0.0])
x_vi = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1])
out_vi = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 2])
nodes = [
helper.make_node("ReduceMean", inputs=["x"], outputs=["mn"], axes=[-1], keepdims=1),
helper.make_node("Sub", inputs=["x", "mn"], outputs=["ct"]),
helper.make_node("Pow", inputs=["ct", "two_v"], outputs=["sq"]),
helper.make_node("ReduceMean", inputs=["sq"], outputs=["vr"], axes=[-1], keepdims=1),
helper.make_node("Add", inputs=["vr", "eps_v"], outputs=["ve"]),
helper.make_node("Sqrt", inputs=["ve"], outputs=["sd"]),
helper.make_node("Div", inputs=["ct", "sd"], outputs=["nm"]),
helper.make_node("Mul", inputs=["scale_v", "nm"], outputs=["out"]),
]
graph = helper.make_graph(
nodes, "g", [x_vi], [out_vi], initializer=[eps_t, two_t, scale_t]
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 11)])
sess = ort.InferenceSession(model.SerializeToString())
result = sess.run(
None, {"x": np.random.randn(1, 1).astype(np.float32)}
)
print("Execution: OK")
```
```txt
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running SimplifiedLayerNormalization node. Name:'/SimplifiedLayerNormFusion/' Status Message: Scale and (optional) bias must match X.shape[axis:] or be NumPy-broadcastable to it. X.shape={1,1} scale.shape={2} bias.shape={} and axis=1
```
### Urgency
No. I found this by fuzzing testing.
### Platform
Linux
### OS Version
5.4.0-162-generic
### ONNX Runtime Installation
Released Package
### ONNX Runtime Version or Commit ID
1.27.0
### ONNX Runtime API
Python
### Architecture
X64
### Execution Provider
CUDA
### Execution Provider Library Version
CUDA 13.0, Driver 580.76.05, GPU: NVIDIA GeForce RTX 3080 Ti
Contributor guide
Assessment
This issue has not been assessed yet.