Repository: pytorch/pytorch

Issue:
nn.MultiheadAttention breaks for mask_type=2 when fast path is enabled (e.g. with batched attn_mask)
### 🐛 Describe the bug

**TLDR:** When `nn.MultiheadAttention` is used with a batched `attn_mask` which [should be](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html) shape (N*H, L, S) (where S=L for self-attn) **_and_** fast path is enabled it crashes.

_It works as expected when fast path is **not** enabled_

### Minimal example to reproduce:

```python
import torch
import torch.nn as nn


with torch.no_grad():
    B = 8
    L = 100
    D = 512
    H = 8

    mha = nn.MultiheadAttention(D, H, batch_first=True)

    X = torch.randn(B, L, D)
    M = torch.randn(B * H, L, L) > 0

    mha.train()  # disable fast path
    out, _ = mha(X, X, X, attn_mask=M, need_weights=False)  # works
    mha.eval()  # enable fast path
    out, _ = mha(X, X, X, attn_mask=M, need_weights=False)  # crashes
```

Gives:
```
RuntimeError: For mask_type == 2 mask shape should match input shape
```

I did a little debugging in the torch code around [here](https://github.com/pytorch/pytorch/blob/fe0afc58529bdf148bfda580a60e3d778b5b8f1d/torch/nn/modules/activation.py#L1176) and it seems that when `self.merge_masks` is called the `mask_type` returned is 0. However, when `torch._native_multi_head_attention` is called with `mask_type=0` but a (N*H, L, S) `attn_mask`, it seems to ignore the `mask_type` flag (see error above).

I'm on pytorch 2.0 (conda) on M2 Max macOS 13.2.1.

### Versions

Collecting environment information...
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.2.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.16 (main, Jan 11 2023, 10:02:19)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.2.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Max

Versions of relevant libraries:
[pip3] lion-pytorch==0.0.4
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==2.0.0
[pip3] pytorch-wavelets==1.3.0
[pip3] torch==2.0.0
[pip3] torch-yin==0.1.3
[pip3] torchaudio==2.0.0
[pip3] torchmetrics==0.11.1
[pip3] xitorch==0.4.0.dev0+3327cc0
[conda] lion-pytorch              0.0.4                    pypi_0    pypi
[conda] numpy                     1.23.5           py39h1398885_0  
[conda] numpy-base                1.23.5           py39h90707a3_0  
[conda] pytorch                   2.0.0                   py3.9_0    pytorch
[conda] pytorch-lightning         2.0.0                    pypi_0    pypi
[conda] pytorch-wavelets          1.3.0                    pypi_0    pypi
[conda] torch-yin                 0.1.3                    pypi_0    pypi
[conda] torchaudio                2.0.0                  py39_cpu    pytorch
[conda] torchmetrics              0.11.1                   pypi_0    pypi
[conda] xitorch                   0.4.0.dev0+3327cc0          pypi_0    pypi

cc @ezyang @gchanan @zou3519 @jbschlosser @bhosmer @cpuhrsch @erichan1


Discussion:
+1 having same issue
Not exactly the same, but I am also having issues with the fast path since upgrading to 2.0:
#97128 
cc @mikekgfb 
+1 having same issue
Marking this for high priority due to traffic to make sure it gets discussed during triage review. cc @jisaacso @drisspg @jbschlosser 
@erichan1 - Once a fix is out can you mark it for cherry pick as well? cc @atalman 
@erichan1  Please post a cherry pick by April 14, 5PM PST.

This issue is in the  milestones :[ https://github.com/pytorch/pytorch/milestone/36?closed=1](https://github.com/pytorch/pytorch/milestone/36?closed=1), if you want to see your fix included in this minor release. Please post it as a cherry-pick into the [[v2.0.1] Release Tracker](https://github.com/pytorch/pytorch/issues/97272). 

**The deadline is April 14, 5PM PST.**

**Only issues that have ‘cherry-picks’ will be considered for the release.**

Common FAQs:

Q1: Where can I find more information on the release process and terminology?

A: [pytorch/RELEASE.md at master · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/blob/master/RELEASE.md)

Q2: Am I guaranteed to be included in the cherry-pick if I do above?

A: No, it is not guaranteed, the Release Team will review all submissions against the listed criteria before making the final decision on what to include on 4/17.

Q3: When is 2.1 going to be released?

A: We do not have a formal date at this time but will update the community when we do. Our immediate focus is 2.0.1. Note that 1.12 was released on 6/28/22, 1.13 on 10/28/22 and 2.0 on 3/15/23.

Q4: **I missed the 4/14 5PM PST deadline, is there any option to have an extension?**

A: **No, in order to meet our 4/28 goal, we must hold 4/14 as our deadline and will not accept any requests after the fact. We are over communicating the timelines and process with the community to avoid such issues.**

Q5: Where should I double check to see if my issue is in the cherry pick tracker?

A: [[v2.0.1] Release Tracker · Issue #97272 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/issues/97272)

Q6: Where can I find the Release Compatibility Matrix for PyTorch?

A: [pytorch/RELEASE.md at master · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/blob/master/RELEASE.md#release-compatibility-matrix)

Please contact OSS Releng team members if you have any questions/comments. Again we appreciate everyone’s time and commitment to the community, PyTorch and 2.0 and 2.01 releases!


Please refer to this post for more details: https://dev-discuss.pytorch.org/t/pytorch-release-2-0-1-important-information/1176
Posted guidelines for a proper fix [here](https://github.com/pytorch/pytorch/pull/98375#issuecomment-1499504721) on the open PR.

Resolving patch:
diff --git a/torch/nn/modules/activation.py b/torch/nn/modules/activation.py
--- a/torch/nn/modules/activation.py
+++ b/torch/nn/modules/activation.py
@@ -1229,20 +1229,27 @@ def merge_masks(self, attn_mask: Optional[Tensor], key_padding_mask: Optional[Te
             check_other=False,
         )
 
-        if attn_mask is not None:
-            mask_type = 0
-            merged_mask = attn_mask
         if key_padding_mask is not None:
             mask_type = 1
             merged_mask = key_padding_mask
-        if (attn_mask is not None) and (key_padding_mask is not None):
+
+        if attn_mask is not None:
             # In this branch query can't be a nested tensor, so it has a shape
             batch_size, seq_len, _ = query.shape
             mask_type = 2
-            key_padding_mask_expanded = key_padding_mask.view(batch_size, 1, 1, seq_len) \
-                                                        .expand(-1, self.num_heads, -1, -1)
-            attn_mask_expanded = attn_mask.view(1, 1, seq_len, seq_len).expand(batch_size, self.num_heads, -1, -1)
-            merged_mask = attn_mask_expanded + key_padding_mask_expanded
+
+            # Always expands attn_mask to 4D
+            if attn_mask.dim() == 3:
+                attn_mask_expanded = attn_mask.view(batch_size, -1, seq_len, seq_len)
+            else:  # attn_mask.dim() == 2:
+                attn_mask_expanded = attn_mask.view(1, 1, seq_len, seq_len).expand(batch_size, self.num_heads, -1, -1)
+            merged_mask = attn_mask_expanded
+
+            if key_padding_mask is not None:
+                key_padding_mask_expanded = key_padding_mask.view(batch_size, 1, 1, seq_len).expand(-1, self.num_heads, -1, -1)
+                merged_mask = attn_mask_expanded + key_padding_mask_expanded
+
+        # no attn_mask and no key_padding_mask, returns None, None
         return merged_mask, mask_type
 
 


Test patch:
diff --git a/test/test_transformers.py b/test/test_transformers.py
--- a/test/test_transformers.py
+++ b/test/test_transformers.py
@@ -34,7 +34,6 @@
 if TEST_FAIRSEQ:
     import fairseq.models.transformer as fairseq_transformer
 
-
 @contextlib.contextmanager
 def use_deterministic_algorithims(mode: bool, warn_only: bool):
     r"""
@@ -155,6 +154,37 @@ def test_train_with_pad_and_catch_error(self, device):
                 l1_bool = nn.L1Loss()(test_train_bool[:, 0:2, :], test_eval_bool[:, 0:2, :]).item()
                 self.assertTrue(l1_bool < 1e-4, "Eval/Train difference in pad_mask BOOL")
 
+    @parametrize("device", device_list)
+    @parametrize("attn_mask_dim", [2, 3, None])
+    @parametrize("key_padding_mask_dim", [2, None])
+    def test_multiheadattention_fastpath_attn_mask(self, device, attn_mask_dim, key_padding_mask_dim):
+        with torch.no_grad():
+            B = 2
+            L = 4
+            D = 8
+            H = 4
+
+
+            if attn_mask_dim == 2:
+                attn_mask = torch.randn(L, L, device=device) > 0
+            elif attn_mask_dim == 3:
+                attn_mask = torch.randn(B * H, L, L, device=device) > 0
+            elif attn_mask_dim is None:
+                attn_mask = None
+
+            if key_padding_mask_dim == 2:
+                key_padding_mask = torch.randn(B, L, device=device) > 0
+            elif key_padding_mask_dim is None:
+                key_padding_mask = None
+
+            mha = nn.MultiheadAttention(D, H, batch_first=True, device=device)
+            X = torch.randn(B, L, D, device=device)
+
+            mha.train()  # disable fast path
+            out, _ = mha(X, X, X, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=False)
+            mha.eval()  # enable fast path
+            out, _ = mha(X, X, X, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=False)
+
     @parametrize("device", device_list)
     @parametrize("nhead", [1, 4, 8])
     def test_transformerencoderlayer_src_mask(self, device, nhead):