Some MHA and RoPE refactoring, llama-3.1 rope_scaling

[francoishernandez]

Notes:

• the RoPE refactoring was probably not strictly necessary, but it makes things clearer IMO;
• models using rotary embeddings will need to be reconverted because all rope related settings were moved to a sub config for clarity;
• this is not extensively tested, but seems to work fine (tested on the example prompt here for instance);
• our RoPE implementation is not strictly equivalent to the HF one, because we rely on some of the "original" rope tricks based on complex space computation (.polar/.real/.imag), where HF just applies .sin()/.cos() methods, which is not numerically equivalent;
• additional scaling types can be implemented (e.g. taking inspiration around here... https://github.com/huggingface/transformers/blob/f1a385b1de7e83e2be9b087d1c0646c0c426e2fc/src/transformers/modeling_rope_utils.py)

[vince62s]

our RoPE implementation is not strictly equivalent to the HF one, because we rely on some of the "original" rope tricks based on complex space computation (.polar/.real/.imag), where HF just applies .sin()/.cos() methods, which is not numerically equivalent

not sure about that.
When the model comes from a Hugging Face format, rotary_interleave is False, and in this case we also use cos/sin, (not the Polar formula)
But maybe I am missing something.

[vince62s]

Did you check the speed performance of this refactoring ? recomputing / applying rotary is quite impactful. My only concern is the reassignment of cos/sin which was performed only when shifting each 32 positions before.

[francoishernandez]

our RoPE implementation is not strictly equivalent to the HF one, because we rely on some of the "original" rope tricks based on complex space computation (.polar/.real/.imag), where HF just applies .sin()/.cos() methods, which is not numerically equivalent

not sure about that. When the model comes from a Hugging Face format, rotary_interleave is False, and in this case we also use cos/sin, (not the Polar formula) But maybe I am missing something.

The .polar trick I'm mentioning is in the initial rope computation:
[main]

eole/eole/modules/multi_headed_attn.py

Line 28 in 9c8cc5b

rope = torch.polar(torch.ones_like(rope), rope)

[refactor]

eole/eole/modules/rope.py

Line 139 in 4eb4853

rope = torch.polar(torch.ones_like(rope), rope)

Afterwards, both the "interleave" and "not" use the real/imag parts of the rope tensor to access cos/sin.

Did you check the speed performance of this refactoring ? recomputing / applying rotary is quite impactful. My only concern is the reassignment of cos/sin which was performed only when shifting each 32 positions before.

Not in depth. It might be worth a look indeed. We can probably have some similar sort of cache to prevent unnecessary recomputations.

[francoishernandez]

d79b6c3 -> similarly to the previous implementation, we pre-compute the rope further than needed to prevent computing at each step.

Not sure about the true impact as my setup might not be the most stable, but it seems that we lost ~5-10% inference speed, which we re-take here.

[vince62s]

what is the benefit of setting position_embedding in transformer_decoder.py, transformer_lm_decoder.py, transformer_encoder.py vs directly in mha.py ?

[francoishernandez]

what is the benefit of setting position_embedding in transformer_decoder.py, transformer_lm_decoder.py, transformer_encoder.py vs directly in mha.py ?

That's a good question. It seemed cleaner to have a single "base" RotaryPosition object to compute rope at a higher level, and use it in all underlying layers/MHA, but that's debatable.

[vince62s]

good to merge.