ReduceMean when a part of LayerNormalization is picked a newer version and fails Schema verification

[yuslepukhin]

Describe the issue

LayerNormalization is defined as a function that has the latest version 17.
LN invokes ReduceMean that has the latest version 18.
The behavior of RM has changed between versions of 17 and 18. Specifically, axes attribute is no longer recognized in version 18 and instead, axes are represented by an optional input.
LN invokes RM per version 17 specification. When LN is inlined RM retains axes attribute.
However, due to the fact that the imported opsets include only version 18 for ONNX, schema validation fails for RM.

To reproduce

import onnxruntime
import numpy as np
import onnx

model_name="stable_diffusion_text_encoder_dynamo"
model_path=model_name + ".onnx"

optimized_model_path=model_name + "_opt.onnx"

#model = onnx.load(model_path)

session_options = onnxruntime.SessionOptions()
session_options.log_severity_level = 4
session_options.optimized_model_filepath = optimized_model_path
session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
sess = onnxruntime.InferenceSession(
    model_path,
    sess_options=session_options,
    providers=["CPUExecutionProvider"],
    #providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

Urgency

No response

Platform

Windows

OS Version

Windows 11

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

9479ba5

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[gramalingam]

It would help to know where this LayerNorm is coming from, since there are the following 3 possibilities:

• It is using ONNX's LayerNorm function-op: this should be fine, since the function-body defined should work for both opset versions (see here).
• It is using ORT's LayerNorm custom op ... I believe that the onnxruntime's custom function-op should be updated (just like the ONNX one above) to handle different target opset-versions.
• Or, it is using pytorch-onnx-exporter's own definition of LayerNorm. If it is this, then we need to look at the definition of this function.

[yuslepukhin]

The function-op comes from the standard ONNX. I looked at the code, and yes, it is doing that. However, that code refers mostly on how to invoke ReduceMean, and not what operator version is going to be picked up after it is inlined. I think this is where the problem is. LN is version 17, so it constructs RM for version 17. However, ONNX schema verification is validating it against version 18 because imported opset for ONNX is v18.

P.S. this does not occur with CUDA EP because CUDA EP claims it has an implementation for LN, and it can run it.

[justinchuby]

Potentially related? #16438

[gramalingam]

Potentially related? #16438

Right, thanks for the link! I think that explains it.

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[github-actions]

This issue has been automatically closed due to inactivity. Please reactivate if further support is needed.