llama2 forward pass seemingly not working with padded inputs, unless one element in batch is not padded
Author: joebhakim
Pull request: no

### System Info

- `transformers` version: 4.33.2
- Platform: Linux-3.10.0-1160.99.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.9.16
- Huggingface_hub version: 0.17.2
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0.dev20230621+cu121 (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?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

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

Below is an example script that creates an example input with padding, then uses the first element (replicating when the tokenizer batch dimension is larger than the batch dimension fed into a model). Please let me know if I'm missing something!
```
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizerFast


tokenizer = AutoTokenizer.from_pretrained("togethercomputer/Llama-2-7B-32K-Instruct")

tokenizer = LlamaTokenizerFast.from_pretrained(
    "togethercomputer/Llama-2-7B-32K-Instruct"
)

model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/Llama-2-7B-32K-Instruct",
    trust_remote_code=True,
    torch_dtype=torch.float16,
).cuda()

""" THIS works in both cases
model = MT5ForConditionalGeneration.from_pretrained(
    'google/mt5-xl'
"""

encoded = tokenizer(
    [
        "
\nWrite a poem about cats\n
\n\n",
        "
\nWrite " + "a poem about" * 400 + " cats\n
\n\n",
    ],
    return_tensors="pt",
    padding="longest",
).to(model.device)

encoded_firstelem = {
    "input_ids": encoded["input_ids"][:1, :],
    "attention_mask": encoded["attention_mask"][:1, :],
}
breakpoint()

print(encoded_firstelem)
# {'input_ids': tensor([[    0,     0,     0,  ..., 29962,    13,    13]], device='cuda:0'), 'attention_mask': tensor([[0, 0, 0,  ..., 1, 1, 1]], device='cuda:0')}

# works
print(model(**encoded))

# breaks
print(model(**encoded_firstelem))
```


This gives the following:
```
Traceback (most recent call last):
  File "<console>", line 1, in <module>
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jh499/.cache/huggingface/modules/transformers_modules/togethercomputer/LLaMA-2-7B-32K/3c84db12268fac86081ec1229a5ef48414478c88/modeling_flash_llama.py", line 812, in forward
    outputs = self.model(
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jh499/.cache/huggingface/modules/transformers_modules/togethercomputer/LLaMA-2-7B-32K/3c84db12268fac86081ec1229a5ef48414478c88/modeling_flash_llama.py", line 696, in forward
    layer_outputs = decoder_layer(
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jh499/.cache/huggingface/modules/transformers_modules/togethercomputer/LLaMA-2-7B-32K/3c84db12268fac86081ec1229a5ef48414478c88/modeling_flash_llama.py", line 447, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jh499/env_39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jh499/.cache/huggingface/modules/transformers_modules/togethercomputer/LLaMA-2-7B-32K/3c84db12268fac86081ec1229a5ef48414478c88/modeling_flash_llama.py", line 380, in forward
    attn_output = pad_input(
RuntimeError: shape '[1, 1215, 4096]' is invalid for input of size 73728
```

### Expected behavior

This should work with the padded example, as it does in the mt5-xl model.

Comments:
WDYT @ArthurZucker @Rocketknight1 ?

Running on CPU + float32, but I can't reproduce this issue, unfortunately! The pasted code works for me. So either the issue is specific to the device + dtype, which seems unlikely, or it's some weird dependency/version problem.

@Rocketknight1 which vesion of CUDA are you on (I'm on 11.7 due to compatibility stuff)? Also, how did you get float32 to work? I get 
`RuntimeError: FlashAttention only support fp16 and bf16 data type` when I set `torch_dtype=torch.float32`, when it's on a cuda device. When it's not, I get `ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)`, so I assumed it *had* to be initialized like I did it?

Ah, sorry! I realized the issue - your reproduction script used `trust_remote_code=True`, but the code worked for me without it, so I thought that was just a mistake and I didn't need it. I realize now that the repo actually does have remote code, it just uses the `LLaMA` class names, and as a result will map to the code in the `transformers` library when `trust_remote_code=False`. When I use `trust_remote_code=True`, the error occurs.

Unfortunately, this means the error is specific to the user code in that repo - I'd suggest opening [an issue on the repo](https://huggingface.co/togethercomputer/Llama-2-7B-32K-Instruct/discussions) instead and pinging the authors!

Opened an issue on their page, will close this one. Thanks @Rocketknight1!

Hello , Did you solve thie problem @joebhakim