Skip to content

Do not fuse LayerNorm when ReduceMean has keepdims=0 #31144

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
guptaishaan wants to merge 1 commit into
base: main
Choose a base branch
from
+99 −2
Open

Do not fuse LayerNorm when ReduceMean has keepdims=0 #31144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 12 additions & 2 deletions onnxruntime/core/optimizer/layer_norm_fusion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,14 @@ static bool CheckAxesOnReduceMean(std::vector<int64_t>& axes_values, int64_t ran
return true;
}

// The fused ops broadcast the reduced result back over the reduced axes, so the sub-graph is only
// equivalent to a layer norm when ReduceMean keeps the reduced dims.
static bool KeepDimsOnReduceMean(const Node& reduce_mean_node) {
const onnxruntime::NodeAttributes& attributes = reduce_mean_node.GetAttributes();
auto it = attributes.find("keepdims");
return it == attributes.end() || it->second.i() != 0;
}

static std::vector<int64_t> GetAxesFromReduceMeanNode(Node& reduce_mean_node, const Graph& graph) {
const onnxruntime::NodeAttributes& attributes = reduce_mean_node.GetAttributes();
std::vector<int64_t> axes_values;
Expand Down Expand Up @@ -263,6 +271,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
!graph_utils::IsSupportedProvider(reduce_mean_node, GetCompatibleExecutionProviders()) ||
(reduce_mean_node.GetOutputEdgesCount() != 1 && reduce_mean_node.GetOutputEdgesCount() != 2) ||
graph.NodeProducesGraphOutput(reduce_mean_node) ||
!KeepDimsOnReduceMean(reduce_mean_node) ||
!IsSupportedDataType(reduce_mean_node, 1)) {
continue;
}
Expand Down Expand Up @@ -426,6 +435,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level,
if (!graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean2_node, "ReduceMean", {1, 11, 13, 18}) ||
reduce_mean2_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() ||
!optimizer_utils::CheckOutputEdges(graph, reduce_mean2_node, 1) ||
!KeepDimsOnReduceMean(reduce_mean2_node) ||
!IsSupportedDataType(reduce_mean2_node, 1) ||
reduce_mean2_node.GetInputEdgesCount() == 0) {
continue;
Expand Down Expand Up @@ -667,8 +677,8 @@ Status SimplifiedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int gr
Node& reduce_mean_node = *graph.GetNode(p_reduce_mean->Index());
if (!graph_utils::IsSupportedOptypeVersionAndDomain(reduce_mean_node, "ReduceMean", {1, 11, 13, 18}) ||
reduce_mean_node.GetExecutionProviderType() != pow_node.GetExecutionProviderType() ||
!optimizer_utils::CheckOutputEdges(graph, reduce_mean_node, 1) || !IsSupportedDataType(reduce_mean_node, 1) ||
reduce_mean_node.GetInputEdgesCount() == 0) {
!optimizer_utils::CheckOutputEdges(graph, reduce_mean_node, 1) || !KeepDimsOnReduceMean(reduce_mean_node) ||
!IsSupportedDataType(reduce_mean_node, 1) || reduce_mean_node.GetInputEdgesCount() == 0) {
continue;
}
nodes_to_remove.push_back(reduce_mean_node);
Expand Down
87 changes: 87 additions & 0 deletions onnxruntime/test/optimizer/graph_transform_test_layernorm.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,6 +466,93 @@ TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest_9) {
TransformerLevel::Level2, 1, nullptr, post_graph_checker));
}

