File size: 2,619 Bytes
f8ccbd6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from __future__ import annotations
import math
import torch
from torch import nn
class RotaryEmbedding(nn.Module):
def __init__(self, d_head: int, max_seq_len: int, theta: float = 10_000.0) -> None:
super().__init__()
if d_head < 2 or d_head % 2 != 0:
raise ValueError("RoPE requires a positive, even attention head dimension.")
if max_seq_len < 1:
raise ValueError("max_seq_len must be positive.")
theta = float(theta)
if not math.isfinite(theta) or theta <= 0:
raise ValueError("RoPE theta must be positive.")
self.d_head = d_head
self.max_seq_len = max_seq_len
self.theta = theta
self.rope_dim = d_head // 2
self.register_buffer("_cos", torch.empty(0), persistent=False)
self.register_buffer("_sin", torch.empty(0), persistent=False)
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
if query.shape != key.shape:
raise ValueError("RoPE query and key tensors must have the same shape.")
return self.rotate(query), self.rotate(key)
def rotate(self, x: torch.Tensor) -> torch.Tensor:
if x.ndim != 4 or x.shape[-1] != self.d_head:
raise ValueError(f"RoPE expects shape [batch, heads, sequence, {self.d_head}].")
seq_len = x.shape[-2]
if seq_len > self.max_seq_len:
raise ValueError(f"Sequence length {seq_len} exceeds RoPE limit {self.max_seq_len}.")
cos, sin = self._cos_sin(x.device)
cos = cos[:, :, :seq_len].to(dtype=x.dtype)
sin = sin[:, :, :seq_len].to(dtype=x.dtype)
first_half = x[..., : self.rope_dim]
second_half = x[..., self.rope_dim :]
return torch.cat(
(
cos * second_half + sin * first_half,
-sin * second_half + cos * first_half,
),
dim=-1,
)
def _cos_sin(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
if self._cos.numel() == 0 or self._cos.device != device:
inverse_frequencies = self.theta ** (
-torch.arange(self.rope_dim, dtype=torch.float32, device=device) / self.rope_dim
)
frequencies = torch.outer(
torch.arange(self.max_seq_len, dtype=torch.float32, device=device),
inverse_frequencies,
)
self._cos = frequencies.cos()[None, None, :, :]
self._sin = frequencies.sin()[None, None, :, :]
return self._cos, self._sin
|