From b0f7f74e3a5c16e002181fe0f950ba4f5d24c629 Mon Sep 17 00:00:00 2001 From: hdharpure9922 Date: Mon, 26 Oct 2026 05:34:57 -0700 Subject: [PATCH 1/2] Fix SkipLayerNormalization fusion safety for residual consumers --- .../transformers/fusion_skiplayernorm.py | 18 +++----- .../test_skip_layer_norm_fusion.py | 44 +++++++++---------- 2 files changed, 28 insertions(+), 34 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 743bf50f6c608..34e9d2b140c49 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -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]) diff --git a/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py b/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py index 454880f5b33b2..0206213604a3d 100644 --- a/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py +++ b/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py @@ -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": 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) @@ -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": 1, + "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) @@ -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": 1, + "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) @@ -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": 1, + "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) @@ -260,13 +260,13 @@ 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": 1, + "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) @@ -396,9 +396,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) From de0487500732d37801e0f5857b0703c15b300751 Mon Sep 17 00:00:00 2001 From: hdharpure9922 Date: Mon, 26 Oct 2026 22:58:26 -0700 Subject: [PATCH 2/2] Fix SkipLayerNormalization fusion regression tests --- .../test_skip_layer_norm_fusion.py | 58 +++++++++++++++++-- 1 file changed, 53 insertions(+), 5 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py b/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py index 0206213604a3d..71726ba9f010c 100644 --- a/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py +++ b/onnxruntime/test/python/transformers/test_skip_layer_norm_fusion.py @@ -188,7 +188,7 @@ def test_skip_layer_norm_graph_output(self): self.verify_skip_layer_norm_fusion( model_name, { - "Add": 1, + "Add": 2, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 0, @@ -205,7 +205,7 @@ def test_skip_layer_norm_graph_output_bias1(self): self.verify_skip_layer_norm_fusion( model_name, { - "Add": 1, + "Add": 2, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 0, @@ -222,7 +222,7 @@ def test_skip_layer_norm_graph_output_bias2(self): self.verify_skip_layer_norm_fusion( model_name, { - "Add": 1, + "Add": 2, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 0, @@ -241,7 +241,7 @@ def test_skip_layer_norm_graph_output_cast_bias1(self): self.verify_skip_layer_norm_fusion( model_name, { - "Add": 1, + "Add": 2, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 1, @@ -260,7 +260,7 @@ def test_skip_layer_norm_graph_output_cast_bias2(self): self.verify_skip_layer_norm_fusion( model_name, { - "Add": 1, + "Add": 2, "LayerNormalization": 1, "SkipLayerNormalization": 0, "Cast": 1, @@ -270,6 +270,54 @@ def test_skip_layer_norm_graph_output_cast_bias2(self): ) 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( self, batch_size: int = 2,