Bug in RoPE

[gvlassis]

Hello!

I think your RoPE implementation is wrong. The get_sinusoidal_embeddings() function looks correct, but there are two problems with the apply_rotary_position_embeddings() function:

1. sinusoidal_emb contains sines in even indices and cosines in odd indices. But sin, cos = sinusoidal_pos.chunk(2, dim=-1) extracts the first and second half of sinusoidal_emb, not the even and odd indices (i.e. there is a problem with the unpacking)
2. I understand all the tensor manipulations you are doing, but you still never add cosines with sines (which you should do at some point in the canonical RoPE implementation)

ChatGPT seems to agree with me (detects the same two problems even when I do not hint them), but it could be a misunderstanding on my part.

@sashkboos
@iloshchilov
@borisgin

[tjcashman]

I agree; thanks for sharing your observations!

One possibility is that the multiplication in apply_rotary_position_embeddings might have been intended to be a complex multiplication, as in the reference implementation for RoPE at https://github.com/meta-llama/llama/blob/689c7f261b9c5514636ecc3c5fefefcbb3e6eed7/llama/model.py#L159

Here's an alternative version of the function that I think addresses both of the issues you noticed, with minimal modifications: 44e7ca2
Note that that in apply_rotary_position_embeddings:

torch.reshape(q_rot, q.shape[:-1] + (q.shape[-1]//2, 2))

is a no-op, as the shape of q_rot is already the shape requested (so I removed the reshape just to make it easier to visually compare apply_rotary_position_embeddingsand apply_rotary_position_embeddings_complex).

The only other change in this commit is to swap the order of the cos and sin packing into sinusoidal_emb; I suspect this makes little difference in practice, but it does make the conversion to complex consistent with the polar form intended in get_sinusoidal_embeddings.