Repository: facebookresearch/xformers

Issue:
NystromAttention sometimes results in NaN values when using key_padding_mask
# 🐛 Bug

NystromAttention sometimes results in NaN values when using key_padding_mask

## To Reproduce
Steps to reproduce the behavior:

```python
from xformers.components.attention import NystromAttention
from xformers.components.feedforward import MLP
from xformers.components import MultiHeadDispatch
from xformers.components.attention.utils import (bool_mask_to_additive)
import torch


input_tensor = torch.load(file_path).unsqueeze(0)
input_tensor[:, -3:] = 0
print('input shape: ', input_tensor.shape)
print('input: ', input_tensor, '\n')

torch.manual_seed(0)
num_heads = 8
proj_dropout = 0.2
bias = (False, False, False, True)
attn = NystromAttention(dropout = 0., 
                num_heads = num_heads, 
                num_landmarks = 256)

msa = MultiHeadDispatch(dim_model = 1024, 
                  num_heads = num_heads, 
                  attention = attn, 
                  bias = bias,
                  residual_dropout = proj_dropout, 
                  use_separate_proj_weight = True)

padding_mask = torch.zeros(1, input_tensor.size(1), dtype=torch.bool)
padding_mask[:, -3:] = True
key_padding_mask = bool_mask_to_additive(~padding_mask)
output = msa(input_tensor, key_padding_mask = key_padding_mask)
print('output with key_padding_mask: ', output)
print()

output = msa(input_tensor, key_padding_mask = None)
print('output without key_padding_mask: ', output)
```

Printout of results:
input shape:  torch.Size([1, 24798, 1024])
input:  tensor([[[0.0864, 0.0117, 0.0049,  ..., 0.0104, 0.0113, 0.0380],
         [0.0842, 0.0025, 0.0103,  ..., 0.0043, 0.0168, 0.0368],
         [0.1197, 0.0151, 0.0195,  ..., 0.0133, 0.0086, 0.0430],
         ...,
         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000]]]) 

output with key_padding_mask:  tensor([[[ 0.0269, -0.0089,  0.0000,  ..., -0.0080,  0.0022,  0.0241],
         [ 0.0269, -0.0089,  0.0468,  ..., -0.0080,  0.0022,  0.0241],
         [ 0.0269, -0.0000,  0.0468,  ..., -0.0080,  0.0022,  0.0241],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]]],
       grad_fn=\<MulBackward0\>)

output without key_padding_mask:  tensor([[[ 0.0000, -0.0089,  0.0468,  ..., -0.0080,  0.0022,  0.0000],
         [ 0.0269, -0.0089,  0.0468,  ..., -0.0080,  0.0022,  0.0241],
         [ 0.0269, -0.0000,  0.0000,  ..., -0.0000,  0.0022,  0.0241],
         ...,
         [ 0.0269, -0.0089,  0.0468,  ..., -0.0000,  0.0022,  0.0241],
         [ 0.0269, -0.0089,  0.0000,  ..., -0.0080,  0.0000,  0.0241],
         [ 0.0269, -0.0089,  0.0000,  ..., -0.0080,  0.0022,  0.0241]]],
       grad_fn=\<MulBackward0\>)

## Expected behavior
No nan values produced

## Environment

- xformers version: 0.0.11
- PyTorch Version (e.g., 1.0): 1.10.2
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): conda
- Python version: 3.8.12
- CUDA/cuDNN version: 11.1
- GPU models and configuration: NVIDIA 3090 (i reproduced the bug on cpu as well as seen in the above code snippet where no cuda device is used).


## Additional context
Happy to share the tensor I used by email or any other means (it's a .pt file)



Discussion:

Resolving patch:
diff --git a/xformers/components/attention/core.py b/xformers/components/attention/core.py
--- a/xformers/components/attention/core.py
+++ b/xformers/components/attention/core.py
@@ -91,7 +91,8 @@ def _matmul_with_mask(
         return torch.ops.xformers.matmul_with_mask(a, b, mask)
 
     # Non optimized codepath
-    assert not isinstance(mask, SparseCS)
+    if _is_sparse_available:
+        assert not isinstance(mask, SparseCS)
 
     att = a @ b
     if mask.dtype == torch.bool:
diff --git a/xformers/components/attention/nystrom.py b/xformers/components/attention/nystrom.py
--- a/xformers/components/attention/nystrom.py
+++ b/xformers/components/attention/nystrom.py
@@ -185,6 +185,7 @@ def forward(
                 logging.warning(
                     "Bool mask found, but an additive mask is expected. Converting but this is slow"
                 )
+
                 key_padding_mask = bool_mask_to_additive(key_padding_mask)
 
             if key_padding_mask.ndim == 2:
@@ -192,6 +193,20 @@ def forward(
                     key_padding_mask, batched_dim
                 )
 
+            zeros = torch.zeros_like(key_padding_mask)
+            ones = torch.ones_like(key_padding_mask)
+            is_masked = torch.isinf(-key_padding_mask)
+
+            # _mask takes 1 if the token is not padded, otherwise 0.
+            _mask = torch.where(is_masked, zeros, ones)
+            _mask = _mask.transpose(2, 1)
+            assert _mask.shape == (batched_dim, q.shape[1], 1)
+
+            # Mask q and k before pooling
+            # https://github.com/mlpen/Nystromformer/blob/main/code/attention_nystrom.py#L31
+            q = q * _mask
+            k = k * _mask
+
             assert key_padding_mask.size() == (batched_dim, 1, seq_len), (
                 f"key_padding_mask has invalid dimensions {key_padding_mask.size()}."
                 f" Must have dimensions {batched_dim, 1, seq_len} or (batch_size, {seq_len})."
@@ -227,22 +242,15 @@ def forward(
                     batched_dim, self.num_landmarks, seq_len, **tt
                 )
 
-            mask_1: Optional[torch.Tensor] = self.causal_mask_1
-            mask_2: Optional[torch.Tensor] = self.causal_mask_2
             mask_3: Optional[torch.Tensor] = self.causal_mask_3
             if key_padding_mask is not None:
-                mask_1 = (
-                    key_padding_mask.transpose(-2, -1)
-                    if mask_1 is None
-                    else mask_1 + key_padding_mask.transpose(-2, -1)
-                )
                 mask_3 = (
                     key_padding_mask if mask_3 is None else mask_3 + key_padding_mask
                 )
 
-            kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=mask_1)
+            kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=None)
             kernel_2 = scaled_query_key_softmax(
-                q=q_landmarks, k=k_landmarks, att_mask=mask_2
+                q=q_landmarks, k=k_landmarks, att_mask=None
             )
             kernel_3 = scaled_dot_product_attention(
                 q=q_landmarks, k=k, v=v, att_mask=mask_3


Test patch:
diff --git a/tests/test_nystrom_attention.py b/tests/test_nystrom_attention.py
--- a/tests/test_nystrom_attention.py
+++ b/tests/test_nystrom_attention.py
@@ -137,11 +137,6 @@ def test_masking():
         )
         r_nystrom = nystrom_attention(a, a, a, key_padding_mask=key_padding_mask)
         r_sdp = sdp_attention(a, a, a, att_mask=mask)
-        # account for when nan != nan
-        if r_nystrom.isnan().any() or r_sdp.isnan().any():
-            rand = random.uniform(0, 1)
-            r_nystrom = r_nystrom.masked_fill(r_nystrom.isnan(), rand)
-            r_sdp = r_sdp.masked_fill(r_sdp.isnan(), rand)
 
         # Not very close, but more so testing functionality.
         assert torch.allclose(