Inconsistency in Llama RoPE implementation?
Author: tigranfah
Pull request: no
Labels: name: Discussion

name: Core: Modeling

name: bug

### System Info

transformers==4.45.1

### Who can help?

@ArthurZucker

### Reproduction

I came across something unexpected with llama RoPE implementation, specifically in `apply_rotary_pos_emb` function in `modeling_llama.py`

```python
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    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
```

The meta-llama/llama3 [implementation](https://github.com/meta-llama/llama3/blob/main/llama/model.py) from meta uses the function `apply_rotary_emb` and complex numbers to perform to dot product with the rotation matrix, but I am not sure I fully understand huggingface's implementation.

```python
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, 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)
```


Currently, I am running llama 3 with PyTorch's [torchtitan](https://github.com/pytorch/torchtitan) codebase and with huggingface and comparing the logits.

I have the following script where the `model` is the torchtitan's llama 3 implementation

```python3
from transformers import AutoModelForCausalLM, AutoTokenizer
weights_path = 'meta-llama/Llama-3.2-1B'
hf_model = AutoModelForCausalLM.from_pretrained(weights_path)
tok = AutoTokenizer.from_pretrained(weights_path)

device = "cuda"
hf_model.to(device)

model.eval()

text = "Hello world"
data = tok(text, return_tensors="pt").to(device)
hf_logits = hf_model(**data).logits
logits = model(data.input_ids)
print(torch.allclose(hf_logits, logits, atol=1e-4)) # False
```

I am going into the forward passes for both huggingface's implementation and torchtitan's implementation and the query's and key's embeddings don't match between the two implementations after the rotary positional embedding is applied.

However if I change the implementation of `apply_rotary_pos_emb` (in `modeling_llama.py`) to the one given below, which tries to implement the rotation matrix multiplication in the RoPE paper, I get almost an exact match between the logits of the torchtitan's implementation and huggingface's implementation.

```python
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    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

    cos = cos[..., :cos.shape[-1] // 2]
    sin = sin[..., :sin.shape[-1] // 2]
    q_embed = torch.empty_like(q, device=q.device)
    q_embed[..., ::2] = (q[..., ::2] * cos) - (q[..., 1::2] * sin)
    q_embed[..., 1::2] = (q[..., ::2] * sin) + (q[..., 1::2] * cos)
    k_embed = torch.empty_like(k, device=k.device)
    k_embed[..., ::2] = (k[..., ::2] * cos) - (k[..., 1::2] * sin)
    k_embed[..., 1::2] = (k[..., ::2] * sin) + (k[..., 1::2] * cos)
    return q_embed, k_embed
```

### Expected behavior

Ideally I would expect to get the logits matching between the two implementations.

Please let me if I am missing something or there is actually an issue.

Comments:
Hey, https://github.com/huggingface/transformers/issues/25199 should answer your question :)

Thanks for pointing it out; i guess you can also implement the reverse of premute function to convert form huggingface to meta's implementation.

```python
# reverse_permute for sliced rotary
def reverse_permute(w, n_heads, dim1, dim2):
    return w.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2)
```

I have the logits matching with this. Thanks.