Can not sample next tokens with GPT-2 model with GPT2Config `reorder_and_upcast_attn=True`
Author: hogru
Pull request: no

### System Info

- `transformers` version: 4.24.0
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.9.12
- Huggingface_hub version: 0.11.0
- PyTorch version (GPU?): 1.12.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?: no (for debugging/generating, training done on GPU)
- Using distributed or parallel set-up in script?: no


### Who can help?

@patrickvonplaten (referenced in both GPT-2 and Text generation)

### 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

The following code uses a GPT-2 model trained from scratch. It works without problems when `reorder_and_upcast_attn=False`(the default) but with `reorder_and_upcast_attn=True` it throws `RuntimeError: probability tensor contains either ´inf´, ´nan´ or element < 0` when calling `model.sample()`. You might need to call `sample()` several times (I did 100 calls in my experiments).

```python
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)

mol = "CCC(C)(C)" 
mol_encoded = tokenizer(
        mol,
        add_special_tokens=True,
        padding=True,
)
input_ids = mol_encoded["input_ids"]
logits = model(**mol_encoded).logits[0]
assert not torch.any(torch.isinf(logits))  # just a safety check for debugging
assert not torch.any(torch.isnan(logits))  # just a safety check for debugging

# "Manual" sampling, works in both cases
probs = torch.nn.functional.softmax(logits, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)

# Hugging Face sampling
next_tokens = model.sample(
        input_ids,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=100,
) # This raises the RuntimeError exception when `reorder_and_upcast_attn=True`
```


### Expected behavior

I can think of the following options
- IF it's a bug, fix it ;-), i.e. `reorder_and_upcast_attn=True` should work for text generation
- if "massaging" the logits with `logits_processor` or `logits_warper` is required, update the docs
- improve my understanding of the option...


Comments:
cc @gante 

Hi @hogru 👋 Having a popular project like `transformers` means we get many support and feature requests — if we want to maximize how much we help the community, the community has to help us stay productive 🙏

To that end, please share a *short* script where the issue is clearly reproducible on *any* computer. In your particular case, your example is missing the model itself, which can influence the `sample` call in many ways (e.g. depending on the model config). Thank you 🤗

Hi @gante, I get this. Since I can avoid the issue by not using that option and you hint at the issue being specific to my environment/config/model/... I save you and me some time and "close" the issue. I assumed this to be a generic issue and wanted to let you know. Solving it for my specific use case is not a priority.

Hey @hogru -- actually it may be an issue that happens on all sorts of environments and models :) 

I didn't mean to sound dismissive. The limitation here is manpower: we have many issues per maintainer, so our focus is on 1) common issues; 2) issues where we can pin the issue. This is the first time I see this issue, so 1) doesn't apply. For 2) to happen, I need to be able to reproduce the issue quickly, otherwise it will be a huge time sink to find the exact problem so it can be fixed 🤗 That's where the short script comes in!

Hi @gante, thanks for reaching out, I did not perceive it as dismissive and my answer was intended to be in a friendly voice. But English is not my native language... And I am new to hugging face which means that (a) there's a chance that I overlook something obvious and (b) I need to figure out how to push the model to the hub, ahem. I know, probably pretty simple, but I am uncertain if it's intended to hold models for debugging. And again, following your argument, your time is likely better invested in other areas. So, all good here.