We want to add a method which is called immediately after RoPE (apply_rotary_pos_emb) in Llama's attention module.
Key considerations:
[Con] We're copying verbatim ~60-80 lines of forward code for each attention module (3 for Llama-2).
[Pro] @sashkboos thinks we'll be able to avoid some unnecessary expensive reshaping when applying the hadamards by having access to this part of the code.
We want to add a method which is called immediately after RoPE (
apply_rotary_pos_emb
) in Llama's attention module.Key considerations:
forward
code for each attention module (3 for Llama-2).