用差分测试找 ONNX 导出器的 bug：哪些有效，哪些骗了我
PyTorch 的 ONNX 导出器会把每个 ATen 算子交给 onnxscript 的 torchlib 里对应的函数来翻译。这些函数都有单元测试，但测试大多只喂普通的随机浮点数。我想知道在边界情况下会发生什么：NaN、无穷大、带符号的零、整数和布尔输入、float16。于是我写了一个小的差分测试，让每个算子分别在 PyTorch eager 模式下运行，以及经过 torch.onnx.export(dynamo=True) 再用 ONNX Runtime 运行，然后对比两边的结果。
做了三轮，测了大约 300 个算子和参数的组合，找到了好几个真 bug，其中 4 个我已经提了修复 PR。它也产生了四类误报，而这才是这篇文章更有价值的部分。
测试脚本
每个用例是一个只调用一个算子的小 nn.Module。脚本先用 eager 模式跑一遍，再导出、用 ONNX Runtime 跑导出的模型，然后对比形状、dtype 和数值。浮点数按容差比较，并设置 equal_nan=True；整数和布尔值必须完全一致。
class M(torch.nn.Module):
def __init__(self, fn): super().__init__(); self.fn = fn
def forward(self, *args): return self.fn(*args)
expected = M(fn)(*inputs)
program = torch.onnx.export(M(fn), inputs, dynamo=True)
actual = program(*inputs) # 在 ONNX Runtime 上运行
compare(actual, expected) # 形状、dtype、数值
输入比算子更重要。每个算子都要跑这些输入：包含 NaN、±inf、±0 和极小值的浮点张量，int32、int64 和 uint8 张量，布尔张量，空张量，以及 float16、bfloat16 和 float64。大多数 bug 都出在这样的地方：PyTorch 的算子接受某种 dtype 或特殊值，而对应的 ONNX 算子不接受。
找到了什么
| 问题 | 后果 | 状态 |
|---|---|---|
ceil、floor、round、trunc、isnan、isfinite 的整数输入 | 导出非法的 ONNX 模型 | 我的修复：onnxscript#3178 |
bool 输入的 max(x, dim) 和 min(x, dim) | 导出非法的 ONNX 模型 | 我的修复：onnxscript#3179 |
heaviside(NaN, v) | 返回 1，PyTorch 返回 0 | 我的修复：onnxscript#3180 |
带 dtype= 的 softmax/log_softmax | 转换的是输出而不是输入：整数输入得到全 0 或非法模型 | 我的修复：onnxscript#3196 |
atan2 遇到带符号的零、无穷大和 NaN | 角度算错 | 已有人在修：onnxscript#3152 |
cross_entropy(label_smoothing=...) | 损失算错 | 已有人在修：onnxscript#3069 |
scatter_reduce(reduce="mean") | 直接覆盖，没有求平均 | 已有人在修：onnxscript#3094 |
大输入下的 special.erfcx | 输出 0 或 NaN | 已有人在修：onnxscript#3038 |
同样的思路用在模型级别上，导出 40 多个 Hugging Face 模型并对比输出，又找到一个：Swin 和 BEiT 在使用 SDPA 注意力时导不出来。这个问题我单独写了一篇文章。
真正的问题里有一半已经有人在修了。我学到的教训是：动手写代码之前，先把开着的和已关闭的 PR 都搜一遍。有一次，我把 atan2 的修复连同测试都写完了，才发现另一位贡献者几天前已经提了同样的修复。
它骗了我的四种方式
- 常量在 ONNX Runtime 看到之前就被折叠掉了。第二轮我在
forward里面构造输入。导出器的优化器在导出时就用 ONNX 参考实现把整段计算折叠成了常量，所以我测的有一部分其实是这个参考实现，而不是 ONNX Runtime。整行都是-inf的logsumexp就是这样得到了 NaN。ONNX Runtime 本身返回的是-inf，而参考实现在上游已经修好了（onnx#8441）。第三轮我把所有输入都改成了图的输入。 - 容差带来的假象。
logit(sigmoid(10 * x), eps=1e-6)显示有 0.05 的差异。换成精确的输入，导出结果和 PyTorch 完全一致：sigmoid在 float32 下会饱和到非常接近 1，而logit会放大最后几位的差异。每个不一致都要先用手挑的输入重新测一遍，再下结论。 - 运行时可能掩盖一张错误的图。float16 输入的
softmax(x, dtype=torch.float32)被导出成了Softmax(float16) -> Cast，但在 CPU 上数值却是对的。ONNX Runtime 的 CPU 后端没有 float16 的Softmax内核，会自己先转成 float32，所以这个 bug 只会在有 float16 内核的后端上显现。因此回归测试检查的是图结构，而且关掉了优化器。 - 不是每个报错都是 bug。第三轮的大多数失败都是还不支持的算子（
digamma、linalg_inv、bessel_i0、masked_select等），属于功能需求。还有一些是我自己的错误：在forward里用随机数生成器，本来就不可能被导出。
几点体会
- 一个简短的测试脚本，加上一份好的边界输入清单，就能很快找到真 bug。罕见的输入比罕见的算子更重要。
- 把数值作为图的输入传进去，每个发现都要手动复核，数字好得不正常时就去看导出的图。
- 写修复之前，先搜一遍已有的 issue 和 PR。
搭建测试脚本、排查这些问题和撰写本文时，我使用了 Claude（Anthropic）辅助。所有结果都经过我本人审阅，各项检查也由我亲自运行。
包启明，现居新西兰奥克兰的 AI 研究者，研究方向为大模型推理与模型部署。博士毕业于奥克兰大学 Strong AI Lab，参与 PyTorch ONNX 导出器、ONNX Runtime GenAI 和 Microsoft PyRIT 等开源项目的贡献。