apache
[Bug] Error converting operator Attention
问题描述
在使用TVM 0.22.0版本导入ONNX模型（特别是使用opset 14的BGE模型）时，遇到了Attention操作符转换失败的问题。错误显示mask index形状不符合要求，即使已经尝试将所有mask相关输入设置为(batch_size, seq_len)的2D形状。
环境信息
- TVM版本：0.22.0
- ONNX模型：BGE模型，opset 14
- 操作系统：Linux
错误详情
完整错误堆栈
Error converting operator Attention, with inputs: [lv8, metadata["relax.expr.
Constant"][0]
# Metadata omitted. Use show_meta=True in script() method to show it., metadata
["relax.expr.Constant"][0]
# Metadata omitted. Use show_meta=True in script() method to show it., lv10]
模型处理错误: mask index should be in shape of (batch_size, seq_len),
or (batch_size, seq_len, seq_len)
Traceback (most recent call last):
File "/home/cyd/workspace/tvm/my_test/optimize_model_v3.py", line 85, in
load_and_optimize_onnx_model
relax_module, params = from_onnx(onnx_model, shape_dict)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 4235, in from_onnx
return g.from_onnx(graph, opset)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 3865, in from_onnx
self._construct_nodes(graph)
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 4046, in _construct_nodes
raise err
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 4041, in _construct_nodes
op = self._convert_operator(op_name, inputs, attr, self.opset)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 4141, in _convert_operator
sym = op_function(self.bb, inputs, attrs, [self._nodes, self._params])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/cyd/workspace/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py",
line 2109, in _impl_v1
assert mask_index_shape in (
^^^^^^^^^^^^^^^^^^^^^
AssertionError: mask index should be in shape of (batch_size, seq_len),
or (batch_size, seq_len, seq_len)
关键错误点
- 断言失败发生在 onnx_frontend.py 第2109行，错误消息为："mask index should be in shape of (batch_size, seq_len), or (batch_size, seq_len, seq_len)"
- 即使已经将所有mask相关输入设置为2D形状 (batch_size, seq_len)，仍然无法通过断言
复现步骤
- 准备一个使用opset 14的ONNX模型（如BGE embedding模型）
- 使用tvm.relax.frontend.onnx.from_onnx导入模型
- 为模型输入设置形状字典，包括将所有mask相关输入设置为2D形状：
shape_dict = { 'input_ids': (1, 512), 'attention_mask': (1, 512), 'dummy_mask_index': (1, 512), 'EmbedLayerNormalization_0_dummy_mask_index': (1, 512) } - 尝试转换模型时遇到断言错误
预期行为
TVM应该能够正确处理符合标准形状要求的mask索引输入，或者提供更清晰的错误信息和解决方案。
可能的解决方案
- 检查 _impl_v1 函数中的断言逻辑，确保它正确处理模型中传递的mask索引形状
- 考虑支持更多类型的mask索引形状，特别是对于opset 14及以上版本的ONNX模型
- 提供更明确的文档说明TVM支持的mask索引形状格式
附加信息
该问题在使用ONNX Runtime作为替代方案时不会出现，ONNX Runtime能够完全支持此类模型。
To better understand the error and reproduce the behavior on my side, could you please provide the exact inference script / code snippet you used when loading and running the ONNX model? If possible, a minimal reproducible example (model + script) would be ideal.
To better understand the error and reproduce the behavior on my side, could you please provide the exact inference script / code snippet you used when loading and running the ONNX model? If possible, a minimal reproducible example (model + script) would be ideal.
and the model, named bge-small. But it is too big to upload.
Gotcha, I will first attempt to reproduce this bug locally. Once I successfully reproduce the issue or make any new discoveries during the process, I will promptly update this issue.
@MoringLotus Add more information, hoping it will be of help.
convert.py and the download link for the model is: https://www.modelscope.cn/models/BAAI/bge-small-zh-v1.5/summary
Thanks for the extra information! It really clarified the issue for me and helped a lot.