Repository: pytorch
Issue: [MPS] Probable incorrect gradient for einsum
Labels: triaged,module: mps

### 🐛 Describe the bug

I was trying to implement a network with several multi-head linear layers, where the input tensor has a shape of [NumBatch, NumHead, DimInput], and the intermediate output of each layer should have a shape of [NumBatch, NumHead, DimHidden]. One simple implementation is use nn.ModuleList to stack multiple (NumHead) of Linear layers:

```python
self.nets = []
for _ in range(n_heads):
    net = [nn.Linear(dim_in, dim_hidden, bias=do_bias)]
    if act is not None:
        net.append(act)
    net += [nn.Linear(dim_hidden, dim_out, bias=do_bias)]
    self.nets.append(nn.Sequential(*net))
self.nets = nn.ModuleList(self.nets)
```

and in the forward function, just simply traverse the MLP in the ModuleList:

```python
hs = []
for i_head in range(self.n_heads):
    hs.append(self.nets[i_head](xs[:,i_head]))
hs = torch.stack(hs, dim=1)
```

Another possible way is using torch.einsum which is faster and easy to write:

```python
hs = torch.einsum("bij,ijk->bik", xs, self.weight) + self.bias
```

But the problem is when I use the latter one I find it hard to train on "mps", but is totally fine with "cpu". The reason is the gradient calculated by "mps" is probably wrong that:
```python
grad(naive_cpu)==grad(naive_mps)==grad(einsum_cpu)!=grad(einsum_mps)
```

To reproduce the problem:

