ROPE implementation differs from official meta implementation 
Author: nreHieW
Pull request: no
Labels: name: duplicate

### System Info

- `transformers` version: 4.26.1
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.9.15
- Huggingface_hub version: 0.13.2
- PyTorch version (GPU?): 1.13.1 (False)
- 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?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker 

### Information

- [ ] 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

```python
# taken from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
def hf_impl(x, head_dim, seq_len, base_theta):
    def rotate_half(x):
        """Rotates half the hidden dims of the input."""
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)
    def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
        """Applies Rotary Position Embedding to the query and key tensors.

        Args:
            q (`torch.Tensor`): The query tensor.
            k (`torch.Tensor`): The key tensor.
            cos (`torch.Tensor`): The cosine part of the rotary embedding.
            sin (`torch.Tensor`): The sine part of the rotary embedding.
            position_ids (`torch.Tensor`, *optional*):
                Deprecated and unused.
            unsqueeze_dim (`int`, *optional*, defaults to 1):
                The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
                sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
                that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
                k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
                cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
                the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
        Returns:
            `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
        """
        cos = cos.unsqueeze(unsqueeze_dim)
        sin = sin.unsqueeze(unsqueeze_dim)
        q_embed = (q * cos) + (rotate_half(q) * sin)
        k_embed = (k * cos) + (rotate_half(k) * sin)
        return q_embed, k_embed
    class LlamaRotaryEmbedding(torch.nn.Module):
        def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
            super().__init__()
            self.scaling_factor = scaling_factor
            self.dim = dim
            self.max_position_embeddings = max_position_embeddings
            self.base = base
            inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
            self.register_buffer("inv_freq", inv_freq, persistent=False)
            # For BC we register cos and sin cached
            self.max_seq_len_cached = max_position_embeddings

        @torch.no_grad()
        def forward(self, x, position_ids):
            # x: [bs, num_attention_heads, seq_len, head_size]
            inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
            position_ids_expanded = position_ids[:, None, :].float()
            # Force float32 since bfloat16 loses precision on long contexts
            # See https://github.com/huggingface/transformers/pull/29285
            device_type = x.device.type
            device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
            with torch.autocast(device_type=device_type, enabled=False):
                freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
                emb = torch.cat((freqs, freqs), dim=-1)
                cos = emb.cos()
                sin = emb.sin()
            return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
    
    x = x.transpose(1, 2)
    emb = LlamaRotaryEmbedding(head_dim, base=base_theta)
    cos_, sin_ = emb(x, torch.arange(seq_len).unsqueeze(0))
    x, _ = apply_rotary_pos_emb(x,x , cos_, sin_)
    return x.transpose(1,2)
```
```python
# taken from https://github.com/meta-llama/llama3/blob/main/llama/model.py
def meta_impl(x, head_dim, seq_len, base_theta):
    def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
        t = torch.arange(end, device=freqs.device, dtype=torch.float32)
        freqs = torch.outer(t, freqs)
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
        return freqs_cis


    def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
        ndim = x.ndim
        assert 0 <= 1 < ndim
        assert freqs_cis.shape == (x.shape[1], x.shape[-1])
        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
        return freqs_cis.view(*shape)


    def apply_rotary_emb(
        xq: torch.Tensor,
        xk: torch.Tensor,
        freqs_cis: torch.Tensor,
    ):
        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
        xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
        freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
        xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
        xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
        return xq_out.type_as(xq), xk_out.type_as(xk)
    freqs_cis = precompute_freqs_cis(head_dim, seq_len, base_theta)
    x, _ = apply_rotary_emb(x, x, freqs_cis=freqs_cis)
    return x
```
```python 
# run below
head_dim = 32
base_theta = 10000.0
seq_len = 100
num_heads = 16
x = torch.randn(1, seq_len, num_heads, head_dim)
a = hf_impl(x, head_dim, seq_len, base_theta)
b = meta_impl(x, head_dim, seq_len, base_theta)
torch.allclose(a, b)
```

### Expected behavior

The implementation of Rope in `transformers` is different from the official meta llama repo [here](https://github.com/meta-llama/llama3/blob/main/llama/model.py)

The reason seems to be how `rotate_half` is used. 
<img width="587" alt="image" src="https://github.com/huggingface/transformers/assets/81154837/de8fd9dc-8b85-4897-9915-833017f1b996">
Per the [paper](https://arxiv.org/pdf/2104.09864), `rotate_half` is used to get the third term of the above equation. 

```python
rotate_half(torch.arange(1,5))
# > tensor([-3, -4,  1,  2])
```
It should however be `-2, 1, -4, 3` per the above equation so 
```python
x = torch.arange(1,5)
rotate_half(x.reshape(*x.shape[:-1], -1, 2)).reshape(x.shape)
# > tensor([-2,  1, -4,  3])
```

Basically chunk the final dimension into pairs, apply `rotate_half` then concat

Would appreciate a second pair of eyes to fact check me on this! Thanks! 

Comments:
#25199 should be everything you need here! 🤗 


> #25199 should be everything you need here! 🤗

ah thanks for the reply @ArthurZucker! Just to clarify my understanding:
- the weights are already converted to accommodate this 
- it is more optimised 

Is this accurate? 



yes sir! 🫡 

Got it thanks for the reply!

Closing 

Sorry to interject here but how exactly can we make sure the Q and K weights are the same if the rotary embeddings are different (i.e. the q and k vectors are rotated to different positions in the HF implementation vs. the Meta implementation)?

yes, that is done in the conversion script