SDPA gives nans/infs during sampling on ROCM w/ float16
Author: cjekel
Pull request: no
Labels: name: bug

name: SDPA

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-4.18.0 RHEL 8
- Python version: 3.11.5
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.0.dev20240401+rocm6.0 (True)
- 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, 4x AMD Mi250x with `device_map='auto'`
- Using distributed or parallel set-up in script?: no


### Who can help?

@ArthurZucker and @younesbelkada

### 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

## observations

- SDPA gives nans/infs during sampling on ROCM w/ float16 or bfloat16
- Noticed models: Phi-2, Mistral 7b, Mixtral 8x7b, and likely more!
- using float32 make the issue go away
- Generation is complete gibberish until a NaN or INF is hit
- perhaps related https://github.com/huggingface/transformers/issues/30020 showed me that reverting to 'eagar' attention implementation get's rid of the problem

# reproduction

1. create a virtual environment
```bash
module load rocm/6.0.2

python3.11 -m venv mlvenv --system-site-packages
source mlvenv/bin/activate

python -m pip install --upgrade pip
python -m pip install --upgrade setuptools lit ipython
python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0
python -m pip install --no-cache-dir transformers==4.39.3 accelerate==0.28.0
```
2. Try to sample from mistral 7b (or mixtral 8x7b)
```python
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = 'mistralai/Mistral-7B-Instruct-v0.2'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map='auto',
    attn_implementation='sdpa',
)

inputs = tokenizer(
    "The world is",
    return_tensors="pt",
).to(model.device)
out = model.generate(
    **inputs,
    do_sample=True,
    max_new_tokens=256,
)

```
## Traceback
```python
RuntimeError                              Traceback (most recent call last)
Cell In[13], line 1
----> 1 out = model.generate(
      2     **inputs,
      3     do_sample=True,
      4     max_new_tokens=256,
      5 )

File /usr/WS2/jekel1/Repos/llmchat-latest/mlvenv/lib/python3.11/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File /usr/WS2/jekel1/Repos/llmchat-latest/mlvenv/lib/python3.11/site-packages/transformers/generation/utils.py:1575, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   1567     input_ids, model_kwargs = self._expand_inputs_for_generation(
   1568         input_ids=input_ids,
   1569         expand_size=generation_config.num_return_sequences,
   1570         is_encoder_decoder=self.config.is_encoder_decoder,
   1571         **model_kwargs,
   1572     )
   1574     # 13. run sample
-> 1575     result = self._sample(
   1576         input_ids,
   1577         logits_processor=prepared_logits_processor,
   1578         logits_warper=logits_warper,
   1579         stopping_criteria=prepared_stopping_criteria,
   1580         pad_token_id=generation_config.pad_token_id,
   1581         eos_token_id=generation_config.eos_token_id,
   1582         output_scores=generation_config.output_scores,
   1583         output_logits=generation_config.output_logits,
   1584         return_dict_in_generate=generation_config.return_dict_in_generate,
   1585         synced_gpus=synced_gpus,
   1586         streamer=streamer,
   1587         **model_kwargs,
   1588     )
   1590 elif generation_mode == GenerationMode.BEAM_SEARCH:
   1591     # 11. prepare beam search scorer
   1592     beam_scorer = BeamSearchScorer(
   1593         batch_size=batch_size,
   1594         num_beams=generation_config.num_beams,
   (...)
   1599         max_length=generation_config.max_length,
   1600     )

File /usr/WS2/jekel1/Repos/llmchat-latest/mlvenv/lib/python3.11/site-packages/transformers/generation/utils.py:2735, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, logits_warper, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, output_logits, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   2733 # sample
   2734 probs = nn.functional.softmax(next_token_scores, dim=-1)
-> 2735 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
   2737 # finished sentences should have their next token be a padding token
   2738 if eos_token_id is not None:

RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
```

Note that if we try phi-2, the generation does not throw inf or nan, but rather is just kind of giberish