// ReduceMean with keepdims=0 drops the reduced dim, so the sub-graph does not compute a layer norm.
TEST_F(GraphTransformationTests, LayerNormFusion_ReduceMeanNoKeepDims_NoFusion) {
auto build_test_case = [](ModelTestBuilder& builder) {
auto* input = builder.MakeInput<float>({{4, 4}});
auto* pow_exponent = builder.MakeInitializer<float>({}, {2.0f});
auto* epsilon = builder.MakeInitializer<float>({}, {1e-5f});
auto* scale = builder.MakeInitializer<float>({4}, {1.0f, 1.0f, 1.0f, 1.0f});
auto* bias = builder.MakeInitializer<float>({4}, {0.0f, 0.0f, 0.0f, 0.0f});

auto* mean_out = builder.MakeIntermediate();
auto* sub_out = builder.MakeIntermediate();
auto* pow_out = builder.MakeIntermediate();
auto* reduce_mean_out = builder.MakeIntermediate();
auto* add_out = builder.MakeIntermediate();
auto* sqrt_out = builder.MakeIntermediate();
auto* div_out = builder.MakeIntermediate();
auto* mul_out = builder.MakeIntermediate();
auto* output = builder.MakeOutput();

Node& mean_node = builder.AddNode("ReduceMean", {input}, {mean_out});
mean_node.AddAttribute("axes", std::vector<int64_t>{-1});
mean_node.AddAttribute("keepdims", static_cast<int64_t>(0));
builder.AddNode("Sub", {input, mean_out}, {sub_out});
builder.AddNode("Pow", {sub_out, pow_exponent}, {pow_out});
Node& reduce_mean_node = builder.AddNode("ReduceMean", {pow_out}, {reduce_mean_out});
reduce_mean_node.AddAttribute("axes", std::vector<int64_t>{-1});
reduce_mean_node.AddAttribute("keepdims", static_cast<int64_t>(0));
builder.AddNode("Add", {reduce_mean_out, epsilon}, {add_out});
builder.AddNode("Sqrt", {add_out}, {sqrt_out});
builder.AddNode("Div", {sub_out, sqrt_out}, {div_out});
builder.AddNode("Mul", {div_out, scale}, {mul_out});
builder.AddNode("Add", {mul_out, bias}, {output});
};

auto post_graph_checker = [](Graph& graph) {
const auto op_to_count = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_to_count.find("LayerNormalization") == op_to_count.end());
const auto reduce_mean_it = op_to_count.find("ReduceMean");
TEST_RETURN_IF_NOT(reduce_mean_it != op_to_count.end() && reduce_mean_it->second == 2);
return Status::OK();
};

const InlinedHashSet<std::string_view> no_limit_empty_ep_list = {};
ASSERT_STATUS_OK(TestGraphTransformer(
build_test_case, 13, *logger_,
std::make_unique<LayerNormFusion>(no_limit_empty_ep_list, TransformerLevel::Level2),
TransformerLevel::Level2, 1, nullptr, post_graph_checker));
}

// Same as above for the simplified pattern, which is the one hit by the model in issue #30513.
TEST_F(GraphTransformationTests, SimplifiedLayerNormFusion_ReduceMeanNoKeepDims_NoFusion) {
auto build_test_case = [](ModelTestBuilder& builder) {
auto* input = builder.MakeInput<float>({{4, 4}});
auto* pow_exponent = builder.MakeInitializer<float>({}, {2.0f});
auto* epsilon = builder.MakeInitializer<float>({}, {1e-5f});
auto* scale = builder.MakeInitializer<float>({4}, {1.0f, 1.0f, 1.0f, 1.0f});

auto* pow_out = builder.MakeIntermediate();
auto* reduce_mean_out = builder.MakeIntermediate();
auto* add_out = builder.MakeIntermediate();
auto* sqrt_out = builder.MakeIntermediate();
auto* div_out = builder.MakeIntermediate();
auto* output = builder.MakeOutput();

builder.AddNode("Pow", {input, pow_exponent}, {pow_out});
Node& reduce_mean_node = builder.AddNode("ReduceMean", {pow_out}, {reduce_mean_out});
reduce_mean_node.AddAttribute("axes", std::vector<int64_t>{-1});
reduce_mean_node.AddAttribute("keepdims", static_cast<int64_t>(0));
builder.AddNode("Add", {reduce_mean_out, epsilon}, {add_out});
builder.AddNode("Sqrt", {add_out}, {sqrt_out});
builder.AddNode("Div", {input, sqrt_out}, {div_out});
builder.AddNode("Mul", {div_out, scale}, {output});
};

auto post_graph_checker = [](Graph& graph) {
const auto op_to_count = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_to_count.find("SimplifiedLayerNormalization") == op_to_count.end());
const auto reduce_mean_it = op_to_count.find("ReduceMean");
TEST_RETURN_IF_NOT(reduce_mean_it != op_to_count.end() && reduce_mean_it->second == 1);
return Status::OK();
};

ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 17, *logger_,
std::make_unique<SimplifiedLayerNormFusion>(),
TransformerLevel::Level2, 1, nullptr, post_graph_checker));
}

TEST_F(GraphTransformationTests, SimplifiedLayerNormFusionTest) {
constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/layer_norm_t5.onnx";
std::shared_ptr<Model> p_model;
Expand Down