```python
import collections
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim

class MultiHeadLinear(nn.Module):

    def __init__(self, n_groups, dim_in, dim_out, init_params=None, bias=True):
        super().__init__()
        self.has_bias = bias
        self.n_groups = n_groups

        if init_params is None:
            self.weight = nn.Parameter(torch.empty((n_groups, dim_in, dim_out)))
            if self.has_bias:
                self.bias = nn.Parameter(torch.empty((n_groups, dim_out)))
            self.reset_parameters()
        else:
            self.weight = nn.Parameter(data=init_params[0], requires_grad=True)
            if self.has_bias:
                self.bias = nn.Parameter(data=init_params[1], requires_grad=True)

    def reset_parameters(self):
        for i_group in range(self.n_groups):
            nn.init.kaiming_uniform_(self.weight[i_group], a=np.sqrt(5))
            if self.has_bias:
                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight[i_group])
                b = 1 / np.sqrt(fan_in) if fan_in>0 else 0
                nn.init.uniform_(self.bias[i_group], -b, b)

    def forward(self, xs):
        hs = torch.einsum("bij,ijk->bik", xs, self.weight) + self.bias[None]
        return hs

class Model_Naive(nn.Module):

    def __init__(self, dim_hidden, n_heads, do_bias=True, act=nn.Tanh()):
        super().__init__()
        self.dim_hidden, self.n_heads = dim_hidden, n_heads

        self.nets = []
        for _ in range(n_heads):
            net = [nn.Linear(3, dim_hidden, bias=do_bias)]
            if act is not None:
                net.append(act)
            net += [nn.Linear(dim_hidden, 1, bias=do_bias)]
            self.nets.append(nn.Sequential(*net))
        self.nets = nn.ModuleList(self.nets)
        self.act_out = nn.Sigmoid()

    def forward(self, xs):
        batch_size = xs.shape[0]
        assert xs.shape[1]==self.n_heads
        hs = []
        for i_head in range(self.n_heads):
            hs.append(self.nets[i_head](xs[:,i_head]))
        hs = torch.stack(hs, dim=1)
        hs = hs.reshape([batch_size, -1])
        return hs

class Model(nn.Module):

    def __init__(self, dim_hidden, n_heads, do_bias=True, act=nn.Tanh()):
        super().__init__()
        self.dim_hidden, self.n_heads = dim_hidden, n_heads

        layers = [MultiHeadLinear(n_heads, 3, dim_hidden)]
        if act is not None:
            layers.append(act)
        layers.append(MultiHeadLinear(n_heads, dim_hidden, 1))

        self.gls = nn.Sequential(*layers)
        # self.act_out = nn.Sigmoid()

    def forward(self, xs):
        batch_size = xs.shape[0]
        hs = self.gls(xs)
        hs = hs.reshape([batch_size, -1])
        return hs

for device in ["cpu", "mps"]:
    for act in [nn.ReLU(), nn.LeakyReLU(), nn.Tanh()]:
        print("device", device, "activation", act)
        dim_hidden = 4
        n_heads = 2
        lr = 5e-4

        my_model_n = Model_Naive(dim_hidden, n_heads, act=act).to(device)
        my_model = Model(dim_hidden, n_heads, act=act).to(device)
        optimizer_n = optim.Adam(list(my_model_n.parameters()), lr=lr)
        optimizer = optim.Adam(list(my_model.parameters()), lr=lr)
        criterion = nn.MSELoss()

        # Transfering the weights, make sure two networks have the same weights and bias
        _state_dict = {}
        for i in range(2):
            _state_dict[f"gls.{i * 2}.weight"] = []
            _state_dict[f"gls.{i * 2}.bias"] = []
        for k, v in my_model_n.named_parameters():
            gls_layer = int(k.split(".")[2])
            if "weight" in k:
                _state_dict[f"gls.{gls_layer}.weight"].append(v.data.T)
            else:
                _state_dict[f"gls.{gls_layer}.bias"].append(v.data)
        state_dict = {}
        for k,v in _state_dict.items():
            if len(v)>0:
                state_dict[k] = torch.stack(v)
        state_dict = collections.OrderedDict(state_dict)
        my_model.load_state_dict(state_dict)

        # Start training
        for step in range(2):
            xs = torch.randn(10, n_heads, 3).to(device)
            ys = torch.randn(10).to(device)
            hs_n, hs = my_model_n(xs), my_model(xs)
            hs_n, hs = torch.mean(hs_n, dim=1), torch.mean(hs, dim=1)
            print(f"    Step {step}: Output difference (should close to 0)", torch.sum(torch.abs(hs_n - hs)).item())

            loss_n = criterion(hs_n, ys)
            loss = criterion(hs, ys)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            optimizer_n.zero_grad()
            loss_n.backward()
            optimizer_n.step()

            print(f"    Step {step}: Loss: {np.round(loss.item(), 6)}, LossNaive: {np.round(loss_n.item(), 6)}")

            grads = [None, None]
            for k, v in my_model_n.named_parameters():
                if k=="nets.0.0.weight":
                    grads[0] = v.grad
                    break
            for k, v in my_model.named_parameters():
                if k=="gls.0.weight":
                    grads[1] = v.grad[0].T
                    break
            print(f"    Step {step}: Gradient difference (should close to 0)", torch.sum(torch.abs(grads[0] - grads[1])).item())
        print("-" * 40)
```
and the output is as follows:
```
device cpu activation ReLU()
    Step 0: Output difference (should close to 0) 1.1920928955078125e-07
    Step 0: Loss: 1.996575, LossNaive: 1.996575
    Step 0: Gradient difference (should close to 0) 0.0
    Step 1: Output difference (should close to 0) 1.1920928955078125e-07
    Step 1: Loss: 0.634939, LossNaive: 0.634939
    Step 1: Gradient difference (should close to 0) 1.3969838619232178e-09
----------------------------------------
device cpu activation LeakyReLU(negative_slope=0.01)
    Step 0: Output difference (should close to 0) 1.6391277313232422e-07
    Step 0: Loss: 1.444707, LossNaive: 1.444707
    Step 0: Gradient difference (should close to 0) 3.725290298461914e-08
    Step 1: Output difference (should close to 0) 1.0803341865539551e-07
    Step 1: Loss: 1.072882, LossNaive: 1.072882
    Step 1: Gradient difference (should close to 0) 0.0
----------------------------------------
device cpu activation Tanh()
    Step 0: Output difference (should close to 0) 1.1920928955078125e-07
    Step 0: Loss: 0.822748, LossNaive: 0.822748
    Step 0: Gradient difference (should close to 0) 4.912726581096649e-08
    Step 1: Output difference (should close to 0) 5.21540641784668e-08
    Step 1: Loss: 1.780892, LossNaive: 1.780892
    Step 1: Gradient difference (should close to 0) 1.862645149230957e-09
----------------------------------------
device mps activation ReLU()
    Step 0: Output difference (should close to 0) 0.0
    Step 0: Loss: 0.248852, LossNaive: 0.248852
    Step 0: Gradient difference (should close to 0) 0.0
    Step 1: Output difference (should close to 0) 0.0
    Step 1: Loss: 0.62069, LossNaive: 0.62069
    Step 1: Gradient difference (should close to 0) 0.0
----------------------------------------
device mps activation LeakyReLU(negative_slope=0.01)
    Step 0: Output difference (should close to 0) 1.3240678310394287
    Step 0: Loss: 0.621392, LossNaive: 0.723184
    Step 0: Gradient difference (should close to 0) 17.308197021484375
    Step 1: Output difference (should close to 0) 1.0551173686981201
    Step 1: Loss: 1.249558, LossNaive: 1.352456
    Step 1: Gradient difference (should close to 0) 17.392276763916016
----------------------------------------
device mps activation Tanh()
    Step 0: Output difference (should close to 0) 0.0
    Step 0: Loss: 1.248354, LossNaive: 1.248354
    Step 0: Gradient difference (should close to 0) 5.823102951049805
    Step 1: Output difference (should close to 0) 0.00257999449968338
    Step 1: Loss: 1.059988, LossNaive: 1.059785
    Step 1: Gradient difference (should close to 0) 25.54180145263672
----------------------------------------
```
One thing to notice is that it seems like different activation function gives different errors. ReLU works fine, LeakyReLU gives the wrong output at the first step, nn.Tanh gives the correct output at the first step but the wrong gradient, which gives rise to wrong output from the second step

### Versions

PyTorch version: 2.3.0.dev20240117
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.2.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.26.3
Libc version: N/A

Python version: 3.12.0 | packaged by Anaconda, Inc. | (main, Oct 2 2023, 12:22:05) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.2.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.2
[pip3] torch==2.3.0.dev20240117
[pip3] torchaudio==2.2.0.dev20240117
[pip3] torchvision==0.18.0.dev20240117
[conda] numpy 1.26.2 pypi_0 pypi
[conda] torch 2.3.0.dev20240117 pypi_0 pypi
[conda] torchaudio 2.2.0.dev20240117 pypi_0 pypi
[conda] torchvision 0.18.0.dev20240117 pypi_0 pypi

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr