Skip to content

Fix SkipLayerNormalization fusion for residual Add outputs #33611

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Harshal Dharpure (hdharpure9922) wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix SkipLayerNormalization fusion for residual Add outputs #33611

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b0f7f74
Fix SkipLayerNormalization fusion safety for residual consumers
hdharpure9922 Oct 26, 2026
de04875
Fix SkipLayerNormalization fusion regression tests
hdharpure9922 Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 6 additions & 12 deletions onnxruntime/python/tools/transformers/fusion_skiplayernorm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,26 +131,20 @@ def fuse(self, node, input_name_to_nodes, output_name_to_node):
)
return

# This means that the residual Add before the LayerNormalization produces an output
# that is consumed by some other nodes or graph output other than the LayerNormalization itself
# We can still go ahead with the SkipLayerNormalization fusion but we need to
# preserve the output of Add and that needs to be produced by SkipLayerNormalization.
# If the residual Add output is consumed elsewhere, we cannot fuse it into a single-output
# SkipLayerNormalization without creating a multi-output custom op. Some backends, including
# OpenVINO, reject this shape at load time. In that case, skip the fusion and leave the original
# Add + LayerNormalization pattern intact so the residual output remains a normal tensor.
add_has_graph_output = self.model.find_graph_output(add.output[0]) is not None
residual_add_has_multiple_consumers = (
add_has_graph_output or len(self.model.get_children(add, input_name_to_nodes)) > 1
)

outputs_to_keep = node.output

if residual_add_has_multiple_consumers:
outputs_to_keep.extend([add.output[0]])
return

outputs_to_keep = node.output
outputs = [node.output[0]]

# Skip the other optional outputs of SkipLayerNormalization before adding the Add's output
if residual_add_has_multiple_consumers:
outputs.extend(["", "", add.output[0]])

if self.model.is_safe_to_fuse_nodes([add, node], outputs_to_keep, input_name_to_nodes, output_name_to_node):
self.nodes_to_remove.extend([add, node])

Expand Down
92 changes: 70 additions & 22 deletions onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,13 +188,13 @@ def test_skip_layer_norm_graph_output(self):
self.verify_skip_layer_norm_fusion(
model_name,
{
"Add": 0,
"LayerNormalization": 0,
"SkipLayerNormalization": 1,
"Add": 2,
"LayerNormalization": 1,
"SkipLayerNormalization": 0,
"Cast": 0,
},
["input_1", "input_2", "layer_norm_weight", "layer_norm_bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

Expand All @@ -205,13 +205,13 @@ def test_skip_layer_norm_graph_output_bias1(self):
self.verify_skip_layer_norm_fusion(
model_name,
{
"Add": 0,
"LayerNormalization": 0,
"SkipLayerNormalization": 1,
"Add": 2,
"LayerNormalization": 1,
"SkipLayerNormalization": 0,
"Cast": 0,
},
["matmul_output", "input_2", "layer_norm_weight", "layer_norm_bias", "bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

Expand All @@ -222,13 +222,13 @@ def test_skip_layer_norm_graph_output_bias2(self):
self.verify_skip_layer_norm_fusion(
model_name,
{
"Add": 0,
"LayerNormalization": 0,
"SkipLayerNormalization": 1,
"Add": 2,
"LayerNormalization": 1,
"SkipLayerNormalization": 0,
"Cast": 0,
},
["matmul_output", "input_1", "layer_norm_weight", "layer_norm_bias", "bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

Expand All @@ -241,13 +241,13 @@ def test_skip_layer_norm_graph_output_cast_bias1(self):
self.verify_skip_layer_norm_fusion(
model_name,
{
"Add": 0,
"LayerNormalization": 0,
"SkipLayerNormalization": 1,
"Add": 2,
"LayerNormalization": 1,
"SkipLayerNormalization": 0,
"Cast": 1,
},
["matmul_output_cast", "input_2", "layer_norm_weight", "layer_norm_bias", "bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

Expand All @@ -260,14 +260,62 @@ def test_skip_layer_norm_graph_output_cast_bias2(self):
self.verify_skip_layer_norm_fusion(
model_name,
{
"Add": 0,
"LayerNormalization": 0,
"SkipLayerNormalization": 1,
"Add": 2,
"LayerNormalization": 1,
"SkipLayerNormalization": 0,
"Cast": 1,
},
["matmul_output_cast", "input_1", "layer_norm_weight", "layer_norm_bias", "bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

def test_skip_layer_norm_multiple_consumers_block_fusion(self):
input_1 = helper.make_tensor_value_info("input_1", TensorProto.FLOAT, [1, 2, 3])
input_2 = helper.make_tensor_value_info("input_2", TensorProto.FLOAT, [1, 2, 3])
output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 2, 3])
consumer_output = helper.make_tensor_value_info("consumer_output", TensorProto.FLOAT, [1, 2, 3])

add_before_layer_norm = helper.make_node("Add", ["input_1", "input_2"], ["layernorm_input"], "add_layernorm")
layer_norm = helper.make_node(
"LayerNormalization",
["layernorm_input", "layer_norm_weight", "layer_norm_bias"],
["output"],
"layernorm",
axis=-1,
epsilon=0.000009999999747378752,
)
add_consumer = helper.make_node(
"Add",
["layernorm_input", "consumer_bias"],
["consumer_output"],
"add_consumer",
)

initializers = [
float_tensor("layer_norm_weight", [3]),
float_tensor("layer_norm_bias", [3]),
float_tensor("consumer_bias", [1, 2, 3]),
]

graph = helper.make_graph(
[add_before_layer_norm, layer_norm, add_consumer],
"skip_layer_norm_multiple_consumers",
[input_1, input_2],
[output, consumer_output],
initializers,
)
model = helper.make_model(graph, opset_imports=(helper.make_opsetid("ai.onnx", 16),))
model_name = "skip_layer_norm_add_multiple_consumers.onnx"
onnx.save(model, model_name)

optimized_model = optimize_model(model_name, optimization_options=FusionOptions("bert"), opt_level=0)

self.assertEqual(len(optimized_model.get_nodes_by_op_type("Add")), 2)
self.assertEqual(len(optimized_model.get_nodes_by_op_type("LayerNormalization")), 1)
self.assertEqual(len(optimized_model.get_nodes_by_op_type("SkipLayerNormalization")), 0)
self.assertIn("consumer_output", [node.output[0] for node in optimized_model.get_nodes()])

os.remove(model_name)

def create_broadcast_test_model(
Expand Down Expand Up @@ -396,9 +444,9 @@ def test_skip_layer_norm_broadcast_graph_output(self):
onnx.save(model, model_name)
self.verify_skip_layer_norm_fusion(
model_name,
{"Add": 0, "LayerNormalization": 0, "SkipLayerNormalization": 1, "Cast": 0},
{"Add": 1, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 0},
["input_1", "input_2", "layer_norm_weight", "layer_norm_bias"],
["output", "", "", "layernorm_input"],
["output"],
)
os.remove(model_name)

Expand Down