Qiming Bao · Blog · 2026 年 11 月 11 日 · English / 中文

Swin 和 BEiT 为什么导不出 ONNX：同一个张量，两种注意力内核

Hugging Face 的 Swin（microsoft/swin-tiny-patch4-window7-224）和 BEiT（microsoft/beit-base-patch16-224）在使用默认的 SDPA 注意力时，用 torch.onnx.export(..., dynamo=True) 导出会失败，报错指向一个看起来毫无问题的 reshape。根本原因是：CPU 上的 scaled_dot_product_attention 会根据 attention mask 是否需要梯度，选择不同的内核，而不同内核输出的内存布局不一样。我把它报告成了 pytorch/pytorch#204360，并给出了一行的修复。

这个 bug 是我在批量导出 40 多个 Hugging Face 模型、逐个和 PyTorch 对比时发现的。所有语言模型都通过了，其中 10 个还测了动态形状。Swin 和 BEiT 是仅有的两个由 ONNX 导出器本身导致的失败。

现象

ConversionError: Failed to decompose the FX graph for ONNX compatibility. This is step 2/3 ...
ValueError: Cannot view a tensor with shape torch.Size([64, 49, 3, 32]) and strides (4704, 32, 1568, 1)
as a tensor with shape (64, 49, 96)!
While executing %view_9 = aten.view.default(%transpose_5, [64, 49, -1])

模型代码里写的是 attn_output.reshape(*input_shape, -1)。对非连续张量做 reshape 是合法的，它会自动复制一份。但在导出流程的某个环节，它被换成了 view，而 view 只有在内存布局允许时才合法。

两个简单的检查缩小了范围：改用 attn_implementation="eager" 就能正常导出，说明问题和 SDPA 有关；单独跑 torch.export.export(...).run_decompositions() 也没问题，说明问题出在 ONNX 导出器分解之后的步骤里。

找到第一个出错的节点

导出器先对图做分解，然后跑一个类型提升（type promotion）的 pass，这个 pass 会用 FakeTensor 把整张图重新执行一遍。我自己写了一个小的 torch.fx.Interpreter，把分解后的图重新执行，逐个节点比对实际输出的 stride 和 node.meta["val"] 里记录的 stride，在第一个不一致的地方停下：

阶段	SDPA 输出的 stride，形状 (64, 3, 49, 32)	`transpose(1, 2)` 之后
`torch.export` 追踪时	(4704, 1568, 32, 1)	非连续：保持 `reshape`
ONNX 导出器分解时	(4704, 32, 96, 1)	看起来连续：被降成 `view`
类型提升重新执行时	(4704, 1568, 32, 1)	非连续：`view` 报错

第一个不一致的节点就是 SDPA 本身，而且输入完全相同。因为 torchlib 实现了 aten.scaled_dot_product_attention，导出器在分解时会保留这个高层算子。分解时记录的输出是 CPU flash 内核的布局；类型提升 pass 重新执行它时，得到的却是 math 内核的布局。

是什么让内核变了

找触发条件最快的办法，是在一个最小模块上列一张对照表。这个模块只有三个 Linear 投影、一个带 mask 的 SDPA，再接 transpose(1, 2).reshape(B, S, -1)：

`attn_mask` 的来源	导出
参数本身，直接使用	成功
参数经过一个运算（`p.unsqueeze(0)`、`p + 0`）	失败
buffer 经过一个运算	成功
图的输入，经过或不经过运算	成功

参数和 buffer 的区别就在于 requires_grad。mask 需要梯度时，CPU 上的 SDPA 会退回 math 内核。类型提升 pass 重新执行时用的是占位符的 FakeTensor，参数在那里仍然是 requires_grad=True，所以由参数算出来的 mask 也需要梯度，内核就换了。Swin 和 BEiT 都会把相对位置偏置（一个参数）变成 attention mask。

import torch

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.q = torch.nn.Linear(96, 96)
        self.k = torch.nn.Linear(96, 96)
        self.v = torch.nn.Linear(96, 96)
        self.bias = torch.nn.Parameter(torch.randn(3, 49, 49))  # 比如相对位置偏置

    def forward(self, x):
        B, S, _ = x.shape
        q = self.q(x).view(B, S, 3, 32).transpose(1, 2)
        k = self.k(x).view(B, S, 3, 32).transpose(1, 2)
        v = self.v(x).view(B, S, 3, 32).transpose(1, 2)
        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=self.bias.unsqueeze(0))
        return out.transpose(1, 2).reshape(B, S, -1)

torch.onnx.export(M().eval(), (torch.randn(64, 49, 96),), dynamo=True)  # 失败

我第一次给出的修复是错的

我最初的补丁是用 torch.no_grad() 包住类型提升的重新执行。打上之后 Swin 就能导出了，我差点就把它写进 issue。后来拿最小模块核对这个说法，发现它还是失败。原因是 PyTorch 的一个细节：在 no_grad 下，对需要梯度的张量做 view 类操作（比如 unsqueeze），结果仍然是 requires_grad=True。Swin 的 mask 是先从偏置表里按下标取值得到的，这一步在 no_grad 下会产生一个不需要梯度的新张量，所以 no_grad 在 Swin 上碰巧有效。

真正有效的修法是：重新执行之前先把 FakeTensor 参数 detach()，让它和分解时一样看到 requires_grad=False 的输入：

# torch/onnx/_internal/fx/passes/type_promotion.py，InsertTypePromotion._run
with fake_mode, dispatcher_mode, fx_traceback.preserve_node_meta():
    fake_args = [arg.detach() if isinstance(arg, torch.Tensor) else arg for arg in fake_args]
    self.interpreter.run(*fake_args)

打上这个补丁后，最小模块、Swin 和 BEiT 都能导出。在 ONNX Runtime 上，Swin 和 eager 的误差在 8.7e-6 以内，BEiT 的相对误差在 3.3e-6 以内。我在本地能跑的 ONNX 导出器测试，打补丁前后结果完全一样。在修复合并之前，导出前先调用 model.requires_grad_(False) 就能绕开这个问题。

几点经验

用机械的方法找第一个分歧点。把图重新执行一遍、逐个节点比对元数据和实际输出，直接定位到了 SDPA 节点。在此之前，我对分解表和 opset 版本做了好几个错误的猜测。
用对照表二分触发条件。每次只改 mask 的一个属性（参数、buffer、输入、是否经过 view），就把“Swin 上会失败”变成了一个精确的条件。
每个结论都要拿最小复现来核对。一个在原始模型上有效的修复，仍然可能是错的。我的第一个补丁就是这样。
一次追踪记录下的元数据，是对下一次执行的承诺。当某个 pass 根据记录的 stride 把 reshape 降成 view 时，后面每一次重新执行都必须精确地复现这些 stride。

完整的分析和 BEiT 的确认都在 pytorch/pytorch#204360 里。PyTorch 的分类机器人已经给它打上了 module: onnx、module: sdpa、has workaround 和 release triage 标签。

排查这个 bug 和撰写本文时，我使用了 Claude（Anthropic）辅助。所有结果都经过我本人审阅，各项检查也由我亲自运行。

包启明，现居新西兰奥克兰的 AI 研究者，研究方向为大模型推理与模型部署。博士毕业于奥克兰大学 Strong AI Lab，参与 PyTorch ONNX 导出器、ONNX Runtime GenAI 和 Microsoft PyRIT 等开源项目的贡献。