pytorch - 💡(How to fix) Fix [ONNX Export] TransformerEncoderLayer export fails when MHA fastpath is enabled [2 comments, 2 participants]

🐛 Describe the bug

Problem

PyTorch's TransformerEncoderLayer fails ONNX export. The following script will produce errors.

import torch
from torch.nn import TransformerEncoderLayer
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder_layer = TransformerEncoderLayer(
            d_model=512,
            dim_feedforward=2048, 
            nhead=8, 
            norm_first=True,
            batch_first=True
        )

    def forward(self, x):
        return self.encoder_layer(x)

model = Model()
model.eval()
# Without the following line, export may fail
# torch.backends.mha.set_fastpath_enabled(False)
dummy_input = torch.randn(10, 32, 512)  # (S, N, E)
with torch.no_grad():
    torch.onnx.export(
        model,
        dummy_input,
        "encoder.onnx",
        opset_version=17
    )

Error message

UnsupportedOperatorError: Exporting the operator 'aten::_transformer_encoder_layer_fwd' to ONNX opset version 17 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

Best Practices

Try disable fastpath when export TransformerEncoderLayer to onnx following the below code

torch.backends.mha.set_fastpath_enabled(False)

Versions

Collecting environment information... PyTorch version: 2.5.0+cu124 Is debug build: False CUDA used to build PyTorch: 12.4 ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64) GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0 Clang version: Could not collect CMake version: version 3.16.3 Libc version: glibc-2.31

Python version: 3.12.9 (main, Mar 17 2025, 21:01:58) [Clang 20.1.0 ] (64-bit runtime) Python platform: Linux-5.11.0-27-generic-x86_64-with-glibc2.31 Is CUDA available: True CUDA runtime version: 12.0.76 CUDA_MODULE_LOADING set to: LAZY GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Nvidia driver version: 525.60.13 cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.7 /usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.7 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

The default behavior of TransformerEncoderLayer cannot support onnx export. Should we add more informations for user when export the onnx?

cc @justinchuby @titaiwangms