Training results from using MPS backend are poor compared to CPU and CUDA
Labels: needs reproduction, triaged, module: mps

### 🐛 Describe the bug

I have a Mac M1 GPU and I've been trying to replicate the results in [this google colab notebook](https://colab.research.google.com/drive/1_X7O2BkFLvqyCdZzDZvV2MB0aAvYALLC) on using a transformer-type architecture for time series forecasting. The notebook comes from [this repo](https://github.com/zhouhaoyi/Informer2020). I'm aware that certain issues regarding mps vs cpu and cuda have been raised in the past, such as [this issue using LSTMs on mps](https://github.com/pytorch/pytorch/issues/92615). This code does not utilize lstm and I'm having a hard time identifying the exact PyTorch method that is causing the problem.

In summary, when I run the training phase in the notebook above, I get bad results using the mps backend compared to my Mac M1 CPU as well as CUDA on google colab. (The speed between mps and cuda is a different issue). This is the training loss that I get:

```
Use GPU: mps
>>>>>>>start training : informer_ETTh1_ftM_sl96_ll48_pl24_dm512_nh8_el2_dl1_df2048_atprob_fc5_ebtimeF_dtTrue_mxTrue_exp_0>>>>>>>>>>>>>>>>>>>>>>>>>>
train 8521
val 2857
test 2857
	iters: 100, epoch: 1 | loss: 0.5028459
	speed: 0.2630s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 393.6544s
	iters: 200, epoch: 1 | loss: 0.5656900
	speed: 0.2554s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 356.7680s
Epoch: 1 cost time: 68.72129678726196
Epoch: 1, Steps: 266 | Train Loss: 0.5757653 Vali Loss: 1.2751399 Test Loss: 1.0501673
Validation loss decreased (inf --> 1.275140).  Saving model ...
Updating learning rate to 0.0001
	iters: 100, epoch: 2 | loss: 0.9338713
	speed: 0.6757s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 831.7784s
	iters: 200, epoch: 2 | loss: 0.9569200
	speed: 0.2548s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 288.1538s
Epoch: 2 cost time: 67.94761419296265
Epoch: 2, Steps: 266 | Train Loss: 0.9442326 Vali Loss: 1.4567285 Test Loss: 1.1148015
EarlyStopping counter: 1 out of 3
Updating learning rate to 5e-05
	iters: 100, epoch: 3 | loss: 0.9309775
	speed: 0.6728s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 649.2323s
	iters: 200, epoch: 3 | loss: 1.1810662
	speed: 0.2581s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 223.2210s
Epoch: 3 cost time: 68.36700391769409
Epoch: 3, Steps: 266 | Train Loss: 0.9918775 Vali Loss: 1.4783044 Test Loss: 1.1274033
EarlyStopping counter: 2 out of 3
Updating learning rate to 2.5e-05
	iters: 100, epoch: 4 | loss: 1.0169864
	speed: 0.6707s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 468.7928s
	iters: 200, epoch: 4 | loss: 0.8940969
	speed: 0.2543s[/iter](https://file+.vscode-resource.vscode-cdn.net/iter); left time: 152.3422s
Epoch: 4 cost time: 67.84642291069031
Epoch: 4, Steps: 266 | Train Loss: 0.9954335 Vali Loss: 1.4753704 Test Loss: 1.1228207
EarlyStopping counter: 3 out of 3
Early stopping
>>>>>>>testing : informer_ETTh1_ftM_sl96_ll48_pl24_dm512_nh8_el2_dl1_df2048_atprob_fc5_ebtimeF_dtTrue_mxTrue_exp_0<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
test 2857
test shape: (89, 32, 24, 7) (89, 32, 24, 7)
test shape: (2848, 24, 7) (2848, 24, 7)
mse:1.0500280857086182, mae:0.7747534513473511
```

In contrast, the loss obtained using cuda, as shown in the google colab notebook, is 
```
Use GPU: cuda:0
>>>>>>>start training : informer_ETTh1_ftM_sl96_ll48_pl24_dm512_nh8_el2_dl1_df2048_atprob_fc5_ebtimeF_dtTrue_exp_0>>>>>>>>>>>>>>>>>>>>>>>>>>
train 8521
val 2857
test 2857
	iters: 100, epoch: 1 | loss: 0.3484939
	speed: 0.0800s/iter; left time: 119.7995s
	iters: 200, epoch: 1 | loss: 0.3274963
	speed: 0.0773s/iter; left time: 108.0117s
Epoch: 1 cost time: 20.9348361492157
Epoch: 1, Steps: 266 | Train Loss: 0.3885468 Vali Loss: 0.6522534 Test Loss: 0.6147651
Validation loss decreased (inf --> 0.652253).  Saving model ...
Updating learning rate to 0.0001
	iters: 100, epoch: 2 | loss: 0.2812596
	speed: 0.1925s/iter; left time: 236.9607s
	iters: 200, epoch: 2 | loss: 0.2148246
	speed: 0.0797s/iter; left time: 90.1093s
Epoch: 2 cost time: 21.145679235458374
Epoch: 2, Steps: 266 | Train Loss: 0.2568903 Vali Loss: 0.6256742 Test Loss: 0.5904896
Validation loss decreased (0.652253 --> 0.625674).  Saving model ...
Updating learning rate to 5e-05
	iters: 100, epoch: 3 | loss: 0.2377600
	speed: 0.1962s/iter; left time: 189.2880s
	iters: 200, epoch: 3 | loss: 0.1902070
	speed: 0.0812s/iter; left time: 70.2002s
Epoch: 3 cost time: 21.539507389068604
Epoch: 3, Steps: 266 | Train Loss: 0.2063076 Vali Loss: 0.6280525 Test Loss: 0.5942183
EarlyStopping counter: 1 out of 3
Updating learning rate to 2.5e-05
	iters: 100, epoch: 4 | loss: 0.1939584
	speed: 0.1975s/iter; left time: 138.0229s
	iters: 200, epoch: 4 | loss: 0.1632166
	speed: 0.0813s/iter; left time: 48.7252s
Epoch: 4 cost time: 21.655611753463745
Epoch: 4, Steps: 266 | Train Loss: 0.1804204 Vali Loss: 0.6678267 Test Loss: 0.6165376
EarlyStopping counter: 2 out of 3
Updating learning rate to 1.25e-05
	iters: 100, epoch: 5 | loss: 0.1675630
	speed: 0.1995s/iter; left time: 86.3623s
	iters: 200, epoch: 5 | loss: 0.1715067
	speed: 0.0826s/iter; left time: 27.5224s
Epoch: 5 cost time: 21.931761741638184
Epoch: 5, Steps: 266 | Train Loss: 0.1663280 Vali Loss: 0.6778610 Test Loss: 0.6354805
EarlyStopping counter: 3 out of 3
Early stopping
>>>>>>>testing : informer_ETTh1_ftM_sl96_ll48_pl24_dm512_nh8_el2_dl1_df2048_atprob_fc5_ebtimeF_dtTrue_exp_0<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
test 2857
test shape: (89, 32, 24, 7) (89, 32, 24, 7)
test shape: (2848, 24, 7) (2848, 24, 7)
mse:0.5913315415382385, mae:0.5571437478065491
```
The mse and the mae using cuda are much lower than that of mps backend. The plots obtained from cuda approximate much better the actual time series data as well.

The GitHub code is mostly written for cuda but I used grep to search the repo and enabled the device to be 'mps' by adding the following code in the file `exp_basic.py`:

```
def _acquire_device(self):
        if self.args.use_gpu:
            if torch.cuda.is_available():
                os.environ["CUDA_VISIBLE_DEVICES"] = str(
                    self.args.gpu) if not self.args.use_multi_gpu else self.args.devices
                device = torch.device('cuda:{}'.format(self.args.gpu))
                print('Use GPU: cuda:{}'.format(self.args.gpu))
            elif getattr(torch,'has_mps',False):
                os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 
                device = "mps"
                print('Use GPU: mps')
```

I believe the error is in the training stage, perhaps gradient computation, because if I load the neural network weights obtained from training using the CPU, and I evaluate the neural network model with mps backend on the test data (without gradient computations), I can reproduce the output of using cpu on the test data.

I also used grep on the whole repository to find what methods from PyTorch were being utilized to identify the source of discrepancy.

Here's a list of methods I found: Conv1d, LayerNorm, Dropout, init.kaiming_normal_, Embedding, Parameter, Linear, BatchNorm1d, ELU, MaxPool1d, torch.softmax, MSELoss, functional.relu. 

I checked Github issues for any mps backend issues with these methods but I couldn't find any. 

Does anyone have suggestions how I can locate the method that's causing the difference between mps and cpu/cuda?

Thanks!


### Versions

python -m torch.utils.collect_env
Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.5.1 (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.4 (main, Jul  5 2023, 08:40:20) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.5.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 M1 Max

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2a0
[conda] numpy                     1.25.2          py311he598dae_0  
[conda] numpy-base                1.25.2          py311hfbfe69c_0  
[conda] pytorch                   2.0.1                  py3.11_0    pytorch
[conda] torchaudio                2.0.2                 py311_cpu    pytorch
[conda] torchvision               0.15.2          cpu_py311he74fb5d_0

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