Repository: pytorch/pytorch

Issue:
MPS backend produces bad training results in comparison to other backends
### 🐛 Describe the bug

Using the [MPS backend](https://pytorch.org/docs/stable/notes/mps.html) to train a model produces much worse results than using other backends (e.g. CPU or CUDA). To be clear, I am not talking about the speed of the training, but rather about the metrics for the quality (loss, perplexity) of the model after it has been trained. For example, if you run the training of the [word_language_model](https://github.com/pytorch/examples/tree/40289773aa4916fad0d50967917b3ae8aa534fd6/word_language_model) in the [pytorch/examples](https://github.com/pytorch/examples/tree/40289773aa4916fad0d50967917b3ae8aa534fd6) repository with either the CPU or the CUDA backend, your values for `loss` and `ppl` will be similar to this:

```
$ python main.py --cuda --epochs 1

| epoch   1 |   200/ 2983 batches | lr 20.00 | ms/batch 18.51 | loss  7.63 | ppl  2063.29
| epoch   1 |   400/ 2983 batches | lr 20.00 | ms/batch 17.47 | loss  6.86 | ppl   950.96
| epoch   1 |   600/ 2983 batches | lr 20.00 | ms/batch 17.47 | loss  6.48 | ppl   653.41
| epoch   1 |   800/ 2983 batches | lr 20.00 | ms/batch 17.46 | loss  6.29 | ppl   539.53
| epoch   1 |  1000/ 2983 batches | lr 20.00 | ms/batch 17.50 | loss  6.14 | ppl   465.56
| epoch   1 |  1200/ 2983 batches | lr 20.00 | ms/batch 17.54 | loss  6.07 | ppl   430.74
| epoch   1 |  1400/ 2983 batches | lr 20.00 | ms/batch 17.52 | loss  5.95 | ppl   384.86
| epoch   1 |  1600/ 2983 batches | lr 20.00 | ms/batch 17.55 | loss  5.96 | ppl   387.05
| epoch   1 |  1800/ 2983 batches | lr 20.00 | ms/batch 17.54 | loss  5.82 | ppl   337.38
| epoch   1 |  2000/ 2983 batches | lr 20.00 | ms/batch 17.52 | loss  5.80 | ppl   329.33
| epoch   1 |  2200/ 2983 batches | lr 20.00 | ms/batch 17.54 | loss  5.67 | ppl   289.19
| epoch   1 |  2400/ 2983 batches | lr 20.00 | ms/batch 17.53 | loss  5.67 | ppl   290.80
| epoch   1 |  2600/ 2983 batches | lr 20.00 | ms/batch 17.53 | loss  5.66 | ppl   285.86
| epoch   1 |  2800/ 2983 batches | lr 20.00 | ms/batch 17.54 | loss  5.55 | ppl   256.72
-----------------------------------------------------------------------------------------
| end of epoch   1 | time: 54.77s | valid loss  5.53 | valid ppl   252.11
-----------------------------------------------------------------------------------------
=========================================================================================
| End of training | test loss  5.44 | test ppl   230.14
=========================================================================================
```

Running the same training with MPS enabled consistently results in significantly worse values for `loss` and `ppl` (tested with pytorch `v1.13.1` and `v2.0.0.dev20230119`):

```
$ python main.py --mps --epochs 1

| epoch   1 |   200/ 2983 batches | lr 20.00 | ms/batch 115.26 | loss  7.99 | ppl  2959.03
| epoch   1 |   400/ 2983 batches | lr 20.00 | ms/batch 114.31 | loss  7.52 | ppl  1849.67
| epoch   1 |   600/ 2983 batches | lr 20.00 | ms/batch 114.39 | loss  7.38 | ppl  1603.87
| epoch   1 |   800/ 2983 batches | lr 20.00 | ms/batch 113.73 | loss  7.30 | ppl  1475.25
| epoch   1 |  1000/ 2983 batches | lr 20.00 | ms/batch 113.39 | loss  7.26 | ppl  1421.42
| epoch   1 |  1200/ 2983 batches | lr 20.00 | ms/batch 113.48 | loss  7.25 | ppl  1406.03
| epoch   1 |  1400/ 2983 batches | lr 20.00 | ms/batch 113.49 | loss  7.18 | ppl  1317.59
| epoch   1 |  1600/ 2983 batches | lr 20.00 | ms/batch 113.44 | loss  7.19 | ppl  1330.15
| epoch   1 |  1800/ 2983 batches | lr 20.00 | ms/batch 114.70 | loss  7.16 | ppl  1280.67
| epoch   1 |  2000/ 2983 batches | lr 20.00 | ms/batch 113.67 | loss  7.16 | ppl  1285.51
| epoch   1 |  2200/ 2983 batches | lr 20.00 | ms/batch 113.82 | loss  7.16 | ppl  1288.31
| epoch   1 |  2400/ 2983 batches | lr 20.00 | ms/batch 113.25 | loss  7.13 | ppl  1247.44
| epoch   1 |  2600/ 2983 batches | lr 20.00 | ms/batch 113.25 | loss  7.18 | ppl  1313.82
| epoch   1 |  2800/ 2983 batches | lr 20.00 | ms/batch 113.55 | loss  7.17 | ppl  1301.76
-----------------------------------------------------------------------------------------
| end of epoch   1 | time: 362.77s | valid loss  7.01 | valid ppl  1112.40
-----------------------------------------------------------------------------------------
=========================================================================================
| End of training | test loss  6.93 | test ppl  1019.79
=========================================================================================
```

Increasing the number of epochs will only increase the difference between the results of the two training configurations.

### Versions

**Latest:**

```
PyTorch version: 1.13.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.23.2
Libc version: N/A

Python version: 3.10.9 (main, Jan 11 2023, 09:18:18) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.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

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.13.1
[conda] numpy                     1.23.5          py310hb93e574_0
[conda] numpy-base                1.23.5          py310haf87e8b_0
[conda] pytorch                   1.13.1                 py3.10_0    pytorch
```

**Nightly:**

```
PyTorch version: 2.0.0.dev20230119
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.23.2
Libc version: N/A

Python version: 3.10.9 (main, Jan 11 2023, 09:18:18) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.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

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.0.dev20230119
[conda] numpy                     1.23.5          py310hb93e574_0
[conda] numpy-base                1.23.5          py310haf87e8b_0
[conda] pytorch                   2.0.0.dev20230119        py3.10_0    pytorch-nightly
```

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


Discussion:
Hello, the bad results come from a numerical correctness issue of LSTM on MPS backend. We're working on fixing this.

Resolving patch:
diff --git a/torchgen/api/python.py b/torchgen/api/python.py
--- a/torchgen/api/python.py
+++ b/torchgen/api/python.py
@@ -1109,6 +1109,7 @@ def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument:
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor>",
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor>",
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor>",
+    "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor>",
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor,int64_t>",
     "::std::tuple<at::Tensor,at::Tensor,double,int64_t>",
     "::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,int64_t>",


Test patch:
diff --git a/test/forward_backward_compatibility/check_forward_backward_compatibility.py b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
--- a/test/forward_backward_compatibility/check_forward_backward_compatibility.py
+++ b/test/forward_backward_compatibility/check_forward_backward_compatibility.py
@@ -269,7 +269,10 @@
     ("aten::dsplit.int", datetime.date(2022, 9, 1)),
     ("aten::hsplit.array", datetime.date(2022, 9, 1)),
     ("aten::hsplit.int", datetime.date(2022, 9, 1)),
-    ("aten::lstm_mps_backward.out", datetime.date(2022, 9, 1)),
+    ("aten::lstm_mps_backward.out", datetime.date(2023, 9, 1)),
+    ("aten::lstm_mps_backward", datetime.date(2023, 9, 1)),
+    ("aten::_lstm_mps.out", datetime.date(2023, 9, 1)),
+    ("aten::_lstm_mps", datetime.date(2023, 9, 1)),
     ("aten::miopen_rnn_backward.out", datetime.date(2022, 9, 1)),
     ("aten::quantize_per_tensor.tensors_out", datetime.date(2022, 9, 1)),
     ("aten::split", datetime.date(2022, 9, 1)),
diff --git a/test/test_mps.py b/test/test_mps.py
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -8926,64 +8926,91 @@ def test_cpu_indices(self, device="mps"):
 
 class TestRNNMPS(TestCaseMPS):
     def test_lstm_1(self, device="mps", dtype=torch.float32):
+        for layers in [1] if product_version < 13.0 else [1, 2, 5]:
+            torch.random.manual_seed(42)
+            rnn = nn.LSTM(7, 4, layers, device="cpu")
+            input = torch.randn(2, 3, 7, device="cpu")
+            hx = torch.randn(layers, 3, 4, device="cpu")
+            cx = torch.randn(layers, 3, 4, device="cpu")
+
+            cpu_output, (cpu_hn, cpu_cn) = rnn(input, (hx, cx))
+
+            rnn = rnn.to(device)
+            input = input.to(device)
+            hx = hx.to(device)
+            cx = cx.to(device)
+            output, (hn, cn) = rnn(input, (hx, cx))
+
+            self.assertEqual(cpu_output, output)
+            self.assertEqual(cpu_hn, hn)
+            self.assertEqual(cpu_cn, cn)
+
+            # test batch_first
+            rnn = nn.LSTM(7, 4, layers, device="cpu", batch_first=True)
+            input = torch.randn(3, 2, 7, device="cpu")
+            hx = torch.randn(layers, 3, 4, device="cpu")
+            cx = torch.randn(layers, 3, 4, device="cpu")
+            cpu_output, (cpu_hn, cpu_cn) = rnn(input, (hx, cx))
+
+            rnn = rnn.to(device)
+            input = input.to(device)
+            hx = hx.to(device)
+            cx = cx.to(device)
+            output, (hn, cn) = rnn(input, (hx, cx))
+
+            self.assertEqual(cpu_output, output)
+            self.assertEqual(cpu_hn, hn)
+            self.assertEqual(cpu_cn, cn)
+
+    def test_lstm_backward(self, device="mps", dtype=torch.float32):
+        for layers in [1] if product_version < 13.0 else [1, 2, 5]:
+            lstm = nn.LSTM(2, 4, layers)  # initialized globally for consistent parameters init
+            lstm.train()
+
+            def get_results(device, inp, hx, cx):
+                rnn = lstm.to(device)
+                inp, hx, cx = inp.to(device), hx.to(device), cx.to(device)
+
+                output, _ = rnn(inp, (hx, cx))
+                f = output.sum()
+
+                param_names, params = zip(*rnn.named_parameters())
+                param_grads = zip(param_names, torch.autograd.grad(f, params, retain_graph=True))
+
+                input_grad, hx_grad, cx_grad = torch.autograd.grad(f, [inp, hx, cx])
+                return output, param_grads, input_grad, hx_grad, cx_grad
+
+            inp = torch.randn((5, 3, 2), requires_grad=True, dtype=dtype, device=device)
+            hx = torch.randn((layers, 3, 4), requires_grad=True, dtype=dtype, device=device)
+            cx = torch.randn((layers, 3, 4), requires_grad=True, dtype=dtype, device=device)
+
+            cpu_output, cpu_weights_grad, cpu_input_grad, cpu_hx_grad, cpu_cx_grad = get_results("cpu", inp, hx, cx)
+            mps_output, mps_weights_grad, mps_input_grad, mps_hx_grad, mps_cx_grad = get_results(device, inp, hx, cx)
+
+            self.assertEqual(cpu_hx_grad, mps_hx_grad)
+            self.assertEqual(cpu_cx_grad, mps_cx_grad)
+            self.assertEqual(cpu_output, mps_output)
+            self.assertEqual(cpu_input_grad, mps_input_grad)
+            for (cpu_name, cpu_weight_grad), (mps_name, mps_weight_grad) in zip(cpu_weights_grad, mps_weights_grad):
+                self.assertEqual(cpu_weight_grad, mps_weight_grad, f"mismatch in cpu:{cpu_name} vs mps:{mps_name}")
+
+            # test batch_first backward
+            lstm = nn.LSTM(2, 4, layers, batch_first=True)
+            lstm.train()
+
+            hx = torch.randn((layers, 5, 4), requires_grad=True, dtype=dtype, device=device)
+            cx = torch.randn((layers, 5, 4), requires_grad=True, dtype=dtype, device=device)
+
+            cpu_output, cpu_weights_grad, cpu_input_grad, cpu_hx_grad, cpu_cx_grad = get_results("cpu", inp, hx, cx)
+            mps_output, mps_weights_grad, mps_input_grad, mps_hx_grad, mps_cx_grad = get_results(device, inp, hx, cx)
+
+            self.assertEqual(cpu_hx_grad, mps_hx_grad)
+            self.assertEqual(cpu_cx_grad, mps_cx_grad)
+            self.assertEqual(cpu_output, mps_output)
+            self.assertEqual(cpu_input_grad, mps_input_grad)
+            for (cpu_name, cpu_weight_grad), (mps_name, mps_weight_grad) in zip(cpu_weights_grad, mps_weights_grad):
+                self.assertEqual(cpu_weight_grad, mps_weight_grad, f"mismatch in cpu:{cpu_name} vs mps:{mps_name}")
 
-        rnn = nn.LSTM(1, 4, 2, device="cpu")
-        input = torch.randn(2, 3, 1, device="cpu")
-        hx = torch.zeros(2, 3, 4, device="cpu")
-        cx = torch.zeros(2, 3, 4, device="cpu")
-
-        cpu_output, (cpu_hn, cpu_cn) = rnn(input, (hx, cx))
-
-        rnn = rnn.to(device)
-        input = input.to(device)
-        hx = hx.to(device)
-        cx = cx.to(device)
-        output, (hn, cn) = rnn(input, (hx, cx))
-
-        self.assertEqual(cpu_output, output)
-        self.assertEqual(cpu_hn, hn)
-        self.assertEqual(cpu_cn, cn)
-
-        # test batch_first
-        rnn = nn.LSTM(1, 4, 2, device="cpu", batch_first=True)
-        input = torch.randn(3, 2, 1, device="cpu")
-        hx = torch.zeros(2, 3, 4, device="cpu")
-        cx = torch.zeros(2, 3, 4, device="cpu")
-        cpu_output, (cpu_hn, cpu_cn) = rnn(input, (hx, cx))
-
-        rnn = rnn.to(device)
-        input = input.to(device)
-        hx = hx.to(device)
-        cx = cx.to(device)
-        output, (hn, cn) = rnn(input, (hx, cx))
-
-        self.assertEqual(cpu_output, output)
-        self.assertEqual(cpu_hn, hn)
-        self.assertEqual(cpu_cn, cn)
-
-    @unittest.skipIf(True, "Backward of lstm returns wrong result")
-    def test_lstm_2(self, device="mps", dtype=torch.float32):
-        def get_results(device):
-            rnn = nn.LSTM(1, 4, 1, device=device)
-            inp = torch.randn(2, 3, 1, device=device, requires_grad=True)
-            hx = torch.zeros(1, 3, 4, device=device)
-            cx = torch.zeros(1, 3, 4, device=device)
-
-            output, _ = rnn(inp, (hx, cx))
-            output.sum().backward()
-
-            weight_grad = rnn.weight_ih_l0.grad.clone()
-            input_grad = inp.grad.clone()
-
-            return output, weight_grad, input_grad
-
-
-        cpu_output, cpu_weight_grad, cpu_input_grad = get_results("cpu")
-        mps_output, mps_weight_grad, mps_input_grad = get_results("mps")
-
-        self.assertEqual(cpu_output, mps_output)
-        self.assertEqual(cpu_input_grad, mps_input_grad)
-        self.assertEqual(cpu_weight_grad, mps_weight_grad)
 
     def test_RNN_cell_no_broadcasting(self):
         def test(cell_module, input, hx, input_size, hidden_size):