```python
model_name = 'microsoft/phi-2'
...
out = model.generate(
    **inputs,
    do_sample=True,
    max_new_tokens=256,
)
tokenizer.batch_decode(out)
```
gives me something like:
>>>
['The world is, ", and you might feel sad at 3n, the "not the the best all of its focus and he\'s probably at which is that, we\'ve been working together and learning more\n\n#CI: The, a special day, the the most beautiful town, we are going to talk about some ideas so that the doctors in their next move, a series of events and the tell the world. COVID\n The LES is an important theme in this time.\n whole of Australia from the University of of New Mexico, who is a great way to do the work. a more comfortable bed. this scenario.\n\n\n".., to the main reason for its operations. The Party, on this situation, the first to store data. No, the state\u202ct. So often the owners? That will get into the y\n}\n\nI will tell you the amount is\nisopion of total weight of the elements of interest in a better futureocrin discreetly with your\n\n1,\'\'J\'s. What“b\')\n\n<class and last game-safe. Its Parents are being used for the whole community.\n accidents are a well-adir)) and Tangle in a1.']


### Expected behavior

If we switch to the 'eagar' attention implementation, everything is okay.

```python
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = 'mistralai/Mistral-7B-Instruct-v0.2'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map='auto',
    attn_implementation='eager',
)

inputs = tokenizer(
    "The world is",
    return_tensors="pt",
).to(model.device)
out = model.generate(
    **inputs,
    do_sample=True,
    max_new_tokens=256,
)
```
and if we use phi-2 we get the following.
```python
model_name = 'microsoft/phi-2'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map='auto',
    attn_implementation='eager',
)
out = model.generate(
    **inputs,
    do_sample=True,
    max_new_tokens=256,
)
tokenizer.batch_decode(out)
```
["The world is an incredibly interconnected globe that has always relied on individuals working together towards a shared goal. However, throughout history, we have seen that achieving this goal has not always been easy, and the path has been riddled with obstacles. Today, we are faced with a unique challenge in the form of a global pandemic that has forced us to work together in ways we never have before. Crafting together in the age of global social isolation has become an art, and one that we must master if we are to emerge stronger on the other side.\n\nAt the heart of any successful teamwork is communication and empathy. We must learn to listen to each other, to understand each other's perspectives, and to work together towards a common goal. This can be a challenge, especially when stress and tension are high. However, there are ways we can diffuse tense communications and manage our stress levels. We can use humor to lighten the mood, we can acknowledge and validate each other's feelings, and we can practice active listening to ensure that everyone's voice is heard.\n\nOne example of this successful teamwork is the development of vaccines for COVID-19. Scientists from all over the world have worked tirelessly to develop vaccines in record time, and their efforts have paid off. The vaccines have been"]

which makes a lot more sense...

Comments:
😶‍🌫️ that's annoying... cc @fxmarty for the ROCM part, I do have easy access to machines. 

Thank you @cjekel this is currently being investigated.

Thanks @fxmarty 

I'm confirming this is still and issue on the latest release.

```bash
#!/usr/bin/bash

module load rocm/6.0.3

python3.11 -m venv mlvenv --system-site-packages
source mlvenv/bin/activate

python -m pip install --upgrade pip
python -m pip install --upgrade setuptools lit ipython
python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.0
python -m pip install --no-cache-dir transformers==4.40.2 accelerate==0.30
```

Then running mistral 7b with sdpa gives

```
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
```


I can confirm that the workaround in https://github.com/huggingface/transformers/pull/30850 works! This was in `transformers==4.41.1` It logs a warning stating
> Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends.

Thank you for the fix! 🤗 

@fxmarty I don't see any other github issue. Is fixing the flash attention backend out of scope for this issue? If so we can close the issue.

I have the feeling not using FA2 on AMD is leaving quite a bit of performance on the table.

@siddharth9820 FYSA


@fxmarty Sorry for the spam. I have not tested https://github.com/ROCm/flash-attention and was curious if your team did and ran into the same issue? Otherwise, I'll test at some point and will report back.

we are waiting on https://github.com/pytorch/pytorch/pull/124885 to merge. 
cc: @pruthvistony @xinyazhang

