[Bug] Error converting operator LayerNormalization: InternalError: Op(relax.nn.layer_norm) requires the input gamma to have as many dimensions as the length of input axes.

[coffezhou]

Expected behavior

TVM should run the model correctly.

Actual behavior

For the following model,

it can be executed by onnxruntime and tensorrt, the results are as follows:

 [array([[[[ 0.956883  ,  0.34308553,  0.5195541 ,  1.9632151 ],
         [ 1.883044  , -0.23061049, -0.7294371 ,  2.5681047 ],
         [ 0.28407115,  1.682113  ,  2.5914793 ,  0.4852687 ],
         [ 0.24526536,  0.30622482,  3.0827932 ,  1.4561847 ]],

        [[ 2.1793444 , -2.5670562 , -0.80677223,  2.8984652 ],
         [ 0.08523583, -0.3017456 , -2.035876  ,  2.2447228 ],
         [-0.35428178,  2.6846695 ,  0.35822523,  2.3507211 ],
         [ 1.695939  ,  0.96472025,  0.18676078,  0.34729946]],

        [[ 2.6731548 ,  0.75374746, -0.6920886 ,  1.0468161 ],
         [ 3.229167  , -0.11532289, -2.1251893 ,  0.39142615],
         [-2.1620007 ,  2.5675638 ,  1.9419372 ,  4.1579857 ],
         [ 0.9184834 ,  1.6698728 ,  2.0392544 ,  0.292496  ]]]],
      dtype=float32)]

However, the onnx frontend of TVM cannot import it:

File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3925, in from_onnx
    return g.from_onnx(graph, opset)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3556, in from_onnx
    self._construct_nodes(graph)
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3736, in _construct_nodes
    raise err
  File "/home/carla/Documents/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 3733, in _construct_nodes
    op = self.bb.normalize(op)
         ^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/relax/block_builder.py", line 667, in normalize
    return _ffi_api.BlockBuilderNormalize(self, expr)  # type: ignore
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/ffi/cython/./function.pxi", line 228, in tvm.ffi.core.Function.__call__
tvm.error.InternalError: Op(relax.nn.layer_norm) requires the input gamma to have as many dimensions as the length of input axes. However, the given one has ndim 3, which is other than the length of axes 1

From the onnx specification of LayerNormalization, the attribute 'axis' is the first normalization dimension. In this issue, axis=1, which indicates that the LayerNormalization operator will normalize the last three dimension of the date. According to this understanding, the shape [3,4,4] of Scale should be correct.

Environment

OS: Ubuntu 20.04
TVM: 0.22.dev0 (c6969d7)
onnxruntime: 1.21.0

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment. As shown in the code, the model can be executed by onnxruntime. However, TVM cannot import this model.

import sys

import numpy as np
import onnx
import onnxruntime

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import pickle

            
def main():
    onnx_model = onnx.load("111.onnx")
    
    with open("inputs.pkl", "rb") as fp:
        inputs = pickle.load(fp)

    try:
        ort_session = onnxruntime.InferenceSession(
            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        sys.exit(1)
        
    print("ONNXRuntime:\n", ort_output)   

    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
    

    
if __name__ == "__main__":
    
    main()

testcase.zip

Triage

Please refer to the list of label tags here to find the relevant tags and add them below in a bullet format (example below).

• needs-triage

[ConvolutedDog]

Could you provide the model and tensor file so that we can reproduce it?

[coffezhou]

Could you provide the model and tensor file so that we can reproduce it?

Sorry! The attachment was failed to be uploaded. I have revised the issue and uploaded the attachment 'testcase.zip'.

[ConvolutedDog]

This result from that the backend C++ only support gamma and beta with the same shape as axes for now, and I will try to fix this next week. The corresponding checking code (op->arguments[1] is gamma and op->arguments[2] is beta):

tvm/src/relax/op/nn/nn.cc

Lines 340 to 347 in 27a314e

	} else if (input_sinfo[i]->ndim != n_axis) {
	ctx->ReportFatal(Diagnostic::Error(call)
	<< op << " requires the input " << op->arguments[i]->name
	<< " to have as many dimensions as the length of input axes. However, the "
	"given one has ndim "
	<< input_sinfo[i]->ndim << ", which is other than the length of axes "
	<< n_axis);
	}

[ConvolutedDog]

With the changes in PR #18143, the following script works as expected.

Click here for code

import os
import sys

import numpy as np
import onnx
import onnxruntime

import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import pickle


def main():
    current_dir = os.path.dirname(os.path.abspath(__file__))

    onnx_model = onnx.load(os.path.join(current_dir, "111.onnx"))

    with open(
        os.path.join(current_dir, "inputs.pkl"),
        "rb",
    ) as fp:
        inputs = pickle.load(fp)

    try:
        ort_session = onnxruntime.InferenceSession(
            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        sys.exit(1)

    print("inputs:\n", inputs)

    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
    tvm_model.show()

    tvm_model, params = relax.frontend.detach_params(tvm_model)

    input_list = [
        tvm.nd.array(inputs[key.name_hint])
        for key in tvm_model["main"].params
        if key.name_hint in inputs
    ]
    if params:
        input_list += params["main"]

    ex = relax.build(tvm_model, target="llvm")
    vm = relax.VirtualMachine(ex, tvm.cpu())

    nd_res = vm["main"](*input_list)

    print("ONNX output:\n", ort_output[0])
    print("TVM output:\n", nd_res)

    tvm.testing.assert_allclose(np.array(ort_output[0]), nd_res.numpy(), rtol=1e-5, atol=1e-5)



if __name__ == "__main__":
    main()