Fix LlamaDynamicNTKScalingRotaryEmbedding cache
Author: i4never
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #25306 . Cached pos embed will only be used when cached length is `model.config.max_position_embeddings` and asked `seq_len` is small or equal than this.
It's a little bit slower but correct. Value of sin/cos could be calculated in first decoder layer and passed to others (that's what [TGI flash_llama_modeling](https://github.com/huggingface/text-generation-inference/blob/f9910d13e296989f41e714c43eb60ce051359db3/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py#L443-L460) does)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. #25306 
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?
@gante 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->

## Test
```python
import time
import random
import torch
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding


class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):
    """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""

    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
        self.scaling_factor = scaling_factor
        super().__init__(dim, max_position_embeddings, base, device)

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        if seq_len <= self.max_position_embeddings:
            seq_len = self.max_position_embeddings

        self.max_seq_len_cached = seq_len

        if seq_len > self.max_position_embeddings:
            base = self.base * (
                    (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))
        else:
            base = self.base
        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))

        t = torch.arange(self.max_seq_len_cached, device=device, dtype=inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if not (seq_len <= self.max_seq_len_cached == self.max_position_embeddings):
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )


class LlamaDynamicNTKScalingRotaryEmbeddingOld(LlamaRotaryEmbedding):
    """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""

    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
        self.scaling_factor = scaling_factor
        super().__init__(dim, max_position_embeddings, base, device)

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len

        if seq_len > self.max_position_embeddings:
            base = self.base * (
                    (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))
            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
            self.register_buffer("inv_freq", inv_freq, persistent=False)

        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)


rope_ntk_new = LlamaDynamicNTKScalingRotaryEmbedding(128)
rope_ntk_old = LlamaDynamicNTKScalingRotaryEmbeddingOld(128)

x = torch.rand(1, 9, 256)
seq_lengths = [128, 512, 1024, 2048, 4096, 8192, 4096, 2048, 1024, 512, 128]

print('=' * 32)
print('Result by new version')
result_by_rope_ntk_new = [rope_ntk_new(x, seq_length) for seq_length in seq_lengths]
for i in range(len(seq_lengths) // 2):
    cos_before = result_by_rope_ntk_new[i][0]
    cos_after = result_by_rope_ntk_new[len(seq_lengths) - i - 1][0]
    sin_before = result_by_rope_ntk_new[i][1]
    sin_after = result_by_rope_ntk_new[len(seq_lengths) - i - 1][1]
    print(f"length {seq_lengths[i]} {seq_lengths[len(seq_lengths) - i - 1]}")
    print(f"\tcos is equal: {torch.equal(cos_before, cos_after)}")
    print(f"\tsin is equal: {torch.equal(sin_before, sin_after)}")

print('=' * 32)
print('Result by old version')
result_by_rope_ntk_old = [rope_ntk_old(x, seq_length) for seq_length in seq_lengths]
for i in range(len(seq_lengths) // 2):
    cos_before = result_by_rope_ntk_old[i][0]
    cos_after = result_by_rope_ntk_old[len(seq_lengths) - i - 1][0]
    sin_before = result_by_rope_ntk_old[i][1]
    sin_after = result_by_rope_ntk_old[len(seq_lengths) - i - 1][1]
    print(f"length {seq_lengths[i]} {seq_lengths[len(seq_lengths) - i - 1]}")
    print(f"\tcos is equal: {torch.equal(cos_before, cos_after)}")
    print(f"\tsin is equal: {torch.equal(sin_before, sin_after)}")


seq_lengths = [random.randint(1, 32768) for _ in range(10000)]
start = time.time()
for seq_length in seq_lengths:
    rope_ntk_new(x, seq_length)
end = time.time()
print(f"{end - start:.3f}s total, {(end - start) / 10000:.5f}s each")

start = time.time()
for seq_length in seq_lengths:
    rope_ntk_old(x, seq_length)
end = time.time()
print(f"{end - start:.3f}s total, {(end - start) / 10000:.5f}s each")
```
Output:
```shell
================================
Result by new version
length 128 128
	cos is equal: True
	sin is equal: True
length 512 512
	cos is equal: True
	sin is equal: True
length 1024 1024
	cos is equal: True
	sin is equal: True
length 2048 2048
	cos is equal: True
	sin is equal: True
length 4096 4096
	cos is equal: True
	sin is equal: True
================================
Result by old version
length 128 128
	cos is equal: False
	sin is equal: False
length 512 512
	cos is equal: False
	sin is equal: False
length 1024 1024
	cos is equal: False
	sin is equal: False
length 2048 2048
	cos is equal: False
	sin is equal: False
length 4096 4096
	cos is equal: False
	sin is equal: False

7.573s total, 0.00076s each
0.101s total, 0.00001s each
```

Comments:
Retriggering the test may fix the test timeout.

Yep, rebasing on main should also help!

Ouch 😅 Also a 10x increase in time seems like a lot 

Each forward will take several millseconds more. If we cache cos/sin in first layer and pass it to remains (as [llama in tgi](https://github.com/huggingface/text-generation-inference/blob/96a982ad8fc232479384476b1596a880697cc1d0/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py#L443-L460) doing), more time could be saved.

Yeah for sure, but that is a bit too big of a change 

Then only the issue of throughput remains. 

I have no obvious suspect of what could be causing it. I have two suggestions: 
1. as a sanity check, confirm that there is no throughput degradation as you increase the sequence length (and not as you reduce it). The changes you introduced should cause no changes here.
2. Run your whole test script in a loop a few times, and discard the first 2 iterations. A major part of the slowdown could be due to GPU initialization time.

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.