`MixtralFlashAttention2` subscripts `position_ids` before checking if it is `None`
Author: xenshinu
Pull request: no

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.2.2+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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker @youn

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

https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L421C9-L421C21
`position_ids` is declared as an optional input of `MixtralFlashAttention2.forward`, but it is subscript before checking if it is `None` here,
https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L447
Since it is also used below
https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L450
I think it should be a required variable or it should be calculated automatically somewhere else.


### Expected behavior

I think the `position_ids` should either be a required variable or it should be calculated automatically somewhere else.


Comments:
In fact, this `rotary_seq_len = ... +1` is also causing some issues. It makes the `cos` and `sin` not aligned with `query_states` and `key_states`. Refer to the original `MixtralAttention`
https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L350
It just uses `seq_len=kv_seq_len`. Is there any certain reason for that plus 1?
I modified the code to be 
```
 # Because the input can be padded, the absolute sequence length depends on the max position id.
        if position_ids is None:
            rotary_seq_len = kv_seq_len
        else:
            rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item())
```
Now at least it is runnable.

> ### System Info
> * `transformers` version: 4.41.2
> * Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.35
> * Python version: 3.10.14
> * Huggingface_hub version: 0.23.3
> * Safetensors version: 0.4.3
> * Accelerate version: not installed
> * Accelerate config: not found
> * PyTorch version (GPU?): 2.2.2+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?:
> * Using distributed or parallel set-up in script?:
> 
> ### Who can help?
> @ArthurZucker @youn
> 
> ### 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
> https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L421C9-L421C21 `position_ids` is declared as an optional input of `MixtralFlashAttention2.forward`, but it is subscript before checking if it is `None` here,
> 
> https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L447
> 
> 
> Since it is also used below
> https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L450
> 
> 
> I think it should be a required variable or it should be calculated automatically somewhere else.
> ### Expected behavior
> I think the `position_ids` should either be a required variable or it should be calculated automatically somewhere else.

I think both MixtralFlashAttention2 & MixtralFlashAttention should have the same operation, as padding can be the case in both. Nice find 

Hey! Could you share a snippet for a reproducible error? MixtralFlashAttention2 is working for me as the position ids are automatically created in the `Model`. But yeah, we should probably escape if someone wants to use the layer as it should / could be optional or enforce it as inputs. DO you want to open a PR to clarify this? 

@ArthurZucker Could you clarify why we need a plus one **outside** here?https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/mixtral/modeling_mixtral.py#L447
```
        # Because the input can be padded, the absolute sequence length depends on the max position id.
        rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1 # why + 1 outside? 
        cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
```
In the following function `apply_rotary_pos_emb`, `cos` and `sin` are selected by `position_ids`. If the input is not padded, `rotary_seq_len` is always `kv_seq_len+1 = position_ids.max()+2`, and the last element of `cos` is not used. If the input is padded, we can just make it as `position_ids[:, -1].max().item()+1`. 
I would also push a PR later.

Ah, btw I think this is kind of legacy code, as it was originally copied from `mistral`, but we changed it ot:
```python 
        cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
```
which no longer uses the rotary seq len. Now even with that, it's not a real hard bug as the sequence length is only used to slice the cos and sin cache, and if the dimension of cos and sin was wrong, we would not be able to perform the rope at all. 

You take `+1` because if you have max of cache position that is say 8, you want to take the 9 values (postiion 0 to 8 is 9)