Add Rotary Positional Embeddings (RoPE) - part 2 of parallel attention blocks

tests/modules/layers/test_position_embedding.py

@@ -4,13 +4,16 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import pytest
+import math
 
+import pytest
 import torch
 from tests.test_utils import assert_expected
 from torch import nn
+
 from torchmultimodal.modules.layers.position_embedding import (
  BroadcastedPositionEmbedding,
+ RotaryPositionalEmbeddings,
  SinusoidalPositionEmbeddings,
 )
 
@@ -112,3 +115,38 @@ def test_forward(self, data, emb):
  actual = emb(data)
  expected = torch.Size([3, 5])
  assert_expected(actual.shape, expected)
+
+
+def test_rotary_embeddings_math():
+ q = (
+ torch.tensor([[1, 0], [1, 0]], dtype=torch.float).unsqueeze(0).unsqueeze(0)
+ ) # b h s e
+
+ k = 2 * torch.tensor([[1, 0], [1, 0]], dtype=torch.float).unsqueeze(0).unsqueeze(
+ 0
+ ) # b h s e
+
+ rotary_embeddings = RotaryPositionalEmbeddings(2, 2, 1)
+ qr, kr = rotary_embeddings(q, k, 0)
+ rot0 = torch.tensor([[math.cos(0), -math.sin(0)], [math.sin(0), math.cos(0)]])
+ rot1 = torch.tensor([[math.cos(1), -math.sin(1)], [math.sin(1), math.cos(1)]])
+
+ assert_expected(torch.matmul(rot0, q[..., 0, :].squeeze()), qr[..., 0, :].squeeze())
+ assert_expected(torch.matmul(rot1, q[..., 1, :].squeeze()), qr[..., 1, :].squeeze())
+ assert_expected(torch.matmul(rot0, k[..., 0, :].squeeze()), kr[..., 0, :].squeeze())
+ assert_expected(torch.matmul(rot1, k[..., 1, :].squeeze()), kr[..., 1, :].squeeze())
+
+
+def test_rotary_embeddings_left_padding():
+ q = torch.ones(2, 1, 4, 16, dtype=torch.float) # b h s e
+ k = 2 * torch.ones(2, 1, 4, 16, dtype=torch.float) # b h s e
+ rotary_embeddings = RotaryPositionalEmbeddings(16, 32)
+
+ qr, kr = rotary_embeddings(q, k, 0)
+ qr2, kr2 = rotary_embeddings(q, k, torch.tensor([0, 1]))
+
+ assert_expected(qr[0], qr2[0])
+ assert_expected(qr[0, :, 1], qr2[1, :, 0])
+
+ assert_expected(kr[0], kr2[0])
+ assert_expected(kr[0, :, 1], kr2[1, :, 0])

torchmultimodal/modules/layers/position_embedding.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import itertools
-from typing import Tuple
+from typing import Tuple, Union
 
 import torch
 from torch import nn, Tensor
@@ -169,3 +169,108 @@ def forward(self, t: Tensor) -> Tensor:
  if self.embed_dim % 2 == 1:
  embeddings = nn.functional.pad(embeddings, (0, 1))
  return embeddings
+
+
+class RotaryPositionalEmbeddings(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ max_position_embeddings: Union[int, float] = 2048,
+ ratio: int = 10000,
+ device: torch.device = None,
+ ):
+ """
+ Implements Rotary Positional Embeddings (RoPE)
+ proposed in: https://arxiv.org/abs/2104.09864
+
+ Args
+ ----
+ dim : int
+ Per-head embedding dimension
+ max_position_embeddings : int
+ Maximum expected sequence length for the model, if exceeded the cached freqs will be recomputed
+ ratio: int
+ The ratio for the geometric progression to compute the rotation angles
+ """
+ super().__init__()
+ self.register_buffer(
+ "freqs",
+ 1.0
+ / (
+ ratio
+ ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim)
+ ),
+ )
+ self.compute_freqs_cis(max_position_embeddings)
+
+ def compute_freqs_cis(
+ self, max_position_embeddings: Union[int, float] = 2048
+ ) -> None:
+ t = torch.arange(
+ max_position_embeddings, device=self.freqs.device, dtype=self.freqs.dtype
+ )
+ freqs = torch.outer(t, self.freqs).float()
+ self.max_seq_len_cached = max_position_embeddings
+ self.register_buffer(
+ "cached_freqs",
+ torch.stack(
+ [
+ torch.cos(freqs),
+ -torch.sin(freqs),
+ torch.sin(freqs),
+ torch.cos(freqs),
+ ],
+ dim=2,
+ ).view(*freqs.shape, 2, 2),
+ )
+
+ def reshape_for_broadcast(
+ self, x: torch.Tensor, cur_freqs: torch.Tensor
+ ) -> torch.Tensor:
+ ndim = x.ndim
+ assert 1 < ndim
+ assert cur_freqs.shape[:2] == (x.shape[2], x.shape[-2])
+ shape = [d if i == 2 or i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]
+ return cur_freqs.view(*shape, 2)
+
+ def forward(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ start_pos: Union[int, float, torch.LongTensor],
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args
+ ----
+ q : torch.Tensor
+ Embedded query tensor, expected size is B x H x S x Eh
+ k : torch.Tensor
+ Embedded query tensor, expected size is B x H x S x Eh
+ start_pos : Union[int, torch.LongTensor]
+ The starting position of the tokens encoded in q and k. This is important in
+ kv-caching and left-padding situations, for which the rotation to be applied might
+ not always be the pre-cached position 0...S. For kv-caching without dynamic batching
+ start_pos is shared for all the batch.
+ """
+ seq_len = q.shape[2]
+ q_ = q.float().reshape(*q.shape[:-1], -1, 2) # B H L D/2 2
+ k_ = k.float().reshape(*k.shape[:-1], -1, 2) # B H L D/2 2
+
+ if isinstance(start_pos, int):
+ if start_pos + seq_len > self.max_seq_len_cached:
+ self.compute_freqs_cis(start_pos + seq_len)
+ cur_freqs = self.cached_freqs[start_pos : start_pos + seq_len]
+ freqs = self.reshape_for_broadcast(q_, cur_freqs)
+ else:
+ max_start_pos = torch.max(start_pos).item()
+ if max_start_pos + seq_len > self.max_seq_len_cached:
+ self.compute_freqs_cis(max_start_pos + seq_len)
+ freqs_idxs = torch.arange(0, seq_len, dtype=torch.long).repeat(
+ start_pos.shape[0]
+ ).view(-1, seq_len) + start_pos.view(-1, 1)
+ freqs = self.cached_freqs[freqs_idxs].unsqueeze(1)
+
+ freqs = freqs.float() # 1 1 L D/2 2 2
+ q_out = freqs.mul(q_.unsqueeze(-2)).sum(5).flatten(3)
+ k_out = freqs.mul(k_.unsqueeze(-2)).sum(5).flatten(3)
+ return q_out.type_as(q).contiguous(), k_out.type_as(k).contiguous()