Memory efficient SDP yields wrong gradients
Labels: triaged, module: multi-headed-attention

### 🐛 Describe the bug

The gradients of samples for models with and without memory efficient SPD should be nearly identical, but are in practice often very different. I have a code example here: https://gist.github.com/lengstrom/fd091020d1cf8e22b2e0caa01c2e9255. `optimum` transforms the model to use memory efficient SDP -- the gradients go back to matching when I disable memory efficient SDP and enable math SDP.

Considering a GPTNeoX model, we fix a sample and record the gradient of the memory efficient SDP model and the gradient of the original model. We then measure the cosine similarity (a vector similarity metric, 0 = uncorrelated, 1 = perfectly correlated) and find that both:

(a) the cosine similarity of gradients between the ME-SDP and standard models is not 1.0 and furthermore
(b) on some parameter groups, this cosine similarity is very low. For example, see `gpt_neox.layers.0.input_layernorm.weight` - the similarity is `0.45` on this parameter group

You can see the logs that show this here (from the script above):

```
gpt_neox.embed_in.weight grad match: False Maxdiff: 2.2251620292663574, relativediff: nan, cosine=0.6628117561340332
gpt_neox.layers.0.input_layernorm.weight grad match: False Maxdiff: 0.23522385954856873, relativediff: 6.147019863128662, cosine=0.4535333514213562
gpt_neox.layers.0.input_layernorm.bias grad match: False Maxdiff: 0.0925818607211113, relativediff: 4.2182841300964355, cosine=0.6757722496986389
gpt_neox.layers.0.post_attention_layernorm.weight grad match: False Maxdiff: 0.09811977297067642, relativediff: 6.601757526397705, cosine=0.7666352987289429
gpt_neox.layers.0.post_attention_layernorm.bias grad match: False Maxdiff: 0.08014828711748123, relativediff: 3.4428935050964355, cosine=0.7390220165252686
gpt_neox.layers.0.attention.query_key_value.weight grad match: False Maxdiff: 1.9186370372772217, relativediff: 12.226367950439453, cosine=0.7150740623474121
gpt_neox.layers.0.attention.query_key_value.bias grad match: False Maxdiff: 0.20297080278396606, relativediff: inf, cosine=0.7175157070159912
gpt_neox.layers.0.attention.dense.weight grad match: False Maxdiff: 0.4924759864807129, relativediff: 11.39132308959961, cosine=0.6795455813407898
```

This code should perfectly reproduce with just pytorch (nightly), optimum (latest -- this library is unversioned?), datasets (2.11.1), and transformers (4.29.2) installed. The full output is commented in the gist. This issue was explored more in https://github.com/huggingface/optimum/issues/1091: it doesn't arise with the 160m Pythia model (only 70m).

### Versions

 Collecting environment information...
PyTorch version: 2.1.0.dev20230609+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.35

Python version: 3.10.11 (main, May 16 2023, 00:28:57) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-73-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.7.64
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA A100 80GB PCIe
GPU 2: NVIDIA A100 80GB PCIe
GPU 3: NVIDIA A100 80GB PCIe
GPU 4: NVIDIA A100 80GB PCIe
GPU 5: NVIDIA A100 80GB PCIe
GPU 6: NVIDIA A100 80GB PCIe
GPU 7: NVIDIA A100 80GB PCIe

Nvidia driver version: 515.43.04
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          128
On-line CPU(s) list:             0-127
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 7513 32-Core Processor
CPU family:                      25
Model:                           1
Thread(s) per core:              2
Core(s) per socket:              32
Socket(s):                       2
Stepping:                        1
Frequency boost:                 enabled
CPU max MHz:                     3681.6399
CPU min MHz:                     1500.0000
BogoMIPS:                        5189.74
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca
Virtualization:                  AMD-V
L1d cache:                       2 MiB (64 instances)
L1i cache:                       2 MiB (64 instances)
L2 cache:                        32 MiB (64 instances)
L3 cache:                        256 MiB (8 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-31,64-95
NUMA node1 CPU(s):               32-63,96-127
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] pytorch-triton==2.1.0+9820899b38
[pip3] torch==2.1.0.dev20230609+cu117
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] magma-cuda117             2.6.1                         1    pytorch
[conda] mkl                       2023.1.0         h6d00ec8_46342
[conda] mkl-include               2023.1.0         h06a4308_46342
[conda] numpy                     1.24.3                   pypi_0    pypi
[conda] pytorch-triton            2.1.0+9820899b38          pypi_0    pypi
[conda] torch                     2.1.0.dev20230609+cu117          pypi_0    pypi
[conda] torchvision               0.15.2                   pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi