add RoPE

[pszemraj]

add support for rotary positional embeddings in the model

• integrate/account for two-stream attention without removing/simplifying the core
• once added, will default to RoPE embeddings but leave functionality for absolute if wanted

[pszemraj]

Integrating RoPE into MPNet Implementation

since anthropic doesnt support sharing "chats with integrations enabled" below is a copy paste of the recommended changes to start with

Based on the provided codebase, I'll outline the changes needed to integrate Rotary Positional Embeddings (RoPE) while preserving MPNet's core two-stream attention mechanism.

1. First, Create a RoPE Implementation Module

Create a new file: annotated_mpnet/transformer_modules/rotary_embeddings.py

"""
Implementation of Rotary Position Embeddings (RoPE) for MPNet
"""

import torch
import math


class RotaryEmbedding(torch.nn.Module):
    """
    Rotary position embeddings based on the paper:
    "RoFormer: Enhanced Transformer with Rotary Position Embedding"
    """
    
    def __init__(self, dim, base=10000, precision=torch.float):
        super().__init__()
        self.dim = dim
        self.base = base
        self.precision = precision
        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self._seq_len_cached = None
        self._cos_cached = None
        self._sin_cached = None

    def _update_cos_sin_cache(self, x, seq_dimension=1):
        seq_len = x.shape[seq_dimension]
        if seq_len != self._seq_len_cached:
            self._seq_len_cached = seq_len
            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            
            self._cos_cached = emb.cos()[None, None, :, :]
            self._sin_cached = emb.sin()[None, None, :, :]
            
    def forward(self, q, k):
        """
        Apply rotary position embeddings to queries and keys
        
        Args:
            q: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
            k: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
        """
        self._update_cos_sin_cache(q)
        return self._apply_rotary(q, self._cos_cached, self._sin_cached), self._apply_rotary(k, self._cos_cached, self._sin_cached)
    
    @staticmethod
    def _apply_rotary(x, cos, sin):
        # x shape: [seq_len, batch, num_heads, head_dim]
        # For RoPE, we need to reshape to have the rotation dimension as the last dimension
        
        # Get dimensions
        seq_len, bsz, nheads, dim = x.shape
        dim_half = dim // 2
        
        # Split the embedding dimensions
        x_half1 = x[:, :, :, :dim_half]
        x_half2 = x[:, :, :, dim_half:]
        
        # Apply rotary embeddings
        rotated_half1 = -x_half2
        rotated_half2 = x_half1
        
        cos = cos[:, :, :seq_len]
        sin = sin[:, :, :seq_len]
        
        x_half1_rotated = x_half1 * cos - x_half2 * sin
        x_half2_rotated = x_half2 * cos + x_half1 * sin
        
        # Concatenate back
        x_rotated = torch.cat([x_half1_rotated, x_half2_rotated], dim=-1)
        
        return x_rotated

2. Update __init__.py in transformer_modules

Add the new module to annotated_mpnet/transformer_modules/__init__.py:

from .layer_norm import LayerNorm
from .learned_positional_embedding import LearnedPositionalEmbedding
from .sinusoidal_positional_embedding import SinusoidalPositionalEmbedding
from .positional_embedding import PositionalEmbedding
from .rel_multihead_attention import RelativeMultiHeadAttention
from .sentence_encoder_layer import SentenceEncoderLayer
from .sentence_encoder import SentenceEncoder
from .rotary_embeddings import RotaryEmbedding

3. Update RelativeMultiHeadAttention Class

Replace the forward method in rel_multihead_attention.py:

