Scaled_dot_product_attention is not numerically stable

ElToto: <p>I am unable to make the <code>scaled_dot_product_attention</code> function numerically stable, so that it<br>
returns the same values as a manual implementation of attention.<br>
I provide a colab snippet that shows the problem.</p>
<p>One time i calculate attention manually and one time I use <code>scaled_dot_product_attention</code>.</p>
<p>Even with</p>
<pre><code class="lang-auto">  with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=False):
</code></pre>
<p>the values don’t match.</p>
<p>A little bit of floating point error is to be expected, but as you can see, I get a difference of 3e-07 which is larger than the standard 1e-07.<br>
Is this much deviation expected?</p>
<p>Mismatched elements: 76893 / 172032 (44.7%)<br>
Max absolute difference: 2.9802322e-07<br>
Max relative difference: 5.89578e-07</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/17stXOoKF5bdEdEvjDJK6oFEPJXMW3ZGT?usp=sharing">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/a1a808af752ea90e8a01d019d94384c8/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/17stXOoKF5bdEdEvjDJK6oFEPJXMW3ZGT?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="260" height="260">

<h3><a href="https://colab.research.google.com/drive/17stXOoKF5bdEdEvjDJK6oFEPJXMW3ZGT?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">import torch
import numpy as np

# should be unneccessary, not running on GPU
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.deterministic =  True


# compares if two tensors are equal
def assert_equal(output1, output2):
  a1, a2 = output1.detach().numpy(), output2.detach().numpy()
  np.testing.assert_allclose(a1, a2, rtol=1e-07)

# regular values for Vision Transformer
head_dim = 64
no_heads = 12
no_input_vectors = 14*16

#random input in shape batchsize, number input vectors, embedding dimension
input = torch.rand(1, no_input_vectors, head_dim*no_heads)

# First, calculate attention manually

x_manual = torch.clone(input)
x_manual = x_manual.reshape(1, no_input_vectors, no_heads, head_dim).permute(0, 2, 1, 3)
q, k, v = x_manual, x_manual, x_manual
k_t = k.transpose(-2, -1)
s_dot_product = ( q @ k_t) * (head_dim**-0.5)
attn = s_dot_product.softmax(dim=-1)
x_manual = attn @ v



## Now do it with the scaled_dot_product function

x_auto = torch.clone(input)
# Disable all "improvements"
# Also doesn't work with enable_math=True
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=False):
  x_auto = x_auto.reshape(1, no_input_vectors, no_heads, head_dim).permute(0, 2, 1, 3)
  x_auto = torch.nn.functional.scaled_dot_product_attention(x_auto,x_auto,x_auto)


assert assert_equal(x_manual, x_auto)
</code></pre>

ptrblck: <aside class="quote no-group" data-username="ElToto" data-post="1" data-topic="182582">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/e/b5a626/48.png" class="avatar"> ElToto:</div>
<blockquote>
<p>A little bit of floating point error is to be expected, but as you can see, I get a difference of 3e-07 which is larger than the standard 1e-07.</p>
</blockquote>
</aside>
<p>Yes, your error is still in the expected range. Could you explain where the hard limit of <code>1e-7</code> comes from?</p>

ElToto: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="182582">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Yes, your error is still in the expected range. Could you explain where the hard limit of <code>1e-7</code> comes from?</p>
</blockquote>
</aside>
<p>Thanks for the reply, good question</p>
<p><code>np.testing.assert_allclose</code> uses as default value 1e-07 so I assumed, this is the accuracy we want to achieve if two calculations should be identical. But this might be wrong.<br>
I was using this value to test if a custom transformer implementation is identically to a timm.VisionTransformer but now, timm switched to  <code>F.scaled_dot_product_attention</code> so my test is failing.<br>
I guess i can set the value to 1e-06.</p>

ptrblck: <p>Thanks for the background. Yes, I think default values make sense, but I would not see them as a hard limiting factor. E.g. even accumulating <code>100 * 100</code> elements on the CPU in different orders already yields a higher error:</p>
<pre><code class="lang-python">x = torch.randn(100, 100)
s1 = x.sum()
s2 = x.sum(0).sum(0)

print((s1 - s2).abs().max())
# tensor(4.3869e-05)
print((s1 - s2).abs()/s2.abs())
# tensor(1.4508e-06)
</code></pre>