Repository: pytorch
Issue: nn.MultiheadAttention softmax inconsistent in training mode 
Labels: module: nn,triaged

### 🐛 Describe the bug

I was playing around with multiple head attention, and stumbled upon a weird interaction when having a key which has only one token:
```
import torch
from torch import nn
d=5
n=3
b=3
x=torch.normal(0,1,(b,n,d))
x_cls=torch.normal(0,1,(b,1,d))
attn = nn.MultiheadAttention(
            d,
            1,
            dropout=0.2,
            batch_first=True
        )
print(attn( x,x_cls, x_cls,)[1])
print(attn.eval()( x,x_cls, x_cls,)[1])
```
which outputs:
```
tensor([[[1.2500],
         [0.0000],
         [1.2500]],

        [[1.2500],
         [1.2500],
         [1.2500]],

        [[1.2500],
         [1.2500],
         [1.2500]]], grad_fn=<DivBackward0>)
tensor([[[1.],
         [1.],
         [1.]],

        [[1.],
         [1.],
         [1.]],

        [[1.],
         [1.],
         [1.]]], grad_fn=<DivBackward0>)
```
From the theory I would expect that both give one as I am calculating the softmax over a nx1 matrix rowwise - however dropout makes it deviate from that and I get weird values away from 1, which seem to be determined by the number tokens. I know this is not relevant in practice but maybe it should still be a sanity test that should be past

### Versions

```

Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 9.3.0
Clang version: 3.4.2 (tags/RELEASE_34/dot2-final)
CMake version: version 3.20.3
Libc version: glibc-2.17

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.66.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 2080 Ti
GPU 1: NVIDIA GeForce RTX 2080 Ti
GPU 2: NVIDIA GeForce RTX 2080 Ti
GPU 3: NVIDIA GeForce RTX 2080 Ti

Nvidia driver version: 515.48.07
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.0
[pip3] performer-pytorch==1.1.4
[pip3] pytorch-lightning==1.6.1
[pip3] torch==1.11.0
[pip3] torch-cluster==1.6.0
[pip3] torch-geometric==2.0.4
[pip3] torch-scatter==2.0.9
[pip3] torch-sparse==0.6.13
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.11.0
[pip3] torchinfo==1.7.0
[pip3] torchmetrics==0.7.3
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.12.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] libblas                   3.9.0            14_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            14_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            14_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] numpy                     1.21.0           py38h9894fe3_0    conda-forge
[conda] performer-pytorch         1.1.4                    pypi_0    pypi
[conda] pytorch                   1.11.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-lightning         1.6.1              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-cluster             1.6.0                    pypi_0    pypi
[conda] torch-geometric           2.0.4                    pypi_0    pypi
[conda] torch-scatter             2.0.9                    pypi_0    pypi
[conda] torch-sparse              0.6.13                   pypi_0    pypi
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.11.0               py38_cu113    pytorch
[conda] torchinfo                 1.7.0                    pypi_0    pypi
[conda] torchmetrics              0.7.3              pyhd8ed1ab_0    conda-forge
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchvision               0.12.0               py38_cu113    pytorch
```

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1 @drisspg