[FA-2] inference dtype mismatch bug
Author: wizyoung
Pull request: no

### System Info

- `transformers` version: 4.35.0.dev0
- Platform: Linux-4.14.0_1-0-0-44-x86_64-with-glibc2.27
- Python version: 3.9.17
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.2
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 1.13.1+cu117 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@younesbelkada @ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I installed the transformers from the latest main branch (commit bdbcd5d4) due to lots of bug fixes related to flash attn2. However, I encountered several dtype mismatch issues during inference. I've pinpointed the cause of the first issue, but the second one remains elusive.
(1) cat dtype mismatch when use_cache is set True:
https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L454-L457
The past_key_value tensors' dtype might be upcast to fp32 in ops afterwards (and it clearly will), so we need to make sure the dtype matches. The dtype mismatch is 100% reproducible in my llama inference code where the LLM weights are frozen and the .generate function is called.
correct code should be:
```python
key_states = torch.cat([past_key_value[0].to(key_states.dtype), key_states], dim=2)
value_states = torch.cat([past_key_value[1].to(value_states.dtype), value_states], dim=2)
```
(2) I also met "get error  tensors used as indices must be long, byte or bool tensors" in LLM inference with use_cache set as True. This error occasionally occurs and it seems to point to the flash-attn inner func. 

Have you tested inference mode when implementing FA-2? By the way, I wonder whether FA2 supports batch inference with padding and the padding mode is "longest", where I tested on xformers but got poor inference speed.

### Expected behavior

FA-2 batch inference works normal.

Comments:
Hello! Can you share a small snippet with a reproducer? Would help use know what is happening and could be useful to improve our tests! 

@ArthurZucker My team project is quite complicated and I'm trying to extract a reproducible snippet and post it here later.

closed as latest version of hf fixed it.

Another problem:
https://github.com/huggingface/transformers/blob/66b088faf01a795a7e0ddafafa1838f065f42f86/src/transformers/models/llama/modeling_llama.py#L719-L723

The dtype of indices_qhere shoud be torch.int64 as later we will use it by calling 
`attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)`

pad_input requires indices_q to be long, byte or bool dtype and it cannot be int32. So we need to change it to:
` indices_q = cu_seqlens_q[:-1].long()`
![image](https://github.com/huggingface/transformers/assets/13296106/de37c0d2-36ed-4893-9238-4521df0ce5f8)




which version of FA are you using @wizyoung ?

> which version of FA are you using @wizyoung ?

I'm using the latest 2.3.3，more specificly，it's https://github.com/Dao-AILab/flash-attention/releases/download/v2.3.3/flash_attn-2.3.3+cu117torch1.13cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

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.

Do you still face the same issue on latest transformers @wizyoung ?

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.