Repository: pytorch
Issue: mps training is slower than cpu
Labels: module: performance,triaged,module: mps

### 🐛 Describe the bug

```
import torch
import time
import torch.nn as nn
import torch.utils.data as Data
torch.manual_seed(10)

class ATT(nn.Module):
    def __init__(self, ):
        super().__init__()

        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=18, nhead=3, dim_feedforward=512, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(
            self.encoder_layer, num_layers=2)
        self.fc = nn.Sequential(
            nn.Linear(18,1),
            nn.ReLU())

    def forward(self, x):
        x = self.transformer_encoder(x)
        output = self.fc(x)
        return output
    

dtr = Data.TensorDataset(torch.rand(int(1e5), 18), torch.randint(0, 2, (int(1e5),)).float())
dva = Data.TensorDataset(torch.rand(int(1e5), 18), torch.randint(0, 2, (int(1e5),)).float())
tr = Data.DataLoader(dtr, batch_size=512, shuffle=True, num_workers=4)
va = Data.DataLoader(dva, batch_size=256*16, shuffle=False, num_workers=4)

st = time.time()
model = ATT()
optim = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fun = nn.MSELoss()
loss_tr, loss_va = [], []

for epoch in range(3):
    model.train()
    for x, y in tr:
        optim.zero_grad()
        y_hat = model(x[:,None])
        loss = loss_fun(y.squeeze(), y_hat.squeeze())
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=100)
        optim.step()
    loss_tr.append(loss.item())

    with torch.no_grad():
        model.eval()
        all_te_loss = 0
        for x, y in va:
            y_hat = model(x[:,None])
            l_va = loss_fun(y.squeeze(), y_hat.squeeze())
            all_te_loss += l_va*y_hat.shape[0]
        temp = all_te_loss.item()
        loss_va.append(temp)
print('cpu time used :', time.time() - st)


st = time.time()
model = ATT().to('mps')
optim = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fun = nn.MSELoss()
loss_tr, loss_va = [], []

for epoch in range(3):
    model.train()
    for x, y in tr:
        optim.zero_grad()
        y_hat = model(x.to('mps')[:,None])
        loss = loss_fun(y.to('mps').squeeze(), y_hat.squeeze())
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=100)
        optim.step()
    loss_tr.append(loss.item())

    with torch.no_grad():   
        model.eval()
        all_te_loss = 0
        for x, y in va:
            y_hat = model(x.to('mps')[:,None])
            l_va = loss_fun(y.to('mps').squeeze(), y_hat.squeeze())
            all_te_loss += l_va*y_hat.shape[0]
        temp = all_te_loss.item()
        loss_va.append(temp)
print('mps time used :', time.time() - st)
```
cpu time used : 14.817430973052979
mps time used : 16.454972743988037


if without training just the inference, the speed is OK:
cpu time used : 5.149087905883789
mps time used : 4.7350358963012695


### Versions

Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.5 (main, Sep 11 2023, 08:31:25) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.4.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] numpy==1.26.1
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torchvision==0.16.0
[conda] numpy                     1.26.1                   pypi_0    pypi
[conda] numpy-base                1.26.0          py311hfbfe69c_0  
[conda] torch                     2.1.0                    pypi_0    pypi
[conda] torchaudio                2.1.0                    pypi_0    pypi
[conda] torchvision               0.16.0                   pypi_0    pypi

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr