MultiHeadAttention export to ONNX fails when using torch.no_grad()

Isaac15B: <p>Hi!</p>
<p>I’m trying to export to ONNX a model that contains a MultiHeadAttention module. However, I’m running into the following error:</p>
<blockquote>
<p>torch.onnx.errors.UnsupportedOperatorError: Exporting the operator ‘aten::_native_multi_head_attention’ to ONNX opset version 17 is not supported.</p>
</blockquote>
<p>I created this small wrapper model to replicate the issue:</p>
<pre data-code-wrap="python"><code class="lang-python">class TestModel(torch.nn.Module):
    
    def __init__(self):
         super().__init__()
         self.self_attn = torch.nn.MultiheadAttention(
                embed_dim=1024,
                num_heads=8,
                dropout=True,
                batch_first=True,
        )
    
    def forward(self, hidden_states, attention_mask):
        x, _ = self.self_attn(
                    query=hidden_states,
                    key=hidden_states,
                    value=hidden_states,
                    key_padding_mask=attention_mask.bool(),
                    need_weights=False,
                )
        return x

model = TestModel().eval()
batch_size = 16
q = torch.randn((batch_size, 50, 1024))
mask = torch.zeros((batch_size, 50))

with torch.no_grad():
    torch.onnx.export(
        model,
        (q, mask),
        "test-model.onnx",
        input_names=["query", "mask"],
        output_names=["attn_output", "attn_output_weights"],
     )
</code></pre>
<p>According to this Github <a href="https://github.com/onnx/onnx/issues/5449" rel="noopener nofollow ugc">issue</a> it looks like the MHA ONNX operator is not implemented yet.</p>
<p>But the thing that I really don’t get is that <strong>this issue is happening only when I use the <code>torch.no_grad()</code></strong> <strong>context manager</strong>. When I don’t use it the export is a success, why?</p>
<p>I need to use this context manager because otherwise my real model is raising CUDA OOM issue when exporting to ONNX on GPU (more precisely during JIT graph creation in <code>torch.onnx.utils._create_jit_graph()</code> function).</p>
<p>Environment:</p>
<p>I’m using NVIDIA PyTorch NGC container <a href="http://nvcr.io/nvidia/pytorch:24.01-py3" rel="noopener nofollow ugc">nvcr.io/nvidia/pytorch:24.01-py3</a> with <code>torch.__version__ = '2.2.0a0+81ea7a4'</code></p>

JackTemaki: <p>I run into the same problem. My current assumption is that when no gradient is needed and certain conditions are met, PyTorch will use the “fast path” implementation, which creates the <code>aten::_native_multi_head_attention</code> node, which is not exportable.</p>
<p>Trying to force PyTorch to not use it was not successful so far… I hope this is fixed in the upcoming 2.3, if not I need to search more for solutions. Currently we stick to PyTorch 1.13 because there the export still works.</p>

ptrblck: <p>You might be able to disable the “fast path” by wrapping your forward pass into the <code>torch.backends.cuda.sdp_kernel</code> context as described <a href="https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html#explicit-dispatcher-control">here</a> allowing you to disable all optimizations.</p>

wcycqjy: <p>Hi, I encountered the same problem and could you explain it in detail? I try to use " torch.backends.cuda.sdp_kernel(*enable_flash=False , *enable_math=False , *enable_mem_efficient=False , *enable_cudnn=False  )", but it doesn’t work.</p>