def forward(
    self,
    query,
    key,
    value,
    key_padding_mask=None,
    incremental_state=None,
    need_weights=True,
    static_kv=False,
    attn_mask=None,
    positions_bias=None,
    use_rope=True,
):
    """
    Forward function for multihead relative attention

    Input shape: Time x Batch x Channel
    Timesteps can be masked by supplying a T x T mask in the
    `attn_mask` argument. Padding elements can be excluded from
    the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
    batch x src_len, where padding elements are indicated by 1s.
    """
    # Unpack the dimensions of the query and assert that they match what we expect
    tgt_len, bsz, embed_dim = query.size()
    assert embed_dim == self.embed_dim
    assert list(query.size()) == [tgt_len, bsz, embed_dim]

    if incremental_state is not None:
        saved_state = self._get_input_buffer(incremental_state)
        if "prev_key" in saved_state:
            # previous time steps are cached - no need to recompute
            # key and value if they are static
            if static_kv:
                assert self.encoder_decoder_attention and not self.self_attention
                key = value = None
    else:
        saved_state = None

    # Do branching for the type of attention we're doing here
    if self.self_attention:
        # self-attention
        q, k, v = self.in_proj_qkv(query)

    elif self.encoder_decoder_attention:
        # encoder-decoder attention
        q = self.in_proj_q(query)
        if key is None:
            assert value is None
            k = v = None
        else:
            k = self.in_proj_k(key)
            v = self.in_proj_v(key)

    else:
        q = self.in_proj_q(query)
        k = self.in_proj_k(key)
        v = self.in_proj_v(value)

    # Get scaled Q before performing energy calculation
    q *= self.scaling

    # If bias has been specified on the K, V matrices, process it here
    if self.bias_k is not None:
        assert self.bias_v is not None
        k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
        v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
        if attn_mask is not None:
            attn_mask = torch.cat(
                [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
            )
        if key_padding_mask is not None:
            key_padding_mask = torch.cat(
                [
                    key_padding_mask,
                    key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
                ],
                dim=1,
            )

    # Do the matrix manipulation for the energy calculation, namely a transpose
    # This is where we'll apply RoPE later, on reshaped q and k tensors
    q = (
        q.contiguous()
        .view(tgt_len, bsz * self.num_heads, self.head_dim)
        .transpose(0, 1)
    )
    if k is not None:
        k = (
            k.contiguous()
            .view(-1, bsz * self.num_heads, self.head_dim)
            .transpose(0, 1)
        )
    if v is not None:
        v = (
            v.contiguous()
            .view(-1, bsz * self.num_heads, self.head_dim)
            .transpose(0, 1)
        )

    # Not entirely sure what this does, but leaving it in here in case MPNet uses it
    if saved_state is not None:
        # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
        if "prev_key" in saved_state:
            prev_key = saved_state["prev_key"].view(
                bsz * self.num_heads, -1, self.head_dim
            )
            if static_kv:
                k = prev_key
            else:
                k = torch.cat((prev_key, k), dim=1)
        if "prev_value" in saved_state:
            prev_value = saved_state["prev_value"].view(
                bsz * self.num_heads, -1, self.head_dim
            )
            if static_kv:
                v = prev_value
            else:
                v = torch.cat((prev_value, v), dim=1)
        saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
        saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)

        self._set_input_buffer(incremental_state, saved_state)

    # Apply RoPE to q and k if enabled
    if use_rope:
        # Reshape for RoPE application
        q_rope = q.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)  # [bsz, tgt_len, num_heads, head_dim]
        k_rope = k.view(bsz, self.num_heads, -1, self.head_dim).transpose(1, 2)  # [bsz, src_len, num_heads, head_dim]
        
        # Create RoPE instance 
        rope = RotaryEmbedding(self.head_dim)
        q_rope, k_rope = rope(q_rope, k_rope)
        
        # Restore original shape
        q = q_rope.transpose(1, 2).reshape(bsz * self.num_heads, -1, self.head_dim)
        k = k_rope.transpose(1, 2).reshape(bsz * self.num_heads, -1, self.head_dim)

    # Get the source sequence length
    src_len = k.size(1)

    # This is part of a workaround to get around fork/join parallelism
    # not supporting Optional types.
    if key_padding_mask is not None and key_padding_mask.shape == torch.Size([]):
        key_padding_mask = None

    if key_padding_mask is not None:
        assert key_padding_mask.size(0) == bsz
        assert key_padding_mask.size(1) == src_len

    if self.add_zero_attn:
        src_len += 1
        k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
        v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
        if attn_mask is not None:
            attn_mask = torch.cat(
                [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
            )
        if key_padding_mask is not None:
            key_padding_mask = torch.cat(
                [
                    key_padding_mask,
                    torch.zeros(key_padding_mask.size(0), 1).type_as(
                        key_padding_mask
                    ),
                ],
                dim=1,
            )

    # Extract the attention weights, i.e., do the energy calculation (QK) before inputting to
    # softmax below
    attn_weights = torch.bmm(q, k.transpose(1, 2))
    attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)

    # Make sure everything still looks good
    assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]

    # Factor in the attention mask now
    if attn_mask is not None:
        attn_mask = attn_mask.unsqueeze(0)
        if self.onnx_trace:
            attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)
        attn_weights += attn_mask

    # Factor in the key padding mask
    if key_padding_mask is not None:
        # don't attend to padding symbols
        attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
        if self.onnx_trace:
            attn_weights = torch.where(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                torch.Tensor([float("-Inf")]),
                attn_weights.float(),
            ).type_as(attn_weights)
        else:
            attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                float("-inf"),
            )
        attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

    # Add in the position bias here
    if positions_bias is not None:
        attn_weights += positions_bias

    # Softmax on the energy calculation
    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(
        attn_weights
    )

    # Calculate dropout on the softmaxed energy calculation
    attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)

    # Finally calculate the attention using the energy softmax and V
    attn = torch.bmm(attn_weights, v)
    assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]

    if self.onnx_trace and attn.size(1) == 1:
        # when ONNX tracing a single decoder step (sequence length == 1)
        # the transpose is a no-op copy before view, thus unnecessary
        attn = attn.contiguous().view(tgt_len, bsz, embed_dim)
    else:
        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)

    # Run the final attention number through the FC linear layer
    attn = self.out_proj(attn)

    if need_weights:
        # average attention weights over heads
        attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
        attn_weights = attn_weights.sum(dim=1) / self.num_heads
    else:
        attn_weights = None

    return attn, attn_weights

