Python implementation of scaled_dot_product_attention() yields different outcome than `aten::scaled_dot_product_attention`

prdvanderwal: <p>Hi everyone,</p>
<p>I am trying to customize the <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5617" rel="noopener nofollow ugc">scaled_dot_product_attention()</a> which is called in the <code>multi_head_attention_forward()</code> as part of the <code>MultiheadAttention</code> class. The original <code>scaled_dot_product_attention()</code> is part of torch._C._nn and (in my understanding) comes down to <code>aten::scaled_dot_product_attention</code> in onnx.</p>
<p>When I copy the source code of the <code>MultiheadAttention</code> class, <code>multi_head_attention_forward()</code> function, and use the Python code for <code>scaled_dot_product_attention()</code> (provided by PyTorch in a comment block), the training behavior of my ViT on ImageNet changes drastically. The snippet can be found below:</p>
<pre data-code-wrap="python"><code class="lang-python">def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
        is_causal=False, scale=None, enable_gqa=False) -&gt; torch.Tensor:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)
  
    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias = attn_mask + attn_bias
  
    if enable_gqa:
        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
  
    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value
</code></pre>
<p>However when I run my training code for the ViT, with the only difference being calling nn.MultiheadAttention() or the copied class with the snippet above for scaled_dot_product_attention(), the training behaviour changes completely. See screenshot below:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9aeee25fc8ad8fa2042155369249c3e5dd070c8.png" data-download-href="https://discuss.pytorch.org/uploads/default/f9aeee25fc8ad8fa2042155369249c3e5dd070c8" title="Training behaviour ViT changes with Python implementation"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f9aeee25fc8ad8fa2042155369249c3e5dd070c8_2_690x325.png" alt="Training behaviour ViT changes with Python implementation" data-base62-sha1="zCNFWvqw2bSnaFDczTmT8BfACz6" width="690" height="325" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f9aeee25fc8ad8fa2042155369249c3e5dd070c8_2_690x325.png, https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9aeee25fc8ad8fa2042155369249c3e5dd070c8.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9aeee25fc8ad8fa2042155369249c3e5dd070c8.png 2x" data-dominant-color="FAFBFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Training behaviour ViT changes with Python implementation</span><span class="informations">1022×482 29.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I want to emphasize that I did not do any customization yet and simply copied the commented source code above.</p>
<p>Am I correct to assume that this difference in training behavior can be attributed to using the Python implementation and that if I want to maintain good performance I have to use the optimized implementations? Or is the Python implementation supposed to give the same results but slower?</p>
<p>Subsequently, to customize the attention mechanism, do I have to clone PyTorch, change the .cpp files as discussed <a href="https://discuss.pytorch.org/t/i-cant-find-the-source-code-of-scaled-dot-product-attention/169389">here</a>, and build PyTorch from source?</p>
<p>I do not have any experience with this so if there is a PyTorch-based workaround I would love to find out. I looked at jit.trace but what I could find that yields a static computational graph which is fine for inference, but not ideal for training.</p>
<p>Thank you for any help in advance!</p>

Du_ni_kai: <p>torch.nn. functional.scaled_dot_product_attention This function is <strong>NOT</strong> equivalent to the definition in the document. That’s all I know</p>