@cjekel there is a bug in current SDPA + FA2 backend using aotriton (https://github.com/ROCm/aotriton) that is being investigated and fixed.

For https://github.com/ROCm/flash-attention, this is supported using the argument `attn_implementation="flash_attention_2"` when loading a model with `from_pretrained`. For more details you can refer to: https://huggingface.co/docs/transformers/v4.41.3/en/perf_infer_gpu_one?flashattention-2#flashattention-2

> we are waiting on [pytorch/pytorch#124885](https://github.com/pytorch/pytorch/pull/124885) to merge. cc: @pruthvistony @xinyazhang

@amathews-amd @fxmarty Looks like this one has been merged. Any plan to re-enable SDPA on ROCM?

cc @IlyasMoutawwakil 

@mht-sharma 

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.

Update to https://github.com/ROCm/aotriton/releases/tag/0.7.1b (Note must be from AOTriton 0.7b) may fix the problem

Analysis can be found at: https://github.com/ROCm/aotriton/pull/43

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.

I am hitting a similar issue (batched generation with padding producing nonsensical output on MI250) using the latest stable pytorch release `2.6.0+rocm6.2.4` (and `pytorch-triton-rocm==3.2.0`), independently on whether the host has rocm 6.2 or 6.3 installed. It is shipped with `libaotriton_v2.so.0.8.0`.

Using pytorch nightly (`torch==2.7.0.dev20250212+rocm6.3` & `pytorch-triton-rocm=3.2.0+git4b3bb1f8`) which is shipped with `libaotriton_v2.so.0.8.2`, there is no issue.

However, replacing `libaotriton_v2.so.0.8.0` by `libaotriton_v2.so.0.8.2` in `torch/lib` does not fix the issue, so the issue might not come from aotriton this time.

cc @mht-sharma @amathews-amd is this something you are aware of?

More precisely, the `SDPBackend.EFFICIENT_ATTENTION` backend is broken. MATH is fine.

Repro:
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
tokenizer.pad_token = tokenizer.eos_token

# using `attn_implementation="eager"` fixes the issue
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", attn_implementation="sdpa")

# using device = "cpu", there is no issue. Using torch==2.7.0.dev20250212+rocm6.3 there is no issue either.
device = "cuda"

self_tokenized = tokenizer(["Would you like to go to Paris an enjoy eating frogs with my French friends?", "Today I am in Paris", "How is the weather today my friend?", "This is really short"], padding=True, padding_side="left", return_tensors="pt").to(device)

model = model.eval().to(device)
model = model.to(torch.float16)

# the math backend also fixes the issue.
# from torch.nn.attention import SDPBackend, sdpa_kernel
#
# with sdpa_kernel(SDPBackend.MATH):
res = model.generate(**self_tokenized, max_new_tokens=20, do_sample=False, temperature=1, top_p=1)

results = tokenizer.batch_decode(res)

# Batched generation produces garbage with pytorch 2.6.0 stable release
for i, res in enumerate(results):
    print(f"\n\n\n\n ########################## {i}")
    print(res)
```

Hi @fxmarty-amd 
Since AOTriton 0.8, the shared library only serves as a dispatcher, and the GPU kernels are shipped inside the `aotriton.images` directory instead. 
This NaN problem is a GPU kernel problem (https://github.com/ROCm/aotriton/pull/68) and cannot be fixed by replacing the dispatcher.
(However, If you replace the `aotriton.images` directory from 0.8.2 but keep the 0.8.0 dispatcher the problem should be fixed)

Thank you @xinyazhang. I missed there was an `aotriton.images` file in `torch/lib`. Replacing it from 0.8.2b release fixes the issue for me.

@xinyazhang do you know if there are plans to include aotriton 0.8.2 as part of an upcoming torch 2.6 patch release, if any? This seems like a rather important bug, affecting the default execution path for batched inference of at least llama 7B (not sure if dependent on problem size or not) in Transformers, which in turns e.g. affects the evaluation of models in libraries like lm-eval-harness.

If not, I think Transformers should warn or disable SDPA on torch 2.6 (and lower if affected as well) on amd devices.