GPT-2 Model Logits and Loss are different on MPS
Author: Conzel
Pull request: no
Labels: name: mps

### System Info

- `transformers` version: 4.37.2
- Platform: macOS-14.4-arm64-arm-64bit
- Python version: 3.12.1
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes (MPS)
- Using distributed or parallel set-up in script?: No

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

First of all, apologize if this might be better suited for the pytorch repository. I just wanted to get it out there in case anyone else stumbles upon this error.

I followed the example laid out at https://huggingface.co/docs/transformers/en/perplexity, using "mps" instead of "cuda" (on a Mac M2) as a device. The perplexity you achieve is way higher (around 3000) instead of the expected perplexity (15-30 depending on stride and model).

You can boil this down to the following minimal working example (happens with gpt2 und gpt2-large)

```

from transformers import GPT2LMHeadModel, GPT2TokenizerFast
from datasets import load_dataset

model_id = "gpt2"
model = GPT2LMHeadModel.from_pretrained(model_id)
tokenizer = GPT2TokenizerFast.from_pretrained(model_id)

test = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
encodings = tokenizer("\n\n".join(test["text"]), return_tensors="pt")

input_ids = encodings.input_ids[:, 0:1023]

model.to("mps")
out1 = model(input_ids.to("mps"), labels=input_ids.to("mps"))

model.to("cpu")
out2 = model(input_ids.to("cpu"), labels=input_ids.to("cpu"))

print(out1.loss)
print(out2.loss)

```

This outputs
```
tensor(7.6430, device='mps:0', grad_fn=<NllLossBackward0>)
tensor(2.8215, grad_fn=<NllLossBackward0>)
```

### Expected behavior

The device given should not change the behaviour/output of the model.

Comments:
Hi @Conzel 

Could you check the values of `input_ids` after it is moved to `mps` and then back to `cpu`?

I tried with `cuda` / `cpu` they give the same loss. I haven't try with `mps` yet. But it would be nice to see if the `input_ids` is affected first.

Hey, 

for me, if you compare cuda and cpu, everything lines up as well and the outputs of both print statements match. The issue is only on MPS vs CPU/CUDA. I tried on Colab and on a linux cluster, both work.

The input IDs stay the same, exemplary the first 20 ids:

```
print(input_ids.to("mps")[0][:20])
print(input_ids.to("cpu")[0][:20])
```

```
tensor([ 628,  796, 5199,  347, 2852,  353,  796,  220,  628,  628,  198, 5199,
         347, 2852,  353,  318,  281, 3594, 2646,  837], device='mps:0')
tensor([ 628,  796, 5199,  347, 2852,  353,  796,  220,  628,  628,  198, 5199,
         347, 2852,  353,  318,  281, 3594, 2646,  837])
```
Multiple moves also behave as expected:

```
assert torch.allclose(input_ids.to("mps"), input_ids.to("cpu").to("mps"))
assert torch.allclose(input_ids.to("cpu"), input_ids.to("mps").to("cpu"))
```

throws no error.

@amyeroberts Do you know usually who could work on issues that have `mps` involved?

@ydshieh There's no-one who specifically takes on mps issues, normally it comes down to what kind of issue it's causing e.g. is it in pipelines, text models etc. Just skimming the description, I suspect that the problem is a layer in the model e.g. [similar to this issue](https://github.com/huggingface/transformers/issues/22468). 

@Conzel Could you get the hidden states for the two forward passes (pass `output_hidden_states=True` on the forward call) and compare those? This way we can try and see where the differences are creeping in 

@amyeroberts @ydshieh
I tried the same code snippet with the same versions of the libraries installed on an M3 pro and I see that both return the same value,
```
tensor(2.8215, device='mps:0', grad_fn=<NllLossBackward0>)
tensor(2.8215, grad_fn=<NllLossBackward0>)
```

@amyeroberts 

> `output_hidden_states=True`

Great idea, running

```
for t1, t2 in zip(out1.hidden_states, out2.hidden_states):
    print((t1.to("cpu") - t2).abs().max())
```

yields

```
tensor(0., grad_fn=<MaxBackward1>)
tensor(73.1698, grad_fn=<MaxBackward1>)
tensor(56.7828, grad_fn=<MaxBackward1>)
tensor(1461.9398, grad_fn=<MaxBackward1>)
tensor(1434.7031, grad_fn=<MaxBackward1>)
tensor(1395.1794, grad_fn=<MaxBackward1>)
tensor(1334.4070, grad_fn=<MaxBackward1>)
tensor(1293.5073, grad_fn=<MaxBackward1>)
tensor(1271.1041, grad_fn=<MaxBackward1>)
tensor(1236.7236, grad_fn=<MaxBackward1>)
tensor(1221.2423, grad_fn=<MaxBackward1>)
tensor(1146.8076, grad_fn=<MaxBackward1>)
tensor(266.9192, grad_fn=<MaxBackward1>)
```

so the error seems to come up after the first GPT2 block. Checking the other issue you linked, I saw that GPT2 also has a GELU activation in its block at GPT2MLP. 

Following the advice from the issue and changing the activation in the forward pass of the GPT2MLP class to 
```
hidden_states = self.act(hidden_states.contiguous())
```

produces hidden states

```tensor(0., grad_fn=<MaxBackward1>)
tensor(3.4332e-05, grad_fn=<MaxBackward1>)
tensor(13.9683, grad_fn=<MaxBackward1>)
tensor(35.2026, grad_fn=<MaxBackward1>)
tensor(48.0906, grad_fn=<MaxBackward1>)
tensor(51.4316, grad_fn=<MaxBackward1>)
tensor(60.2854, grad_fn=<MaxBackward1>)
tensor(67.5542, grad_fn=<MaxBackward1>)
tensor(59.4678, grad_fn=<MaxBackward1>)
tensor(65.5447, grad_fn=<MaxBackward1>)
tensor(65.7066, grad_fn=<MaxBackward1>)
tensor(96.5678, grad_fn=<MaxBackward1>)
tensor(253.6287, grad_fn=<MaxBackward1>)
```


If I set the activation to identity instead, I get

```
tensor(0., grad_fn=<MaxBackward1>)
tensor(0.0002, grad_fn=<MaxBackward1>)
tensor(48.9532, grad_fn=<MaxBackward1>)
tensor(133.0081, grad_fn=<MaxBackward1>)
tensor(165.9325, grad_fn=<MaxBackward1>)
tensor(172.5591, grad_fn=<MaxBackward1>)
tensor(170.6113, grad_fn=<MaxBackward1>)
tensor(174.9679, grad_fn=<MaxBackward1>)
tensor(182.4258, grad_fn=<MaxBackward1>)
tensor(185.3785, grad_fn=<MaxBackward1>)
tensor(192.6411, grad_fn=<MaxBackward1>)
tensor(195.8365, grad_fn=<MaxBackward1>)
tensor(2.9336, grad_fn=<MaxBackward1>)
```

so the problem must (also) be located somewhere else.

@Conzel Thanks for sharing these details! 

Running the example reproducer on my M1 I wasn't able to reproduce and see the same numbers as @bayllama (thanks for testing on the M3!)

```
tensor(2.8215, device='mps:0', grad_fn=<NllLossBackward0>)
tensor(2.8215, grad_fn=<NllLossBackward0>)
```

I originally was running on torch 2.3.1, however downgrading to 2.2.0 I get the same results. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.