runtime shape mismatch for llama2-70
Author: ZhangShiyue
Pull request: no

### System Info

transformers-4.34.0.dev0
torch-2.1.0.dev20230711+cu121
accelerate-0.23.0.dev0
flash-attn-2.1.1

cuda 12.1
H100 GPUs

python-3.10


### 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

I'm trying to get model logits from llama2-70b for some text.

The following issue only occurs when batch size > 1 + use_flash_attention_2=True + llama2-70. 

In other words, if I use batch=1 or  use_flash_attention_2=False or llama2-7/13, no error will be thrown out.

FYI, dtype=torch.float16

```
Traceback (most recent call last):
...
  File "/workspace/llm_evaluator/tasks/evaluator_base.py", line 41, in compute_logits
    return model(input_ids=input_ids, attention_mask=attn_mask).logits
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1522, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1531, in _call_impl
    return forward_call(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1035, in forward
    outputs = self.model(
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1522, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1531, in _call_impl
    return forward_call(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 922, in forward
    layer_outputs = decoder_layer(
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1522, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1531, in _call_impl
    return forward_call(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 632, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1522, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1531, in _call_impl
    return forward_call(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 490, in forward
    attn_output = self._flash_attention_forward(
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 527, in _flash_attention_forward
    query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
  File "/layers/buildpacks/requirements/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 563, in _upad_input
    query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
RuntimeError: shape '[1400, 8, 128]' is invalid for input of size 11468800
```
note that the shape difference is always a factor of 8: 11468800 / 8 = 1400 * 8 * 128.

I noticed [this relevant issue](https://github.com/facebookresearch/llama/issues/423). I confirmed that in config.json, pretraining_tp=1, and I did not use bitsandbytes.

### Expected behavior

Shouldn't throw out this error.


Comments:
I tested flash-attn-2.3.1. It does not work either.

Hi @ZhangShiyue 
I suspect the issue comes from the fact that Llama-2-70b uses GQA (Grouped Query Attention). I created a dummy llama-2 model that uses GQA [here](https://huggingface.co/ybelkada/tiny-random-LlamaForCausalLM-GQA) (same config as llama-2-70B except for the number of layers) and I ran:

```python
import torch
from transformers import AutoModelForCausalLM

model_id = "ybelkada/tiny-random-LlamaForCausalLM-GQA"

model = AutoModelForCausalLM.from_pretrained(
    model_id, use_flash_attention_2=True, torch_dtype=torch.float16, low_cpu_mem_usage=True
).to(0)

dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).to(0)
attention_mask = torch.LongTensor([[1, 1, 1], [0, 1, 1]]).to(0)
print(model.generate(dummy_input, attention_mask=attention_mask, max_new_tokens=100))
```

And the script seemed to work on transformers main. Can you help me create a small reproducible snippet for your issue?

Thank you @younesbelkada! generate also works for me. What does not work is getting logits.
Here is a snippet of code I just tried and it gave me: RuntimeError: shape '[6, 8, 128]' is invalid for input of size 49152

```
import torch
from transformers import AutoTokenizer, AutoModel

loading_kwargs = {
    "device_map": "balanced_low_0",
    "use_flash_attention_2": True,
    "torch_dtype": torch.float16
}
model = AutoModel.from_pretrained("llama-2-70b/pytorch/", **loading_kwargs)
model.eval()

dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).cuda()
attention_mask = torch.LongTensor([[1, 1, 1], [0, 1, 1]]).cuda()

print(model(input_ids=dummy_input, attention_mask=attention_mask).logits)
```



Hi @ZhangShiyue , I still was not able to reproduce, can you run this snippet and confirm it fails on your end?

```python
import torch
from transformers import AutoModelForCausalLM, AutoModel

model_id = "ybelkada/tiny-random-LlamaForCausalLM-GQA"

model = AutoModelForCausalLM.from_pretrained(
    model_id, use_flash_attention_2=True, torch_dtype=torch.float16, low_cpu_mem_usage=True
).to(0)

dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).to(0)
attention_mask = torch.LongTensor([[1, 1, 1], [0, 1, 1]]).to(0)

print(model(input_ids=dummy_input, attention_mask=attention_mask).logits)
```
The main difference between my setup and yours is that I am using transformers main, can you try to update your transformers package?
```
pip install -U transformers
```

hey @younesbelkada, thank you! This issue has been resolved after I updated transformers.