Parallel inference on generative models throws an exception
Author: paulmakeraiimi
Pull request: no
Labels: name: Usage

name: bug

name: Generation

### System Info

Linux (ubuntu 20.04), transformers==4.40.0 onwards (including latest version), A100 GPU
Llama3 8B Instruct and LLama3.1 8B Instruct

### Who can help?

@zucchini-nlp @ArthurZucker @SunMarc

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

**Steps**

1. Install transformers==4.40.0 or newer
2. Get Llama2 8B Instruct
3. Load model using  AutoModelForCausalLM.from_pretrained onto GPU (I am using A100)
4. Try and use the same model from two threads (I have a web app that allow more than one request to use a single instance of a model)
5. Note that each thread has its own tokenizer as I know these are not thread safe
6. I am using the TextIteratorStreamer but it also blows up calling generate and waiting for the whole completion

The fault is:

  File "/home/AiimiAdmin/aiimi/InsightMaker.Python/AIModelService/venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 1989, in generate
    self._target(*self._args, **self._kwargs)
  File "/home/AiimiAdmin/aiimi/InsightMaker.Python/AIModelService/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/AiimiAdmin/aiimi/InsightMaker.Python/AIModelService/venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 1989, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/AiimiAdmin/aiimi/InsightMaker.Python/AIModelService/venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2969, in _sample
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/AiimiAdmin/aiimi/InsightMaker.Python/AIModelService/venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2969, in _sample
    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

**Note that this works fine if I use transformers==4.39.0 or older.** 

Seems like the logic has changed to use a shared data structure or there is something more fundamental with torch.

### Expected behavior

Inference runs, albeit slower as you are running more than one generate task with the same model and GPU.

Comments:
I think @zucchini-nlp and @gante are the good contacts to help you out :)

Hey @paulmakeraiimi !

Can you share a minimal example to reproduce the error?

Hi, thanks for looking in this. Here is a code snippet. It uses a single instance of a model (llama3) and then kicks off two threads, each with its own tokenizer. With transformers 4.39 it works fine, and with 4.40 onwards it blows up with the exception detailed above.

  from threading import Thread
  from transformers import AutoModelForCausalLM, AutoTokenizer
  import torch
  
  model_folder = "./models/Meta-Llama-3-8B-Instruct"
  print("loading model")
  model = AutoModelForCausalLM.from_pretrained(model_folder, device_map = "auto", torch_dtype = torch.float16)
  print("loaded model")
  
  def doChat(userPrompt):
      tokenizer = AutoTokenizer.from_pretrained(model_folder)
  
      messages = [
          {
              "role": "user",
              "content": userPrompt
          }
      ]
  
      prompt = tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True)
      input = tokenizer(prompt, return_token_type_ids = False, return_tensors = 'pt').to("cuda:0")
  
      print("-> " + userPrompt)
      output = model.generate(**input, max_new_tokens = 500, temperature = 0.1, 
                              repetition_penalty = 1.1, do_sample = True, 
                              eos_token_id = [tokenizer.eos_token_id,  tokenizer.convert_tokens_to_ids("
")])
      
      start_of_new_tokens = len(input[0])
      str = tokenizer.decode(output[0][start_of_new_tokens : ], skip_special_tokens = True)
      print(str)
  
  thread1 = Thread(target =lambda: doChat("Tell me about cheese?"))
  thread1.start()
  
  thread2 = Thread(target =lambda: doChat("Tell me about butter?"))
  thread2.start()
  
  thread1.join()
  thread2.join()

(I don't know why the code formatting is not working....)

![image](https://github.com/user-attachments/assets/65760883-dc3e-42db-ad4e-8746f1743f3c)


I could reproduce the error and located that it's related to SDPA attention specifically. Transformers models dispatch on SDPA by default and forcing the model to use "eager" attention solved the issue. Also using torch_dtype="float32" also solves it 

I cannot say right now why SDPA is causing NaNs when threading. I will try to dig further, but it might take some time. So you can use `attn_implementation="eager"` or `torch_dtype="float32"` when loading the model in the meanwhile :)

That works :) thank you... Be interesting to see what the underlying fault is

The NaNs only occur when we dispatch on Flash Attention 2 for `torch.scaled_dot_product` and pass in `is_causal=False` in decoding stage. cc @ArthurZucker as you know more about difficulties we had with SDPA

```python
attn_output = torch.nn.functional.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=causal_mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=is_causal,
        )
```

Not 100 sure, but a few things here: 
1. #32390 is a bit related, overall the way the masked is passed is very important. If you have padding left and the mask was not correctly created this can happen.
2. the dtype is important, if the model was trained with bfloat16, using float16 can influence the outputs of the model. 
3. If you force use the flash attention path, you have less precision than the `math` one,  -> instability! 

Hi all, is there a fix for this yet?

Did the issue persist? We have not seen other reports like this one

@zucchini-nlp was able to reproduce the issue. if it helps i can upgrade to the latest version of transformers and test this again

cc @gante  as well since you did  a bit of deep dive on parallel runs

If we replace the `scaled_dot_product_attention` call with the lines below in llama, at a first glance we see that the SDPA call is where all trouble starts. However, if we launch the script with `CUDA_LAUNCH_BLOCKING=1`, the issue may also pop up before SDPA. 

No issues at all if we run one thread at a time.

```py
        for tensor in (query_states, key_states, value_states, causal_mask):
            if tensor is None:
                continue
            if torch.isnan(tensor).any() or torch.isinf(tensor).any():
                print("failed before sdpa")
                exit()
        attn_output = torch.nn.functional.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=causal_mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=is_causal,
        )
        if torch.isnan(attn_output).any() or torch.isinf(attn_output).any():
            print("failed on sdpa")
            exit()
```

For reference, I modified the original script to NOT rely on sampling, for better reproducibility:

```py
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_folder = "meta-llama/Meta-Llama-3.1-8B-Instruct"
print("loading model")
model = AutoModelForCausalLM.from_pretrained(model_folder, device_map = "auto", torch_dtype = torch.float16)
print("loaded model")

def doChat(userPrompt):
    tokenizer = AutoTokenizer.from_pretrained(model_folder)

    messages = [
        {
            "role": "user",
            "content": userPrompt
        }
    ]

    prompt = tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True)
    input = tokenizer(prompt, return_token_type_ids = False, return_tensors = 'pt').to("cuda:0")

    print("-> " + userPrompt)
    output = model.generate(
        **input,
        max_new_tokens = 500,
        repetition_penalty = 1.1,
        do_sample = False,
        eos_token_id = [tokenizer.eos_token_id,  tokenizer.convert_tokens_to_ids("
")]
    )

    start_of_new_tokens = len(input[0])
    str = tokenizer.decode(output[0][start_of_new_tokens : ], skip_special_tokens = True)
    print(str)

thread1 = Thread(target =lambda: doChat("Tell me about cheese?"))
thread1.start()

thread2 = Thread(target =lambda: doChat("Tell me about butter?"))
thread2.start()

thread1.join()
thread2.join()
```

👉 Given that there are known workarounds for this issue (pool requests before launching the model) and finding the root cause seems non-trivial (can go all the way to torch's internals), I'd recommend not diving further into the issue 🤗 

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.