GPT2 model getting NaN logits with MPS device
Author: uwu-420
Pull request: no

### System Info

- `transformers` version: 4.39.3
- Platform: macOS-14.4.1-arm64-arm-64bit
- Python version: 3.11.8
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.2.2 (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

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hello,

considering the following snippet

```python
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

set_seed(42)

model_name = "openai-community/gpt2"
device = "mps"
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

inputs = tokenizer("The sky is blue and grass is", return_tensors="pt")
input_ids = inputs.input_ids.to(device)
attention_mask = inputs.attention_mask.to(device)
outputs = model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    max_new_tokens=1,
    pad_token_id=tokenizer.eos_token_id
)
print(tokenizer.batch_decode(outputs))
```

When I run this on the device `"cpu"` everything works as expected, I get reasonable output (i.e. the token for `"green"`) but when using a `"mps"` device I just get the token for `"!"` which has the value `0`. It's because the logits are all NaNs.

At least for me, it's even a Heisenbug, because when I set a breakpoint here and then continue the logits look fine.
https://github.com/huggingface/transformers/blob/8e5f76f51196bb3b537a8d58c986819bb103f0a8/src/transformers/generation/utils.py#L2477C13-L2477C28

### Expected behavior

The GPT2 model loaded from Huggingface works correctly (no NaN logits) with a MPS device.

Comments:
Hi @uwu-420, thanks for opening this issue! 

I suspect the problem here is more likely to do with mps integration on the torch side the Hugging Face if the model works on CPU. The model returned from `AutoModelForCausalLM.from_pretrained` is a pytorch model and so calling `.to("mps")` I'm pretty sure is entirely handled by torch. 

cc @younesbelkada @SunMarc who might have some more insight here and can correct me if I'm wrong  

This is correct ! You might get some issues with Pytorch + MPS, can you follow the solution proposed in https://github.com/huggingface/transformers/issues/30229 ? - i.e. installing pytorch from source as things seems to be stable there for MPS devices on pytorch nightlies

Ok thanks for the hint! :) 

Hi @younesbelkada @amyeroberts. Wondering if there's any more support on this? Getting the same issue and moved to pytorch nightly but still getting same problem.

I am still having some issues with the latest torch/transformers versions:

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2-medium")
tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
model.to("cpu")
sentence = "Hello, world!"
inputs = tokenizer(sentence, return_tensors="pt")
inputs = {k: v.to("cpu") for k, v in inputs.items()}
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    print(logits.isnan().sum())
```
gives `tensor(1428)` nans, while 

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2-medium")
tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
model.to("mps")
sentence = "Hello, world!"
inputs = tokenizer(sentence, return_tensors="pt")
inputs = {k: v.to("mps") for k, v in inputs.items()}
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    print(logits.isnan().sum())
```
gives `tensor(0)`.