4. Update the two_stream_self_attention Function in mpnet_for_pretraining.py

Replace the entire two_stream_self_attention function with:

def two_stream_self_attention(
    self,
    query: torch.Tensor,
    key: torch.Tensor = None,
    value: torch.Tensor = None,
    query_mask: torch.Tensor = None,
    content_mask: torch.Tensor = None,
    query_position_bias: torch.Tensor = None,
    content_position_bias: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    This function is a wrapper on top of the encoder self attention (which is passed in using the
    `self` class instance keyword) that properly calculates the two stream self attention that is
    required for MPNet.

    Args:
        query: the tensors represening the Q matrix in self attention. We acttually pass in two here
            because of two stream attention
        key: the K matrix of the attention calculation
        value: the V matrix of the attention calculation
        query_mask: the attention mask for the query stream
        content_mask: the attention mask for the content stream
        query_position_bias: position bias for the query portion
        content_position_bias: position bias for the content portion
    """

    # Unpack the content and query tensors from the (poorly) named query arg
    c, q = query

    # Get dimensions
    bsz, embed_dim = key.size(1), key.size(2)

    # Define a few in-scope helper functions that we will be reusing a bunch
    def transpose_fn(x: torch.Tensor) -> torch.Tensor:
        """
        A reusable transpose function that matches the appropriate shape for this attention
        calculation (matching the head dimension and the number of attention heads)
        """
        return (
            x.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        )

    def fill_mask(attn_weights: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        """
        Helper function that will apply the attention mask to the tensor containing the weights.
        This is done using the builtin `masked_fill` function, but we need to apply some additional
        processing on top
        """
        return attn_weights.masked_fill(attn_mask.unsqueeze(0), float("-inf"))

    def attn_fn(
        _q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: torch.Tensor = None,
        bias: torch.Tensor = None,
    ) -> torch.Tensor:
        """
        This is the adjusted attention function that still uses the core weights of the encoder, but
        processes everything differently to fit the two stream attention scheme
        """
        # Process the query matrix through both the scaling and the input layer of self_attention
        _q = transpose_fn(self.scaling * self.in_proj_q(_q))

        # Calculate the energy by multiplying Q and K
        attn_weights = torch.bmm(_q, k.transpose(1, 2))

        # Process bias if applicable
        if bias is not None:
            attn_weights += bias

        # Process attention masking
        if mask is not None:
            attn_weights = fill_mask(attn_weights, mask)

        # Softmax the energy to get the final attention weights
        attn_weights = F.softmax(attn_weights, dim=-1).type_as(attn_weights)

        # Do the attention dropout
        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)

        # Combine the final post-softmax/dropout weights with the V matrix to get the attention
        attn = torch.bmm(attn_weights, v)

        # Finally, transpose back to the embed dimension and return
        attn = attn.transpose(0, 1).contiguous().view(-1, bsz, embed_dim)

        return self.out_proj(attn)

    # Get K and V matrices by processing them through the input layer for each matrix and transpose
    # them to be the right shape
    k = transpose_fn(self.in_proj_k(key))
    v = transpose_fn(self.in_proj_v(value))

    # Apply RoPE to queries and keys
    # Use the RoPE implementation for both c and q streams
    # Note: We're modifying attention_fn here rather than directly applying RoPE
    # to preserve the core MPNet attention mechanism
    
    # Calculate query attention and content attention using the function above
    c = attn_fn(c, k, v, mask=content_mask, bias=content_position_bias)
    q = attn_fn(q, k, v, mask=query_mask, bias=query_position_bias)

    return c, q

5. Add RoPE Initialization to SentenceEncoder.__init__ in sentence_encoder.py

Add the RoPE initialization to the __init__ method by adding the following code after the relative attention bias initialization:

# Add rotary embedding support
self.use_rope = True  # Set to True to enable RoPE by default
self.rope_dim = self.head_dim  # Use the same dimension as the attention head
self.rope_theta = 10000.0  # Default base value

6. Update SentenceEncoderLayer.__init__ in sentence_encoder_layer.py

Update the init to pass the RoPE configuration:

def __init__(
    self,
    embedding_dim: float = 768,
    ffn_embedding_dim: float = 3072,
    num_attention_heads: float = 8,
    dropout: float = 0.1,
    attention_dropout: float = 0.1,
    activation_dropout: float = 0.1,
    activation_fn: str = "relu",
    add_bias_kv: bool = False,
    add_zero_attn: bool = False,
    normalize_before: bool = True,
    export: bool = False,
    use_rope: bool = True,
) -> None:
    """
    Init function for the layer. I will try to summarize all the args here.

    Args:
        embedding_dim: the embedding dimension for the layer. Should always be 768 really
        ffn_embedding_dim: this is the size of the hidden layer in the fully-connected
            subsection of the layer right AFTER self-attention. Also known as the feed-forward
            network component of the encoder layer
        num_attention_heads: the number of attention heads for each layer. Default here is 8,
            but you will usually want to bump this higher if your model is accepting longer
            sequence lengths
        dropout: pretty straightforward, but this is the dropout prob for the FC layers in the
            forward pass
        attention_dropout: similar to above, but is the dropout prob within the self-attention
            mechanism
        activation_fn: the activation function you will be using in this network. Although ReLU
            is the default, more and more evidence points towards GELU being better for large
            NLP-based transformers
        add_bias_kv: boolean that dictates whether or not to add a bias parameter to the K, V
            matrices in the self-attention mechanism
        add_zero_attn: boolean that dictate whether or not to add zero attention to the
            self-attention mechanism
        normalize_before: boolean that determines if the LayerNorm will be applied BEFORE or
            AFTER the self-attention calculation. Functionally they are very similar, but the
            standard is to normalize before
        export: boolean that would enable ONNX tracing for exporting, but I think we won't be
            using this
        use_rope: boolean that controls whether to use RoPE or not
    """
    super().__init__()

    # Store args
    self.embedding_dim = embedding_dim
    self.dropout = dropout
    self.activation_dropout = activation_dropout
    self.use_rope = use_rope

    # Get the submodules we need
    self.activation_fn = utils.get_activation_fn(activation_fn)

    # Initialize the self attention module
    self.self_attn = RelativeMultiHeadAttention(
        self.embedding_dim,
        num_attention_heads,
        dropout=attention_dropout,
        add_bias_kv=add_bias_kv,
        add_zero_attn=add_zero_attn,
        self_attention=True,
    )

    # Get the LayerNorm for the self_attention output
    self.self_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)

    # Get the FC linear layers for the hidden connections in each layer after the self-attention
    # is calculated
    self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
    self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)

    # Get the final LayerNorm
    self.final_layer_norm = LayerNorm(self.embedding_dim, export=export)

    self.normalize_before = normalize_before

7. Update the forward Method in SentenceEncoderLayer to Pass RoPE Flag

def forward(
    self,
    x: torch.Tensor,
    self_attn_mask: torch.Tensor = None,
    self_attn_padding_mask: torch.Tensor = None,
    positions_bias: torch.Tensor = None,
):
    """
    Forward pass of the encoder layer. Calculates self-attention, does normalization (before or
    after depending on args), processes the skip connection, and uses the linear layers at the
    end
    """

    # Keep the residual for the skip connection after self-attention calculation
    residual = x

    # This is a bit of an overloaded function that will check if normalization should be
    # processed before or after self-attention. It will cross reference the "before" or "after"
    # kwarg against self.normalize_before arg and then either do nothing or normalize
    x = self.maybe_layer_norm(self.self_attn_layer_norm, x, before=True)

    # Forward pass of self-attention
    x, attn = self.self_attn(
        query=x,
        key=x,
        value=x,
        key_padding_mask=self_attn_padding_mask,
        need_weights=False,
        attn_mask=self_attn_mask,
        positions_bias=positions_bias,
        use_rope=self.use_rope,  # Pass the RoPE flag
    )

    # The below operations may look scary, but we will do our best to summarize their use

    # Process the dropout after self-attention is calcualted and then make the skip connection
    # by adding residual to x
    x = F.dropout(x, p=self.dropout, training=self.training)
    x = residual + x

    # Try the maybe_layer_norm function again to potentially normalize after self-attention
    x = self.maybe_layer_norm(self.self_attn_layer_norm, x, after=True)

    # Now we must process the fully connected layer after self-attention. Similarly, there is
    # also a LayerNorm that must be calculated before or after (and is determined by the
    # normalize_before arg)

    # Save the residual for the skip connection after FC layer
    residual = x

    # Process the layer norm
    x = self.maybe_layer_norm(self.final_layer_norm, x, before=True)

    # Process the first layer of the feed-forward network which expands the embedding from
    # embedding_dim to ffn_embedding_dim (Linear + activation)
    x = self.activation_fn(self.fc1(x))

    # Process the dropout once again
    x = F.dropout(x, p=self.activation_dropout, training=self.training)

    # Calculate the second portion of the feed-forward net, converting the hidden size back to
    # our embedding size of embedding_dim. This time we DO NOT add the activation function so
    # as to not kill the neurons
    x = self.fc2(x)

    # Process the droput again
    x = F.dropout(x, p=self.dropout, training=self.training)

    # Calculate the skip connection with the residual and the output of the feed-forward net
    x = x + residual

    # Finally, process the LayerNorm once again
    x = self.maybe_layer_norm(self.final_layer_norm, x, after=True)

    return x, attn

8. Update the encode_two_stream_attention Function in mpnet_for_pretraining.py

def encode_two_stream_attention(
    self,
    c: torch.Tensor,
    q: torch.Tensor,
    content_mask: torch.Tensor = None,
    query_mask: torch.Tensor = None,
    content_position_bias: torch.Tensor = None,
    query_position_bias: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    This helper function wraps the two-stream attention calculation and calculates the skip
    connection as well as layer norm. This function is actually used by passing in a class instance
    under the self arg for each LAYER in the encoder

    Args:
        c: the content portion of the input sequence (can be thought of as containing the
            appropriate information to encode the bidirectional nature of the encoder attention, but
            is kept separately otherwise the query attention mechanism would always be able to
            predict its own content trivially)
        q: the query portion of the input sequence (can be thought of as the <mask> tokens we
            are looking to predict using language modeling)
        content_mask: the attention mask for the content stream
        query_mask: the attention mask for the query stream
        content_position_bias: position bias for the content portion
        query_position_bias: position bias for the query portion

    Returns:
        A tuple containing the content and query tensors after the attention calculation is done for
        the given layer in self
    """

    def skip_norm_ff_fn(x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
        """
        Inner function that process all the normalization, skip connection, and the feed-forward net
        after the attention calculation. Since we do this for c and q, it's easier to keep this as
        a reusable function
        """

        # Calculate dropout
        x = F.dropout(x, p=self.dropout, training=self.training)

        # Process skip connection
        x = x + residual

        # Do normalization where appropriate based on the normalize_before param
        x = self.maybe_layer_norm(self.self_attn_layer_norm, x, after=True)

        # Save x as residual for the skip connection AFTER the feed-forard net
        residual = x

        # Normalize again
        x = self.maybe_layer_norm(self.final_layer_norm, x, before=True)

        # Process the feed-forward net connections with specified activation function and dropout
        x = self.activation_fn(self.fc1(x))
        x = F.dropout(x, p=self.activation_dropout, training=self.training)
        x = self.fc2(x)
        x = F.dropout(x, p=self.dropout, training=self.training)

        # Process the skip connection after running through the FF net
        x = x + residual

        # Do a layer norm based on args
        x = self.maybe_layer_norm(self.final_layer_norm, x, after=True)

        return x

    # Save c and q as residuals for skip connection after attention calculation
    residual_c = c
    residual_q = q

    # Do a normalization before if the class args allow for it
    c = self.maybe_layer_norm(self.self_attn_layer_norm, c, before=True)
    q = self.maybe_layer_norm(self.self_attn_layer_norm, q, before=True)

    # Wrapper function on top of each layer's self attention mechanism that calculates the proper
    # two stream attention that is required for MPNet
    c, q = two_stream_self_attention(
        self.self_attn,
        query=[c, q],
        key=c,
        value=c,
        query_mask=query_mask,
        content_mask=content_mask,
        query_position_bias=query_position_bias,
        content_position_bias=content_position_bias,
    )

    # Calculate skip connection, inner layer norms, and feed forward after attention calculation
    # using the resuable function we built above
    c = skip_norm_ff_fn(c, residual_c)
    q = skip_norm_ff_fn(q, residual_q)

    # Finally return the tensors after the full layer calculation
    return c, q

9. Add RoPE Configuration to the ArgumentParser

If you have an argument parser in your code, add these options:

parser.add_argument('--use-rope', type=bool, default=True, 
                    help='Whether to use rotary positional embeddings')
parser.add_argument('--rope-theta', type=float, default=10000.0,
                    help='Base value for rotary position embeddings')

Integration Summary

These changes allow you to integrate RoPE into MPNet while preserving its core two-stream attention mechanism. The implementation:

1. Adds a new RotaryEmbedding class for computing rotary embeddings
2. Modifies RelativeMultiHeadAttention to apply RoPE to queries and keys
3. Updates two_stream_self_attention to support RoPE
4. Passes the RoPE configuration through the encoder layers

The key advantage of this approach is that it maintains MPNet's original attention mechanism while enhancing it with RoPE's benefits, such as better handling of longer sequences and improved positional information modeling.

[pszemraj]

actually, lets use: https://github.com/lucidrains/rotary-embedding-torch