Making the unsqueeze dimension parameterized in the apply_rotary_pos_emb function in modeling_llama.py
Author: ShashankMosaicML
Pull request: no

### Feature request

Hi Huggingface Team, 
We would like to request to make the 'unsqueeze(1)' in [these two lines ](https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L209C1-L210C41) parameterized. To be precise, we would like to request that the following lines

```
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    cos = cos[position_ids].unsqueeze(1) 
    sin = sin[position_ids].unsqueeze(1)
    ...
```
be converted to something like the following
```
def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    ...
```

### Motivation

We are trying to import and use the [apply_rotary_pos_emb function](https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L208) and the [LlamaRotaryEmbedding class](https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L119). However, our query and key tensors have the shape `[batch_size, sequence_len, heads, dim]`. To make them compatible with the `apply_rotary_pos_emb` function, we have to transpose the tensors to `[batch_size, heads, sequence_len, dim]`, call `apply_rotary_pos_emb` on them and then transpose the tensors back to `[batch_size, sequence_len, heads, dim]`. These unnecessary transposes could be avoided if the `apply_rotary_pos_emb` function had a parameter which controlled the dimension on which the `unsqueeze`'s [here](https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L209C1-L210C41) were applied. 

_Please note that the Llama Huggingface code also [does similar back and forth transposes](https://github.com/huggingface/transformers/blob/08a2edfc6629a323effd7a85feafed9e6701e2dd/src/transformers/models/llama/modeling_llama.py#L442C1-L463C52), and hence could benefit from this **very small** code change as well._

### Your contribution

We are willing to submit a PR for this.

Comments:
MMM changing the transpose would change the past key values which is not backward compatible. The unsqueeze dimension could be a nice to have but would rather have it in the init of the rope embedding than in the forward as this is another not backward compatible change 😓 
we refrain from adding custom usage code as `transformers` is not meant to be a tool box with 10K config arguments, but if removing the 2 transpose gives a very good performance gain we could consider this IMO. 

cc @gante for ROPE 🤗 

@ShashankMosaicML I do not oppose adding `unsqueeze_dim=1` as an argument if it makes life easier for others 🤗 

As for backward compatibility -- this is a good question. As of now, as @ArthurZucker said, we don't want to break it. However, we are currently working on a new cache structure (with a flag for backwards compatibility), so making the new cache structure optimal for FA2 might be a good idea 🤔 

If we can remove 2 transpose (it's not a bottlneck but still) would be nice. Let's keep that in mind when refactoring the cache cc @tomaarsen as well. 

@ShashankMosaicML feel free to open a PR and tag us 🤗 

Great, will do this soon!

Here is the [link to the pull request](https://github.com/huggingface/transformers/pull/27117).