Llama 3.1: replace for loop by tensor ops at inv_freq initialization
Author: gante
Pull request: yes

# What does this PR do?

Fixes #32187 

The old `inv_freq` creation path consisted in pulling the values out of a tensor, and manually looping over them. This doesn't work if we want to initialize the model on a `meta` device, as shown in the issue above.

This PR replaces the loop by equivalent tensor operations 🤗 

________________________________________________________

### Equivalence check

```py
from transformers import AutoConfig
from transformers.modeling_rope_utils import _compute_default_rope_parameters
import torch
import math

model_repo = "meta-llama/Meta-Llama-3.1-8B-Instruct"
config = AutoConfig.from_pretrained(model_repo)

inv_freq, attention_factor = _compute_default_rope_parameters(config, None, None)

# OLD
factor = config.rope_scaling["factor"]  # `8` in the original implementation
low_freq_factor = config.rope_scaling["low_freq_factor"]  # `1` in the original implementation
high_freq_factor = config.rope_scaling["high_freq_factor"]  # `4` in the original implementation
old_context_len = config.rope_scaling["original_max_position_embeddings"]  # `8192` in the original implementation

low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor
new_freqs = []
for freq in inv_freq:
    wavelen = 2 * math.pi / freq
    if wavelen < high_freq_wavelen:
        new_freqs.append(freq)
    elif wavelen > low_freq_wavelen:
        new_freqs.append(freq / factor)
    else:

        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
        new_freqs.append((1 - smooth) * freq / factor + smooth * freq)

inv_freq_old = torch.tensor(new_freqs, dtype=inv_freq.dtype, device=inv_freq.device)
print(inv_freq_old)

# NEW
factor = config.rope_scaling["factor"]  # `8` in the original implementation
low_freq_factor = config.rope_scaling["low_freq_factor"]  # `1` in the original implementation
high_freq_factor = config.rope_scaling["high_freq_factor"]  # `4` in the original implementation
old_context_len = config.rope_scaling["original_max_position_embeddings"]  # `8192` in the original implementation

low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor
assert low_freq_wavelen != high_freq_wavelen

wavelen = 2 * math.pi / inv_freq
inv_freq_new = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)

smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
inv_freq_new = torch.where(
    ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen),
    (1 - smooth) * inv_freq_new / factor + smooth * inv_freq_new,
    inv_freq_new
)
print(inv_freq_new)

assert torch.allclose(inv_freq_old, inv_freq_new)
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32244). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

I am curious about the return value of `_compute_llama3_parameters`, which still returns `inv_freq` instead of `inv_freq_new`. Is there any in-place operation I missed? 

@lostkevin oh god that was a HUGE failure on my end. Thank you so much for spotting it! 🙏 

That means we are missing a llama 3.1 integration tests, which would have caught this basic mistake. Opening a PR to fix it