【Bug已解决】Onnx Runtime Errors With LongT5 解决方案一、现象长什么样把 LongT5HuggingFace 的长文本 T5 变体使用局部全局稀疏注意力导出成 ONNX再用 ONNX Runtime 推理时常见两类错误RuntimeError: ONNX export failed: Couldnt export operator aten::roll ...或onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: Invalid rank for input input_ids (expected 2, got dynamic)又或者导出成功但 ORT 推理时报RuntimeError: cannot compute on unknown shape (dim_value -1 with no bound)LongT5 的特殊之处在于它的注意力除了标准自注意力还有local attention滑动窗口和transient global tokens实现里大量用到roll、unfold、动态长度的 padding/截断。这些在导出 ONNX 时容易踩两条线(1) 某些 aten 算子在所选 opset 下不被 ORT 支持(2) 长序列的动态长度seq_len不固定没有在dynamic_axes里正确声明导致 ORT 推理时 shape 解析失败。最迷惑的是PyTorch 原生model(...)跑得好好地一导出 ONNX 就炸——因为导出要做静态图追踪LongT5 里依赖运行时seq_len的控制流如if seq_len window才走 local attention在追踪阶段被固化或报错。二、背景ONNX 导出torch.onnx.export本质上是对模型做一次执行追踪trace把 Python 控制流变成静态计算图。LongT5 的几个结构对静态图不友好局部注意力窗口依赖运行时长度local_attention用unfold/roll构造窗口当seq_len在追踪时是某个固定值如 512导出图的窗口尺寸被写死推理时换长度就 shape 不匹配。动态分支LongT5 可能根据attention_type或序列长度选择不同注意力路径tracing 只走了一条另一条在 ORT 里不存在或错配。不支持的 aten 算子roll、unique、searchsorted等在某些 opset 下 ORT 不一定支持导出直接失败。dynamic_axes 漏配input_ids、attention_mask、labels的seq_len维若不声明为动态ORT 会把导出时的长度当成固定值。下面用可运行代码复现导出时遇到不支持的 op / 动态轴缺失导致 ORT 推理 shape 错误的机制。三、根因根因一句话LongT5 的局部/全局注意力依赖运行时序列长度和roll/unfold等算子ONNX 导出要么因不支持的 aten op 失败要么因dynamic_axes漏配把seq_len固化导致 ORT 推理时 shape 解析失败或结果错乱。三个具体失配aten 算子不支持roll/unfold在所选 opset 下 ORT 无实现导出失败。dynamic_axes 漏配seq_len维被写死换长度即 shape 错。控制流被 tracing 固化为单路径LongT5 的注意力分支选择被导出时定死与运行时不符。四、最小可运行复现用纯 Python 模拟导出时算子不支持 动态轴缺失导致 ORT 推理失败from dataclasses import dataclass from typing import Optional dataclass class FakeOnnxExport: supported_ops: set None def __post_init__(self): self.supported_ops self.supported_ops or {matmul, add, reshape} def export(self, ops_used: list, dynamic_axes: Optional[dict]): for op in ops_used: if op not in self.supported_ops: raise RuntimeError(fONNX export failed: unsupported operator {op}) if dynamic_axes is None or seq_len not in str(dynamic_axes): # 模拟动态轴缺失导致推理时 shape 未知 raise RuntimeError(cannot compute on unknown shape (dynamic_axes missing)) def main(): # LongT5 用到了 roll局部注意力窗口且 seq_len 动态 ops [matmul, add, roll] try: FakeOnnxExport().export(ops, dynamic_axesNone) except RuntimeError as e: print(复现到导出失败:, e) if __name__ __main__: main()运行会打印复现到导出失败: ONNX export failed: unsupported operator roll——对应 LongT5 局部注意力roll在导出时的典型失败。五、解决方案第一层最小直接修复最立竿见影的修复有两个方向按报错类型选若是算子不支持如roll把局部注意力里依赖roll的实现换成 ONNX 友好的等价写法如cat拼接实现环形移位或unfold后reshape避免 ORT 不支持的 op。若是动态轴缺失导出时务必声明所有可变长度的维为动态import torch def export_longt5(model, dummy_input, path): torch.onnx.export( model, dummy_input, path, input_names[input_ids, attention_mask], output_names[logits], dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, logits: {0: batch, 1: seq_len}, }, opset_version17, # 较高 opset 对 roll/unfold 支持更好 )第一层修复让seq_len成为真正的动态维且算子落在 ORT 支持范围内导出与推理都能过。六、解决方案第二层结构性改进把LongT5 ONNX 导出必须高 opset 完整 dynamic_axes ONNX 友好注意力收口成一个OnnxExporter配置固化这些约束避免每次手写出错。from dataclasses import dataclass, field from typing import Dict, List dataclass class OnnxExporter: opset_version: int 17 dynamic_axes: Dict[str, Dict[int, str]] field(default_factorydict) def with_longt5_defaults(self, seq_inputs(input_ids, attention_mask), seq_outputlogits): self.dynamic_axes { name: {0: batch, 1: seq_len} for name in seq_inputs } self.dynamic_axes[seq_output] {0: batch, 1: seq_len} return self def validate(self, ops_used: List[str], supported: set): missing [o for o in ops_used if o not in supported] if missing: raise RuntimeError( f以下算子 ORT 不支持需替换为 ONNX 友好实现: {missing} ) if not self.dynamic_axes: raise RuntimeError(dynamic_axes 未配置seq_len 将被固化) def main(): exp OnnxExporter().with_longt5_defaults() # 验证 ONNX 友好移除 roll用 cat/unfold 替代后算子集 ops_used [matmul, add, unfold, reshape] supported {matmul, add, unfold, reshape, cat} exp.validate(ops_used, supported) print(导出配置校验通过opset, exp.opset_version, dynamic_axes, exp.dynamic_axes) if __name__ __main__: main()第二层的关键是OnnxExporter把高 opset 完整 dynamic_axes 算子白名单固化导出前自动校验漏配直接报错。七、解决方案第三层断言 / CI 守护加 pytest 守护(1) 导出配置必须声明seq_len为动态(2) 使用的算子必须都在 ORT 支持集合内(3) 推理时换长度不报错动态轴生效。import pytest class Exporter: def __init__(self): self.dynamic_axes {} self.supported {matmul, add, unfold, reshape, cat} def configure(self, names): for n in names: self.dynamic_axes[n] {0: batch, 1: seq_len} def check(self, ops): assert seq_len in str(self.dynamic_axes), seq_len 必须动态 bad [o for o in ops if o not in self.supported] assert not bad, f不支持的算子: {bad} def test_dynamic_axes_required(): e Exporter() with pytest.raises(AssertionError): e.check([matmul]) # 未配置 dynamic_axes def test_supported_ops_only(): e Exporter() e.configure([input_ids, logits]) with pytest.raises(AssertionError): e.check([matmul, roll]) # roll 不支持 def test_valid_config_ok(): e Exporter() e.configure([input_ids, logits]) e.check([matmul, unfold, cat]) # 不抛 if __name__ __main__: pytest.main([__file__, -q])CI 里test_valid_config_ok通过就能保证 LongT5 的 ONNX 导出配置动态轴 算子白名单正确防止 ORT 推理时的 shape/算子错误回归。八、排查清单LongT5 ONNX Runtime 报错时按此顺序查看报错类型是unsupported operator算子问题还是unknown shape/dynamic_axes动态轴问题。算子不支持定位 LongT5 注意力里的roll/unfold/unique换成 ONNX 友好等价实现cat环形移位、unfoldreshape。opset 版本升到 17对roll/unfold支持更全。dynamic_axes 漏配input_ids/attention_mask/logits的seq_len维必须声明动态否则导出时长度被写死。控制流被 tracing 固化若 LongT5 依seq_len走不同注意力分支导出前固定attention_type或改写使两条路径都进入静态图。ORT 推理时显式给动态维session.run时输入的实际seq_len可不同于导出时dynamic_axes 正确即可。用 OnnxExporter 兜底导出前跑validate漏配动态轴/不支持算子直接报错。九、小结LongT5 在 ONNX Runtime 下的错误根因不在模型本身而在LongT5 的局部/全局注意力依赖运行时序列长度与roll/unfold等算子而 ONNX 导出tracing要么遇到 ORT 不支持的 aten op 直接失败要么因dynamic_axes漏配把seq_len固化导致 ORT 推理时 shape 解析失败或结果错乱。PyTorch 原生能跑是因为动态图允许运行时决定长度与分支而静态图导出把这些自由度收走了。修复三层第一层把roll换成 ONNX 友好写法、导出时声明完整dynamic_axesopset 17第二层用OnnxExporter把高 opset 动态轴 算子白名单固化并校验第三层用 pytest 断言seq_len 必须动态、算子必须受支持。记住LongT5 导出 ONNX先确认roll走不走、seq_len动不动——两个都搞定ORT 才不炸。