Qiming Bao · Blog · 11 November 2026 · English / 中文

Why Swin and BEiT fail to export to ONNX: one tensor, two attention kernels

Hugging Face Swin (microsoft/swin-tiny-patch4-window7-224) and BEiT (microsoft/beit-base-patch16-224) cannot be exported with torch.onnx.export(..., dynamo=True) when they use their default SDPA attention. The export stops with an error that points at an innocent reshape. The cause is that CPU scaled_dot_product_attention picks a different kernel, with a different output layout, depending on whether the attention mask requires grad. I reported it as pytorch/pytorch#204360 with a one-line fix.

I found it while exporting more than 40 Hugging Face models and comparing each against PyTorch. Every language model passed, ten of them also with dynamic shapes. Swin and BEiT were the only failures that came from the ONNX exporter itself.

The symptom

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])

The model code says attn_output.reshape(*input_shape, -1). A reshape of a non-contiguous tensor is legal: it copies. Somewhere in the export pipeline it had become a view, which is only legal when the memory layout allows it.

Two quick checks narrowed it down. With attn_implementation="eager" the export works, so SDPA is involved. And torch.export.export(...).run_decompositions() on its own also works, so the problem is in what the ONNX exporter does after decomposition.

Finding the first wrong node

The exporter decomposes the graph, then runs a type-promotion pass that re-executes the whole graph with fake tensors. I re-ran the decomposed graph myself with a small torch.fx.Interpreter that compared each node's output strides with the strides recorded in node.meta["val"], and stopped at the first mismatch:

The first mismatch was the SDPA node itself, with identical inputs. The exporter keeps aten.scaled_dot_product_attention undecomposed because torchlib implements it. During decomposition its recorded output had the CPU flash kernel's layout; when the type-promotion pass re-ran it, it got the math kernel's layout.

What makes the kernel change

The quickest way to find the trigger was a small table of variants on a minimal module: three Linear projections, SDPA with a mask, then transpose(1, 2).reshape(B, S, -1):

The difference between a parameter and a buffer is requires_grad. CPU SDPA falls back to the math kernel when the mask requires grad. The type-promotion pass re-runs the graph with the placeholders' fake values, and parameters keep requires_grad=True there, so a mask computed from a parameter requires grad and the kernel switches. Swin and BEiT both turn a relative position bias, which is a parameter, into the 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))  # e.g. relative position bias

    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)  # fails

The fix I got wrong first

My first patch wrapped the type-promotion re-run in torch.no_grad(). Swin then exported, and I nearly wrote it up. Checking the claim against the minimal module showed it still failed. The reason is a PyTorch detail: under no_grad, a view of a tensor that requires grad, such as unsqueeze, still has requires_grad=True. Swin's mask starts with an indexing lookup into its bias table, which under no_grad produces a new tensor that does not require grad, so no_grad happened to work there.

The fix that works is to detach the fake arguments before the re-run, so it sees the same requires_grad=False inputs as the decomposition:

# 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)

With it, the minimal module, Swin and BEiT all export. On ONNX Runtime, Swin matches eager within 8.7e-6 and BEiT within 3.3e-6 relative. The ONNX exporter tests I could run locally have the same results with and without the change. Until a fix lands, model.requires_grad_(False) before exporting avoids the problem.

Lessons

1. Find the first divergence mechanically. Re-running the graph and comparing each node's metadata with what it actually produces pointed straight at the SDPA node, after several wrong guesses about decompositions and opsets.
2. Bisect the trigger with a table of variants. Changing one property of the mask at a time (parameter, buffer, input, view or not) turned "it fails on Swin" into one precise condition.
3. Check every claim against the minimal repro. A fix that works on the model you started with can still be wrong. My first patch was.
4. Metadata from one tracing pass is a promise about the next. When a pass lowers reshape to view based on recorded strides, every later re-run has to reproduce those strides exactly.

The issue, pytorch/pytorch#204360, has the full analysis and the BEiT confirmation. The PyTorch triage bot labelled it module: onnx, module: sdpa, has workaround and release triage.

I used Claude (Anthropic) to help investigate this bug and draft this post. I reviewed the results and ran the checks myself.

Qiming Bao is an AI researcher in Auckland, New Zealand, working on LLM reasoning and model deployment. He received his PhD from the University of Auckland's Strong AI Lab and contributes to the PyTorch ONNX exporter, ONNX Runtime GenAI and Microsoft PyRIT.