Repository: pytorch
Issue: MPS memory usage significantly higher than on CPU
Labels: module: memory usage,triaged,module: mps

### 🐛 Describe the bug

Thanks for the MPS backend! I've started to use it more but one issue I've been running into is that I OOM much more easily when using `device=mps` than `device=cpu`. I've included a fairly small example that reproduces this behavior:

```python
# bug.py
import argparse

import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalMultiheadAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_ctx):
        super().__init__()

        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_head = d_model // n_heads
        self.n_ctx = n_ctx

        self.linear_q = nn.Linear(d_model, n_heads * d_head)
        self.linear_k = nn.Linear(d_model, n_heads * d_head)
        self.linear_v = nn.Linear(d_model, n_heads * d_head)
        self.output = nn.Linear(n_heads * d_head, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.size()

        q = self.linear_q(x).view(B, T, self.n_heads, self.d_head)
        k = self.linear_k(x).view(B, T, self.n_heads, self.d_head)
        v = self.linear_v(x).view(B, T, self.n_heads, self.d_head)
        q = q.transpose(1, 2)  # [B, n_heads, T, d_head]
        k = k.transpose(1, 2).transpose(2, 3)  # [B, n_heads, d_head, T]
        v = v.transpose(1, 2)  # [B, n_heads, T, d_head]

        x = q @ k  # [B, n_heads, T, T]
        x.mul_(self.d_head**-0.5)
        x = F.softmax(x, dim=-1)
        x = x @ v  # [B, n_heads, T, d_head]

        x = x.transpose(1, 2).contiguous()  # [B, T, n_heads, d_head]
        x = x.view(B, T, self.n_heads * self.d_head)

        x = self.output(x)
        assert x.size() == (B, T, D)
        return x


d_model = 128
n_heads = 8
n_ctx = 2048
batch_size = 32

argparse = argparse.ArgumentParser()
argparse.add_argument("--device", type=str, default="cpu")
args = argparse.parse_args()
device = torch.device(args.device)

attn = CausalMultiheadAttention(d_model, n_heads, n_ctx).to(device)
for idx in range(2000):
    print(idx)
    ins = torch.randn(batch_size, n_ctx, d_model).to(device)
    out = attn(ins)
```

Running it on my M1 Max gives me roughly:

* `python bug.py --device cpu` has a peak memory usage of about 10GB
* `python bug.py --device mps` has a peak memory usage of about 17GB

(I'm actually unsure how to best report the memory size in use since MPS will use shared GPU memory; the above figures are from what Activity Monitor reports for the python process)

Beyond this toy example, I find that I consistently cannot train models on MPS that I can fit when training on CPU (with the exact same code, the only difference being the `device` used).

### Versions

Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.102)
CMake version: version 3.24.2
Libc version: N/A

Python version: 3.10.5 (main, Jul  9 2022, 10:15:34) [Clang 13.1.6 (clang-1316.0.21.2.5)] (64-bit runtime)
Python platform: macOS-12.6-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

Versions of relevant libraries:
[pip3] mypy==0.971
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] torch==1.12.1
[conda] Could not collect

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev