Inconsistent Rotation Base for Dynamic NTK Scaling RoPE
### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.19.0-42-generic-x86_64-with-glibc2.27
- Python version: 3.8.0
- PyTorch version (GPU?): 2.0.0+cu117 (True)

### Who can help?

@ArthurZucker @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

### Inconsistent problem

There is a subtle rotation inconsistency in the base factor of the DynamicNTKRope implemented in [transformers 4.31.0](https://github.com/huggingface/transformers/blob/b257c46a075419c09e5ce5c5aa39bc346ecdb9a5/src/transformers/models/llama/modeling_llama.py#L147)

Suppose we have a decoder model, like LLaMA-1, that utilizes DynamicNTKRope for interpolation and we want to evaluate it using perplexity. In any layer of this decoder model, after the key_states and query_states are computed from the hidden features, they are then rotated based on a fixed seq_len, which is the context length. 

However, while generating token by token beyond its maximum trained length at the inference stage, LLM usually reuses previous **cached keys** which are rotated based on factors associated with the previous seq_len. As the seq len keeps increasing, each cached key is rotated with respect to a different base, and consequently, the inconsistency between keys and values arises. 

### Expected behavior

I have conducted some experiments on the inconsistency and edited the codes about applying rotation on the keys and values to ensure the rotation base consistent [here](https://github.com/NormXU/Consistent-DynamicNTKRoPE/blob/main/scale_rope/consistent_rope_for_llama_patch.py). Please check the [repo](https://github.com/NormXU/Consistent-DynamicNTKRoPE) for further details.

While I haven't tested if a consistent rotation will benefit perplexity or downstream tasks in any dataset or language model, I believe that, from a mathematical perspective, keeping consistency in the rotation base could potentially enhance the language model's ability to learn relative positions more effectively. My intuition suggests that this consistency might offer advantages in capturing relative position information.
cc @ArthurZucker who works more on this (kind of) model(s).

However, @NormXU, could you provide a short but self-contained code snippet to demonstrate the `inconsistency` you mentioned? Thanks.<|||||>Sure

The inconsistency happens here:

https://github.com/huggingface/transformers/blob/a5cc30d72ae2dc19af534e4b35c986cc28db1275/src/transformers/models/llama/modeling_llama.py#L314-L330

While LLM generates token by token beyond its maximum trained length at the inference stage, the key_states are first applied RoPE based on cos and sin w.r.t. `kv_seq_len`, then rotated key_states are cached. 

Then when we come to the next token, key_states are applied to RoPE based on cos and sin w.r.t. `kv_seq_len + 1`. Since DynamicNTKScale has a rotation base w.r.t seq_len:

https://github.com/huggingface/transformers/blob/a5cc30d72ae2dc19af534e4b35c986cc28db1275/src/transformers/models/llama/modeling_llama.py#L156-L163

Therefore, we have an inconsistency between cached `keys_states` and between `keys_states` and query_states

<|||||>Actually @gante is the king of dynamic ROPE so will let him handle this! 🤗 <|||||>Hey @NormXU 👋 

I agree with your consistency issue you pointed out. However, our users' biggest concern (and ours, by extension) is empirical results, regardless of correctness. 

If you can share with us a few benchmarks where we see the change has positive benefits (and little to no downsides), I'll be more than happy to include it!<|||||>@gante Of course. 

How about the perplexity experiments I did in my repo [link](https://github.com/NormXU/Consistent-DynamicNTKRoPE), 

The way how we currently compute perplexity is more like we keep the rotation base consistent. Therefore, to bridge such a gap in rotation base between perplexity evaluation and inference with DynamicNTKScale, I modified the codes about how to apply the rotary embedding on keys and queries and do simple experiments on LLama1-7B.

After modification, the perplexity is computed in this way:

![image](https://github.com/huggingface/transformers/assets/33339685/b7936531-a59a-4792-bbbd-624600c72e40)

$K(\alpha(x))$ means, key_states is rotated by a rotation matrix whose base $\alpha$ is a function of sequence length.

Then, I compare the perplexity and the results are shown as below
![image](https://github.com/huggingface/transformers/assets/33339685/88a7888c-542a-400e-99c3-4873003aeb43)
This is about perplexity value on Llama1-7B, an 2k max sequence length model, values above 12.0 are cut off for concise; 
**Vanilla:** RoPE w/o any interpolation;
**NTK:** DynamicNTK when scale=1; 
**Consistent DynamicNTK:** keep rotation base between keys consistent, this is how we currently calculate perplexity 
 **Inconsistent DynamicNTK**: keep rotation base between keys inconsistent w.r.t context length; 

Can this experiment convincingly demonstrate that a consistent DynamicNTK can achieve better perplexity in long context than an inconsistent DynamicNTK?

Besides, could you please give me any advice on what benchmark I need to test this on? I have access to 8 x A100, enabling me to conduct many experiments quickly.

<|||||>@NormXU I see, that makes sense! 

A final question, before I make a decision: have you measured throughput vs the original dynamic scaling? Since we need to reapply RoPE to the cached values, it should introduce slowdowns. The decision to include the technique in `transformers` depends on the extent of the slowdown :) 