nn.MultiheadAttention throwing NaNs for entire batch

username: NumberChiffre
cooked: <p>Hey guys,</p>
<p>I’ve begun using torch’s latest MHA and noticed some differences, where by adding some NaNs as an input tensor for forward pass returns an output tensor full of NaNs. Using my default implementation, I would only get NaNs for the NaNs passed in the input tensor. Here’s how I reproduced this:</p>
<pre><code class="lang-python">from typing import Optional
import torch
import torch.nn as nn
from einops import rearrange

class MultiheadAttention(nn.Module):
    def __init__(self, dim: int, heads: int, dim_head: Optional[int] = None):
        super().__init__()
        self.dim_head = dim // heads if dim_head is None else dim_head
        self.heads = heads
        self.scale_factor = self.dim_head**-0.5
        self.to_qvk = nn.Linear(dim, self.dim_head * heads * 3, bias=False)
        self.to_out = nn.Linear(self.dim_head * heads, dim, bias=False)

    def forward(
        self, x: torch.Tensor, mask: torch.Tensor = None
    ) -&gt; torch.Tensor:
        qkv = self.to_qvk(x)
        q, k, v = tuple(
            rearrange(qkv, "b t (d h k) -&gt; k b h t d ", k=3, h=self.heads)
        )
        scaled_dot_prod = (
            torch.einsum("b h i d , b h j d -&gt; b h i j", q, k)
            * self.scale_factor
        )
        if mask is not None:
            assert mask.shape == scaled_dot_prod.shape[2:]
            scaled_dot_prod = scaled_dot_prod.masked_fill(mask, -torch.inf)
        attention = torch.softmax(scaled_dot_prod, dim=-1)
        context = torch.einsum("b h i j , b h j d -&gt; b h i d", attention, v)
        context = rearrange(context, "b h t d -&gt; b t (h d)")
        out = self.to_out(context)
        return out
</code></pre>
<p>The following would pass:</p>
<pre><code class="lang-python">size = 2
tokens = 2
dim = 8
heads = 2
inputs = torch.randn((size, tokens, dim), requires_grad=True).cuda()

model = MultiheadAttention(dim=dim, heads=heads).cuda()
outputs = model(inputs)
assert not torch.isnan(outputs).any()
assert not torch.isinf(outputs).any()

with_nans = torch.cat((inputs, torch.ones_like(inputs) * float("nan")))
outputs_with_nans = model(with_nans)

print(outputs_with_nans.shape)
print(torch.isnan(outputs_with_nans).sum().item())
print(outputs.shape[0] * outputs.shape[1] * outputs.shape[2])
assert torch.isnan(outputs_with_nans).sum() == outputs.shape[0] * outputs.shape[1] * outputs.shape[2]
</code></pre>
<p>This would fail, which is the problem:</p>
<pre><code class="lang-python">size = 2
tokens = 2
dim = 8
heads = 2
inputs = torch.randn((size, tokens, dim), requires_grad=True).cuda()

model = nn.MultiheadAttention(embed_dim=dim, num_heads=heads).cuda()

with torch.backends.cuda.sdp_kernel(enable_flash=True):
    outputs = model(inputs, inputs, inputs, need_weights=False)[0]
assert not torch.isnan(outputs).any()
assert not torch.isinf(outputs).any()

with_nans = torch.cat((inputs, torch.ones_like(inputs) * float("nan")))

with torch.backends.cuda.sdp_kernel(enable_flash=True):
    outputs_with_nans = model(with_nans, with_nans, with_nans, need_weights=False)[0]

print(outputs_with_nans.shape)
print(torch.isnan(outputs_with_nans).sum().item())
print(outputs.shape[0] * outputs.shape[1] * outputs.shape[2])
assert torch.isnan(outputs_with_nans).sum() == outputs.shape[0] * outputs.shape[1] * outputs.shape[2]
</code></pre>