Error inferencing on 2x A100 GPUs
Author: RonanKMcGovern
Pull request: no

### System Info

I'm getting this issue with 2x A100s for inference (or training too), but I don't get the issue with 4x A6000s.

*Environment*
```
- `transformers` version: 4.36.0.dev0
- Platform: Linux-5.4.0-149-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.18.0
- Safetensors version: 0.4.0
- Accelerate version: 0.25.0.dev0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1+cu118 (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, 2x A100s.
- Using distributed or parallel set-up in script?: Model split across gpus.
```

### Who can help?

@ArthurZucker @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

*Reproduction:*
```
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id  = "tiiuae/falcon-40B"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map={"":"cuda:0"}, # I also tried 'auto'
    # trust_remote_code=False,
    torch_dtype=torch.bfloat16,
    use_flash_attention_2=True, # works with Llama models and reduces memory reqs)

import torch
import gc

tokenizer = AutoTokenizer.from_pretrained(model_id,use_fast=True)

def generate_simple(model, tokenizer, prompt):
    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs, max_length=50)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Basic test of model generation
prompt = "The quick brown fox"
generated_text = generate_simple(model, tokenizer, prompt)
print(generated_text)
```

*Error*
```
The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:11 for open-end generation.
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[9], line 13
     11 # Basic test of model generation
     12 prompt = "The quick brown fox"
---> 13 generated_text = generate_simple(model, tokenizer, prompt)
     14 print(generated_text)
     16 # Clear GPU cache and run garbage collection

Cell In[9], line 8, in generate_simple(model, tokenizer, prompt)
      6 def generate_simple(model, tokenizer, prompt):
      7     inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
----> 8     outputs = model.generate(inputs, max_length=50)
      9     return tokenizer.decode(outputs[0], skip_special_tokens=True)

File /usr/local/lib/python3.10/dist-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/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1753, 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)
   1736     return self.assisted_decoding(
   1737         input_ids,
   1738         assistant_model=assistant_model,
   (...)
   1749         **model_kwargs,
   1750     )
   1751 if generation_mode == GenerationMode.GREEDY_SEARCH:
   1752     # 11. run greedy search
-> 1753     return self.greedy_search(
   1754         input_ids,
   1755         logits_processor=logits_processor,
   1756         stopping_criteria=stopping_criteria,
   1757         pad_token_id=generation_config.pad_token_id,
   1758         eos_token_id=generation_config.eos_token_id,
   1759         output_scores=generation_config.output_scores,
   1760         return_dict_in_generate=generation_config.return_dict_in_generate,
   1761         synced_gpus=synced_gpus,
   1762         streamer=streamer,
   1763         **model_kwargs,
   1764     )
   1766 elif generation_mode == GenerationMode.CONTRASTIVE_SEARCH:
   1767     if not model_kwargs["use_cache"]:

File /usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:2614, in GenerationMixin.greedy_search(self, input_ids, logits_processor, stopping_criteria, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   2611 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
   2613 # forward pass to get next token
-> 2614 outputs = self(
   2615     **model_inputs,
   2616     return_dict=True,
   2617     output_attentions=output_attentions,
   2618     output_hidden_states=output_hidden_states,
   2619 )
   2621 if synced_gpus and this_peer_finished:
   2622     continue  # don't waste resources running the code we don't need

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:1246, in FalconForCausalLM.forward(self, input_ids, past_key_values, attention_mask, position_ids, head_mask, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1237 r"""
   1238 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
   1239     Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
   1240     `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
   1241     are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
   1242 """
   1244 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1246 transformer_outputs = self.transformer(
   1247     input_ids,
   1248     past_key_values=past_key_values,
   1249     attention_mask=attention_mask,
   1250     position_ids=position_ids,
   1251     head_mask=head_mask,
   1252     inputs_embeds=inputs_embeds,
   1253     use_cache=use_cache,
   1254     output_attentions=output_attentions,
   1255     output_hidden_states=output_hidden_states,
   1256     return_dict=return_dict,
   1257 )
   1258 hidden_states = transformer_outputs[0]
   1260 lm_logits = self.lm_head(hidden_states)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:1122, in FalconModel.forward(self, input_ids, past_key_values, attention_mask, position_ids, head_mask, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1110     outputs = self._gradient_checkpointing_func(
   1111         block.__call__,
   1112         hidden_states,
   (...)
   1119         output_attentions,
   1120     )
   1121 else:
-> 1122     outputs = block(
   1123         hidden_states,
   1124         layer_past=layer_past,
   1125         attention_mask=attention_mask,
   1126         position_ids=position_ids,
   1127         head_mask=head_mask[i],
   1128         use_cache=use_cache,
   1129         output_attentions=output_attentions,
   1130         alibi=alibi,
   1131     )
   1133 hidden_states = outputs[0]
   1134 if use_cache is True:

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:796, in FalconDecoderLayer.forward(self, hidden_states, alibi, attention_mask, position_ids, layer_past, head_mask, use_cache, output_attentions, **kwargs)
    793     attention_layernorm_out = self.input_layernorm(hidden_states)
    795 # Self attention.
--> 796 attn_outputs = self.self_attention(
    797     attention_layernorm_out,
    798     layer_past=layer_past,
    799     attention_mask=attention_mask,
    800     position_ids=position_ids,
    801     alibi=alibi,
    802     head_mask=head_mask,
    803     use_cache=use_cache,
    804     output_attentions=output_attentions,
    805     **kwargs,
    806 )
    808 attention_output = attn_outputs[0]
    810 if not self.config.new_decoder_architecture:

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:559, in FalconFlashAttention2.forward(self, hidden_states, alibi, attention_mask, position_ids, layer_past, head_mask, use_cache, output_attentions, **kwargs)
    556     # overwrite attention_mask with padding_mask
    557     attention_mask = kwargs.pop("padding_mask")
--> 559 fused_qkv = self.query_key_value(hidden_states)  # [batch_size, seq_length, 3 x hidden_size]
    560 num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
    561 # 3 x [batch_size, seq_length, num_heads, head_dim]

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:68, in FalconLinear.forward(self, input)
     67 def forward(self, input: torch.Tensor) -> torch.Tensor:
---> 68     hidden_states = input @ self.weight.T
     69     if self.bias is None:
     70         return hidden_states

RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`
```

### Expected behavior

Inference (and training) should work fine, just as is the case with 1x A100 or 6x A6000s.

Comments:
Hi @RonanKMcGovern 
Thanks for the issue! :D 
One thing that is a bit unusual here is that you are forcing the whole model to fit on a single GPU (through `device_map={"": "cuda:0"}`) for a 40B model in bfloat16. I don't think such a model should fit in a single A100 GPU for generation. Since you force the model to be loaded on a single GPU with `{"": "cuda:0"}` it might led to something wrong. 
Can you share more details about your GPU hardware (40GB or 80GB RAM?) and how you ran your script with `"auto"` for the device map

Hi @younesbelkada .

- The A100s are 80 GB.
- You're right the model won't fit on one, so doing '{"":"cuda:0"}' doesn't make sense. I simply misunderstood what that did. In any case, I first had the error with 'auto' and then tried '{"":"cuda:0"}' (which in hindsight is wrong).

When running with auto, I used the same script just with:
```
device_map='auto'
```

when loading the model

@younesbelkada was it possible to replicate?

Hi @younesbelkada , moving to the last stable transformers release seems to have solved the issue:
```
- `transformers` version: 4.35.0
- Platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.25.0.dev0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+cu118 (True)
```
BTW, the same repro above with the dev version is now giving:
```
The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:11 for open-end generation.
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[3], line 11
      9 # Basic test of model generation
     10 prompt = "The quick brown fox"
---> 11 generated_text = generate_simple(model, tokenizer, prompt)
     12 print(generated_text)

Cell In[3], line 6, in generate_simple(model, tokenizer, prompt)
      4 def generate_simple(model, tokenizer, prompt):
      5     inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
----> 6     outputs = model.generate(inputs, max_length=50)
      7     return tokenizer.decode(outputs[0], skip_special_tokens=True)

File /usr/local/lib/python3.10/dist-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/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1753, 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)
   1736     return self.assisted_decoding(
   1737         input_ids,
   1738         assistant_model=assistant_model,
   (...)
   1749         **model_kwargs,
   1750     )
   1751 if generation_mode == GenerationMode.GREEDY_SEARCH:
   1752     # 11. run greedy search
-> 1753     return self.greedy_search(
   1754         input_ids,
   1755         logits_processor=logits_processor,
   1756         stopping_criteria=stopping_criteria,
   1757         pad_token_id=generation_config.pad_token_id,
   1758         eos_token_id=generation_config.eos_token_id,
   1759         output_scores=generation_config.output_scores,
   1760         return_dict_in_generate=generation_config.return_dict_in_generate,
   1761         synced_gpus=synced_gpus,
   1762         streamer=streamer,
   1763         **model_kwargs,
   1764     )
   1766 elif generation_mode == GenerationMode.CONTRASTIVE_SEARCH:
   1767     if not model_kwargs["use_cache"]:

File /usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:2614, in GenerationMixin.greedy_search(self, input_ids, logits_processor, stopping_criteria, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   2611 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
   2613 # forward pass to get next token
-> 2614 outputs = self(
   2615     **model_inputs,
   2616     return_dict=True,
   2617     output_attentions=output_attentions,
   2618     output_hidden_states=output_hidden_states,
   2619 )
   2621 if synced_gpus and this_peer_finished:
   2622     continue  # don't waste resources running the code we don't need

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:164, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    162         output = module._old_forward(*args, **kwargs)
    163 else:
--> 164     output = module._old_forward(*args, **kwargs)
    165 return module._hf_hook.post_forward(module, output)

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:1175, in FalconForCausalLM.forward(self, input_ids, past_key_values, attention_mask, position_ids, head_mask, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1166 r"""
   1167 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
   1168     Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
   1169     `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
   1170     are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
   1171 """
   1173 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1175 transformer_outputs = self.transformer(
   1176     input_ids,
   1177     past_key_values=past_key_values,
   1178     attention_mask=attention_mask,
   1179     position_ids=position_ids,
   1180     head_mask=head_mask,
   1181     inputs_embeds=inputs_embeds,
   1182     use_cache=use_cache,
   1183     output_attentions=output_attentions,
   1184     output_hidden_states=output_hidden_states,
   1185     return_dict=return_dict,
   1186 )
   1187 hidden_states = transformer_outputs[0]
   1189 lm_logits = self.lm_head(hidden_states)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:1054, in FalconModel.forward(self, input_ids, past_key_values, attention_mask, position_ids, head_mask, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1042     outputs = self._gradient_checkpointing_func(
   1043         block.__call__,
   1044         hidden_states,
   (...)
   1051         output_attentions,
   1052     )
   1053 else:
-> 1054     outputs = block(
   1055         hidden_states,
   1056         layer_past=layer_past,
   1057         attention_mask=attention_mask,
   1058         position_ids=position_ids,
   1059         head_mask=head_mask[i],
   1060         use_cache=use_cache,
   1061         output_attentions=output_attentions,
   1062         alibi=alibi,
   1063     )
   1065 hidden_states = outputs[0]
   1066 if use_cache is True:

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:164, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    162         output = module._old_forward(*args, **kwargs)
    163 else:
--> 164     output = module._old_forward(*args, **kwargs)
    165 return module._hf_hook.post_forward(module, output)

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:766, in FalconDecoderLayer.forward(self, hidden_states, alibi, attention_mask, position_ids, layer_past, head_mask, use_cache, output_attentions, **kwargs)
    763     attention_layernorm_out = self.input_layernorm(hidden_states)
    765 # Self attention.
--> 766 attn_outputs = self.self_attention(
    767     attention_layernorm_out,
    768     layer_past=layer_past,
    769     attention_mask=attention_mask,
    770     position_ids=position_ids,
    771     alibi=alibi,
    772     head_mask=head_mask,
    773     use_cache=use_cache,
    774     output_attentions=output_attentions,
    775     **kwargs,
    776 )
    778 attention_output = attn_outputs[0]
    780 if not self.config.new_decoder_architecture:

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /usr/local/lib/python3.10/dist-packages/accelerate/hooks.py:164, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    162         output = module._old_forward(*args, **kwargs)
    163 else:
--> 164     output = module._old_forward(*args, **kwargs)
    165 return module._hf_hook.post_forward(module, output)

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:593, in FalconFlashAttention2.forward(self, hidden_states, alibi, attention_mask, position_ids, layer_past, head_mask, use_cache, output_attentions, **kwargs)
    590     key_layer = key_layer.to(target_dtype)
    591     value_layer = value_layer.to(target_dtype)
--> 593 attn_output = self._flash_attention_forward(
    594     query_layer, key_layer, value_layer, attention_mask, query_length, dropout=attn_dropout
    595 )
    597 attn_weights = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
    598 attn_output = self.dense(attn_weights)

File /usr/local/lib/python3.10/dist-packages/transformers/models/falcon/modeling_falcon.py:653, in FalconFlashAttention2._flash_attention_forward(self, query_states, key_states, value_states, attention_mask, query_length, dropout, softmax_scale)
    651     attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
    652 else:
--> 653     attn_output = flash_attn_func(
    654         query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal
    655     )
    657 return attn_output

File /usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_interface.py:705, in flash_attn_func(q, k, v, dropout_p, softmax_scale, causal, window_size, return_attn_probs)
    652 def flash_attn_func(
    653     q,
    654     k,
   (...)
    660     return_attn_probs=False,
    661 ):
    662     """dropout_p should be set to 0.0 during evaluation
    663     Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
    664     than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
   (...)
    703             pattern (negative means that location was dropped, nonnegative means it was kept).
    704     """
--> 705     return FlashAttnFunc.apply(
    706         q, k, v, dropout_p, softmax_scale, causal, window_size, return_attn_probs
    707     )

File /usr/local/lib/python3.10/dist-packages/torch/autograd/function.py:539, in Function.apply(cls, *args, **kwargs)
    536 if not torch._C._are_functorch_transforms_active():
    537     # See NOTE: [functorch vjp and autograd interaction]
    538     args = _functorch.utils.unwrap_dead_wrappers(args)
--> 539     return super().apply(*args, **kwargs)  # type: ignore[misc]
    541 if cls.setup_context == _SingleLevelFunction.setup_context:
    542     raise RuntimeError(
    543         "In order to use an autograd.Function with functorch transforms "
    544         "(vmap, grad, jvp, jacrev, ...), it must override the setup_context "
    545         "staticmethod. For more details, please see "
    546         "https://pytorch.org/docs/master/notes/extending.func.html"
    547     )

File /usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_interface.py:434, in FlashAttnFunc.forward(ctx, q, k, v, dropout_p, softmax_scale, causal, window_size, return_softmax)
    432 if softmax_scale is None:
    433     softmax_scale = q.shape[-1] ** (-0.5)
--> 434 out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
    435     q,
    436     k,
    437     v,
    438     dropout_p,
    439     softmax_scale,
    440     causal=causal,
    441     window_size=window_size,
    442     return_softmax=return_softmax and dropout_p > 0,
    443 )
    444 ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
    445 ctx.dropout_p = dropout_p

File /usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_interface.py:47, in _flash_attn_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, return_softmax)
     45 maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
     46 q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
---> 47 out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = flash_attn_cuda.fwd(
     48     q,
     49     k,
     50     v,
     51     None,
     52     dropout_p,
     53     softmax_scale,
     54     causal,
     55     window_size[0],
     56     window_size[1],
     57     return_softmax,
     58     None,
     59 )
     60 return out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state

RuntimeError: Number of heads in key/value must divide number of heads in query
```

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.