Is there a way to implement RoPE around `nn.MultiheadAttention` somehow?

AsyncTh: <p>I want to implement Rotary Position Embeddings in PyTorch, however it seems like they need to be applied to the output of linear layers before scaled dot-product attention is computed (this is unlike sinusoidal positional encoding, which is applied to word embeddings directly).</p>
<p>I’m wondering if there is still a way implement Rotary Position Embeddings in a way that works with <code>nn.MultiheadAttention</code> and doesn’t require reimplementing multi-head attention.</p>
<p>The reason I want to do this is because <code>nn.MultiheadAttention</code> will likely perform better than my implementation since it now uses FlashAttention, plus it might have some other optimizations I’m not aware of.</p>

J_Johnson: <p>Was working on a PaLM model and using lucidrain’s Pytorch implementation. This makes use of a rotary embedding between a LayerNorm and Linear as part of a transformer block. So in this implementation, it seems to be before the Linear and attention layers.</p>
<p>You can find that code here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/lucidrains/PaLM-rlhf-pytorch/blob/main/palm_rlhf_pytorch/palm.py">
  <header class="source">

      <a href="https://github.com/lucidrains/PaLM-rlhf-pytorch/blob/main/palm_rlhf_pytorch/palm.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/lucidrains/PaLM-rlhf-pytorch/blob/main/palm_rlhf_pytorch/palm.py" target="_blank" rel="noopener nofollow ugc">lucidrains/PaLM-rlhf-pytorch/blob/main/palm_rlhf_pytorch/palm.py</a></h4>


      <pre><code class="lang-py">import math
import copy
from pathlib import Path
from collections import namedtuple
from functools import wraps
from itertools import zip_longest

from tqdm import tqdm
from beartype import beartype
from beartype.typing import Tuple, Optional

import torch
from torch import einsum, nn
import torch.nn.functional as F

from einops import rearrange, repeat, reduce, pack, unpack
from einops.layers.torch import Rearrange, Reduce

from palm_rlhf_pytorch.attention import Attention
from palm_rlhf_pytorch.utils import top_p, top_k, masked_mean, gumbel_sample, eval_decorator
</code></pre>



  This file has been truncated. <a href="https://github.com/lucidrains/PaLM-rlhf-pytorch/blob/main/palm_rlhf_pytorch/palm.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


AsyncTh: <p>If I’m reading the code correctly, line 209 computes outputs of those pre-attention linear layers, so rotary embeddings are applied after linear layers and before attention like in other implementations.</p>

AsyncTh: <p>The answer so far seems to be “no”, but as it turns out I can just use <code>torch.nn.functional.scaled_dot_product_attention</code> to run efficient implementations of SDPA in my custom implementation of multi-head attention, so I guess it makes this question irrelevant. Not sure if I’m losing any performance by not using <code>nn.Transformer</code> though.</p>

JannisZeller: <p>RoPE-ing the <code>q</code> and <code>k</code> inputs to <code>torch.nn.functional.scaled_dot_product_attention</code> is certainly possible, but I think, if one reads the original RoPE-paper (<a href="https://arxiv.org/pdf/2104.09864.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2104.09864.pdf</a>) carefully, they note in eq. (19) the the denominator should not get rotated, to prevent accidental 0-divisions. I am not sure if <code>torch.nn.functional.scaled_dot_product_attention</code> already handles this somehow. So if one wants to stick to the original paper, there is currently (as far as I managed to understand) no other way, than “inserting” the rotations by re-implementing the SDPA with RoPE. Unfortunately then the layer lacks the enhancements through the optimized attention versions except one implements a custom kernel. I have set up a little demo containing the described approach on my <a href="https://github.com/JannisZeller/rope-multi-head-attention" rel="noopener nofollow ugc">github</a>. Feedback welcome! <code>:)</code></p>

srishti-git1110: <p>Hey <a class="mention" href="/u/janniszeller">@JannisZeller</a><br>
Came across your comment here as I was implementing RoPE in PyTorch.</p>
<p>Just wanted to discuss on eq. 19 in the paper: According to my understanding, not rotating q and k for the denominator calculation is for linear self attention only and not for SDPA. Also, went through some implementations by<br>
hugging face and they seem to use softwax as is on the attention scores calculated using rotated q and k implying that q and k in the denominator of the SDPA formula are rotated too.</p>

JannisZeller: <p>Hello <a class="mention" href="/u/srishti-git1110">@srishti-git1110</a>. Thanks for pointing that out. I did not notice a differentiation between linear attention and SDPA, thanks for pointing that out! Still I guess / hope, that I suggested a correct implementation of eq. 19 as it is.</p>