Repository: pytorch
Issue: Optim.step() is significantly SLOW on MPS
Labels: module: performance,triaged,module: mps

### 🐛 Describe the bug

The first one is running on CPU. It makes sense that loss.back() is low on CPU. The second one is running on MPS, and loss.back() is faster than the first one, which makes sense too.

However, if we look at optim.step(), which should be a very simple step, it is way too slow, especially on MPS, which is nonsense.

```
import time

model = MLP()
optim = torch.optim.RAdam(model.parameters(), lr=1e-4)
loss_fun = nn.MSELoss()
loss_tr, loss_va = [], []

for epoch in range(1):
    model.train()
    for x, y in tr:
        
        optim.zero_grad()
        y_hat = model(x)
        loss = loss_fun(y[:,None], y_hat)
        start = time.time()
        loss.backward()
        print('Backward time: {:.4f}'.format(time.time() - start))
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=100)

        start = time.time()
        optim.step()
        print('step time: {:.4f}'.format(time.time() - start))
    loss_tr.append(loss.item())
```
Backward time: 0.0049
step time: 0.0020
Backward time: 0.0040
step time: 0.0011
Backward time: 0.0034
step time: 0.0010
Backward time: 0.0037
step time: 0.0011
Backward time: 0.0032
step time: 0.0010
Backward time: 0.0034
step time: 0.0022
Backward time: 0.0032
step time: 0.0019
Backward time: 0.0033
step time: 0.0020
Backward time: 0.0033
step time: 0.0022

```
#%%
model = MLP().to('mps')
optim = torch.optim.RAdam(model.parameters(), lr=1e-4)
loss_fun = nn.MSELoss()
loss_tr, loss_va = [], []

for epoch in range(1):
    model.train()
    for x, y in tr:
        
        optim.zero_grad()
        y_hat = model(x.to('mps'))
        loss = loss_fun(y[:,None].to('mps'), y_hat)
        start = time.time()
        loss.backward()
        print('Backward time: {:.4f}'.format(time.time() - start))
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=100)

        start = time.time()
        optim.step()
        print('step time: {:.4f}'.format(time.time() - start))
    loss_tr.append(loss.item())

```
Backward time: 0.0023
step time: 0.0069
Backward time: 0.0017
step time: 0.0062
Backward time: 0.0015
step time: 0.0055
Backward time: 0.0014
step time: 0.0049
Backward time: 0.0014
step time: 0.0047
Backward time: 0.0015
step time: 0.0092
Backward time: 0.0015
step time: 0.0085
Backward time: 0.0013
step time: 0.0083
Backward time: 0.0013
step time: 0.0083




### 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 13.0 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:37:07) [Clang 15.0.7 ] (64-bit runtime)
Python platform: macOS-13.0-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.24.3
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torchimize==0.0.14
[pip3] torchvision==0.16.0
[conda] numpy                     1.25.1                   pypi_0    pypi
[conda] numpy-base                1.24.3          py311h1d85a46_0  
[conda] torch                     2.1.0                    pypi_0    pypi
[conda] torchaudio                2.1.0                    pypi_0    pypi
[conda] torchimize                0.0.14                   pypi_0    pypi
[conda] torchvision               0.16.0                   pypi_0    pypi

cc @ezyang @gchanan @zou3519 @kadeng @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev