Repository: facebookresearch/xformers

Issue:
Shape mismatch between batches if lengths are different.
# 🐛 Bug

For two batches of shape `[B1, S1]` and `[B2, S2]`, if `S1 != S2` then a runtime error is raised:

```python
RuntimeError: The size of tensor a (S2) must match the size of tensor b (S1) at non-singleton dimension 2
```

## Command

## To Reproduce

With file `repro.py`:

```python
import torch
import torch.nn as nn

from xformers.factory.model_factory import xFormer, xFormerConfig


class MTModel(nn.Module):
    def __init__(self, transformer, out_embed_dim, vocab_size):
        super().__init__()
        self.transformer = transformer
        self.out_proj = nn.Linear(out_embed_dim, vocab_size, bias=False)

    def forward(self, src, tgt=None, src_mask=None, tgt_mask=None):
        return self.out_proj(self.transformer(src, tgt, src_mask, tgt_mask))


def build_model() -> MTModel:
    EMB = 384
    # Maximum allowable sentence in our dataset (in # subword tokens + 2 for 
, 
)
    BLOCK_SIZE = 1024
    VOCAB = 16_000

    encoder_configs = {
        "reversible": False,
        "block_type": "encoder",
        "dim_model": EMB,
        "position_encoding_config": {
            "name": "vocab",
            "seq_len": BLOCK_SIZE,
            "vocab_size": VOCAB,
            "dim_model": EMB,
        },
        "num_layers": 3,
        "multi_head_config": {
            "num_heads": 4,
            "residual_dropout": 0,
            "attention": {
                "name": "scaled_dot_product",
                "dropout": 0,
                "causal": True,
            },
            "dim_model": EMB,
        },
        "feedforward_config": {
            "name": "MLP",
            "dropout": 0,
            "activation": "relu",
            "hidden_layer_multiplier": 4,
            "dim_model": EMB,
        },
    }

    decoder_configs = {
        "block_type": "decoder",
        "dim_model": EMB,
        "position_encoding_config": {
            "name": "vocab",
            "seq_len": BLOCK_SIZE,
            "vocab_size": VOCAB,
            "dim_model": EMB,
        },
        "num_layers": 2,
        "multi_head_config_masked": {
            "num_heads": 4,
            "residual_dropout": 0,
            "dim_model": EMB,
            "attention": {
                "name": "scaled_dot_product",
                "dropout": 0,
                "causal": True,
            },
        },
        "multi_head_config_cross": {
            "num_heads": 4,
            "residual_dropout": 0,
            "dim_model": EMB,
            "attention": {
                "name": "scaled_dot_product",
                "dropout": 0,
                "causal": True,
            },
        },
        "feedforward_config": {
            "name": "MLP",
            "dropout": 0,
            "activation": "relu",
            "hidden_layer_multiplier": 4,
            "dim_model": EMB,
        },
    }

    
    my_config = [encoder_configs, decoder_configs]

    config = xFormerConfig(my_config)
    transformer = xFormer.from_config(config)

    return MTModel(transformer, EMB, VOCAB)


if __name__ == "__main__":
    model = build_model().to("cuda:0")

    src1 = torch.randint(0, 16000, (5, 7), device="cuda:0")
    src2 = torch.randint(0, 16000, (11, 13), device="cuda:0")

    # Also happens if src_mask is provided.
    model(src1)
    model(src2)
```

1. `python repro.py`


```log
python repro.py
A matching Triton is not available, some optimizations will not be enabled.
Error caught was: No module named 'triton'
Triton is not available, some optimizations will not be enabled.
This is just a warning: No module named 'triton'
Triton is not available, FusedMLP will not be enabled.
Either FairScale or torch distributed is not available, MixtureOfExperts will not be exposed. Please install them if you would like to use MoE
Traceback (most recent call last):
  File "/home/erip/Code/tmp/src/repro.py", line 107, in <module>
    model(src2)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/Code/tmp/src/repro.py", line 14, in forward
    return self.out_proj(self.transformer(src, tgt, src_mask, tgt_mask))
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/factory/model_factory.py", line 276, in forward
    memory = encoder(memory, input_mask=encoder_input_mask)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/factory/block_factory.py", line 231, in forward
    x = self.wrap_att(inputs=[q, k, v], att_mask=att_mask)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/residual.py", line 162, in forward
    x = self.sublayer(inputs=inputs, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/residual.py", line 89, in forward
    return residue + self.layer(*inputs, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/multi_head_dispatch.py", line 246, in forward
    y = self.attention(q, k, v, **kw_mask_args)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/attention/scaled_dot_product.py", line 130, in forward
    y = scaled_dot_product_attention(
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/attention/core.py", line 334, in scaled_dot_product_attention
    att = scaled_query_key_softmax(q, k, att_mask=att_mask)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/attention/core.py", line 232, in scaled_query_key_softmax
    att = _matmul_with_mask(q, k.transpose(-2, -1), mask)
  File "/home/erip/anaconda3/envs/tmp-dev/lib/python3.10/site-packages/xformers/components/attention/core.py", line 119, in _matmul_with_mask
    att += mask
RuntimeError: The size of tensor a (13) must match the size of tensor b (7) at non-singleton dimension 2
```

## Expected behavior

Mask should be independent batch shapes, especially if provided explicitly.

## Environment

Using `xFormers 0.0.16.dev432+git.bc08bbc`...

```
Collecting environment information...
PyTorch version: 1.13.1
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.79.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.13.1
[pip3] torchdata==0.5.1
[pip3] torchtext==0.14.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.5          py310hd5efca6_0  
[conda] numpy-base                1.23.5          py310h8e6c178_0  
[conda] pytorch                   1.13.1          py3.10_cuda11.7_cudnn8.5.0_0    pytorch
[conda] pytorch-cuda              11.7                 h67b0de4_1    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchdata                 0.5.1                    pypi_0    pypi
[conda] torchtext                 0.14.1                    py310    pytorch
```



Discussion:
It seems like this is a known issue, but I couldn't find a ticket explicitly for this. Apologies if I missed it!
cc @blefaudeux maybe?
It seems like the causal attention mask is tied to the sequence length [here](https://github.com/facebookresearch/xformers/blob/main/xformers/components/attention/scaled_dot_product.py#L57-L58) and [here](https://github.com/facebookresearch/xformers/blob/main/xformers/components/attention/scaled_dot_product.py#L101-L106) (the latter causing this issue). I'm not sure if there's a reason it isn't created on the fly (maybe a triton issue? maybe a general efficiency issue?).

I haven't done a deep look into the code, so maybe there's a reason the mask is a member attribute instead of a local, but the naive fix would be to drop the `self.` bit. 😄 

Resolving patch:
diff --git a/xformers/components/attention/scaled_dot_product.py b/xformers/components/attention/scaled_dot_product.py
--- a/xformers/components/attention/scaled_dot_product.py
+++ b/xformers/components/attention/scaled_dot_product.py
@@ -97,8 +97,9 @@ def forward(
             )
 
         # Handle a possibly deferred causal mask handling
+        mask = self.mask
         if self.causal and self.mask is None:
-            self.mask = AttentionMask.make_causal(
+            mask = AttentionMask.make_causal(
                 seq_len=q.shape[-2],
                 to_seq_len=q.shape[-2],
                 device=q.device,
@@ -106,10 +107,10 @@ def forward(
             )
 
         # Merge the optional causal mask and the user-provided mask
-        if self.mask is not None:
-            self.mask = self.mask.to(dtype=q.dtype, device=q.device)
+        if mask is not None:
+            mask = mask.to(dtype=q.dtype, device=q.device)
 
-            att_mask = att_mask + self.mask if att_mask is not None else self.mask
+            att_mask = att_mask + mask if att_mask is not None else mask
 
         # Try to handle a case where the sequence is smaller than the mask
         if (


Test patch:
diff --git a/tests/test_attentions.py b/tests/test_attentions.py
--- a/tests/test_attentions.py
+++ b/tests/test_attentions.py
@@ -204,6 +204,64 @@ def test_kqv_ordering(
     assert torch.allclose(res_false[0, :, :], res_false[1, :, :])
 
 
+@pytest.mark.parametrize("heads", [1, 4])
+@pytest.mark.parametrize("attention_name", ["scaled_dot_product"])
+@pytest.mark.parametrize("device", DEVICES)
+def test_different_seqlen(
+    attention_name: str,
+    heads: int,
+    device: torch.device,
+):
+    multi_head = _get_multihead(attention_name, 0.0, 0.0, False, heads, device)
+
+    # Check kqv are not flipped
+    # this will not catch all issues, but would catch a V being misplaced
+    # make k and q complimentary, so that QKt is all zero and attention is uniform
+
+    q = torch.cat(
+        (
+            torch.rand((1, MODEL // 2), device=device),
+            torch.zeros((1, MODEL // 2), device=device),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ, MODEL))
+
+    k = torch.cat(
+        (
+            torch.zeros((1, MODEL // 2), device=device),
+            torch.rand((1, MODEL // 2), device=device),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ, MODEL))
+    v = torch.rand(BATCH, SEQ, MODEL, device=device)
+
+    # Normal call
+    res = multi_head(query=q, key=k, value=v)
+
+    # Changing sequence length by dividing by two to simulate differing sequence length
+    q2 = torch.cat(
+        (
+            torch.rand((1, MODEL // 2), device=device),
+            torch.zeros((1, MODEL // 2), device=device),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ // 2, MODEL))
+
+    k2 = torch.cat(
+        (
+            torch.zeros((1, MODEL // 2), device=device),
+            torch.rand((1, MODEL // 2), device=device),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ // 2, MODEL))
+
+    v2 = torch.rand(BATCH, SEQ // 2, MODEL, device=device)
+
+    res2 = multi_head(query=q2, key=k2, value=v2)
+
+    assert res.shape != res2.shape
+
+
 @pytest.mark.parametrize("proj_bias", [False, True])
 @pytest.mark.parametrize("same_sizes", [False, True])
 @pytest.mark.parametrize("same_settings", [False, True])