microsoft / microsoft/onnxruntime

SimplifiedLayerNormFusion Incorrectly Fuses Downstream Mul with Shape-Incompatible Scale, Causing Runtime Shape Error

Open
#31,980 0 comments 1 reaction 2 assignees Claimed by @tianleiwu View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.