""" RoPE (Rotary Position Embeddings) — Codificação posicional rotacional. Referência: Su et al., 2021 — "RoFormer: Enhanced Transformer with Rotary Position Embedding". Propriedade chave: = → a atenção depende apenas da diferença posicional relativa. Ver docs/MATH_ANALYSIS.md seção 4.1. Autor: CNN-BiGRU Project """ from __future__ import annotations import math from typing import Optional, Tuple import torch import torch.nn as nn def precompute_rope_frequencies( head_dim: int, max_seq_len: int, base: float = 10000.0, device: Optional[torch.device] = None, dtype: torch.dtype = torch.float32, ) -> torch.Tensor: """ Pré-computa as frequências angulares theta_i para RoPE. Fórmula: theta_i = base^(-2i / head_dim) para i = 0, 1, ..., head_dim/2 - 1 Args: head_dim: dimensão de cada cabeça (deve ser par) max_seq_len: comprimento máximo da sequência base: base da exponencial (padrão 10000) device: device do tensor dtype: dtype do tensor Returns: cos_sin: tensor [max_seq_len, head_dim] com [cos(theta_0*m), cos(theta_1*m), ..., sin(theta_0*m), sin(theta_1*m), ...] intercalados por par (formato compatível com apply_rope). Matematicamente: Para cada posição m em [0, max_seq_len) e cada par de dimensões 2i: angle = m * theta_i = m * base^(-2i/head_dim) A matriz de rotação no plano (2i, 2i+1) é: [[cos(angle), -sin(angle)], [sin(angle), cos(angle)]] """ if head_dim % 2 != 0: raise ValueError(f"head_dim deve ser par; recebido {head_dim}") # theta_i = base^(-2i / head_dim), i = 0, ..., head_dim/2 - 1 i = torch.arange(0, head_dim, 2, device=device, dtype=dtype) # [head_dim/2] theta = 1.0 / (base ** (i / head_dim)) # [head_dim/2] # Posições m positions = torch.arange(max_seq_len, device=device, dtype=dtype) # [max_seq_len] # angle[m, i] = m * theta_i angles = torch.outer(positions, theta) # [max_seq_len, head_dim/2] # Duplicar para intercalar: [cos_0, cos_1, ..., sin_0, sin_1, ...] cos = torch.cos(angles) # [max_seq_len, head_dim/2] sin = torch.sin(angles) # Formato final: [max_seq_len, head_dim] com pares (cos_i, sin_i) cos_sin = torch.stack([cos, sin], dim=-1).reshape(max_seq_len, head_dim) return cos_sin def apply_rope( x: torch.Tensor, cos_sin: torch.Tensor, ) -> torch.Tensor: """ Aplica RoPE ao tensor x. Args: x: [batch, n_heads, seq_len, head_dim] cos_sin: [seq_len, head_dim] — tensor pré-computado Returns: x_rotated: [batch, n_heads, seq_len, head_dim] Matemática (para cada batch, head, posição m, par de dims (2i, 2i+1)): x'[2i] = x[2i] * cos(angle) - x[2i+1] * sin(angle) x'[2i+1] = x[2i] * sin(angle) + x[2i+1] * cos(angle) O tensor cos_sin tem formato [seq_len, head_dim] com: cos_sin[m, 2i] = cos(m * theta_i) cos_sin[m, 2i+1] = sin(m * theta_i) """ if x.dim() != 4: raise ValueError( f"x deve ser [batch, n_heads, seq_len, head_dim]; recebido {x.shape}" ) B, H, T, D = x.shape if D % 2 != 0: raise ValueError(f"head_dim deve ser par; recebido {D}") # Reshape x para [B, H, T, D/2, 2] — pares de dimensões x_pairs = x.float().reshape(B, H, T, D // 2, 2) # Reshape cos_sin para [T, D/2, 2] cos_sin_t = cos_sin[:T].to(x.device).float().reshape(T, D // 2, 2) cos = cos_sin_t[..., 0] # [T, D/2] sin = cos_sin_t[..., 1] # Broadcast: [1, 1, T, D/2] cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) # Aplicar rotação x_real = x_pairs[..., 0] # [B, H, T, D/2] x_imag = x_pairs[..., 1] rotated_real = x_real * cos - x_imag * sin rotated_imag = x_real * sin + x_imag * cos # Reagrupar para [B, H, T, D] rotated = torch.stack([rotated_real, rotated_imag], dim=-1).reshape(B, H, T, D) return rotated.to(x.dtype) class RotaryPositionEmbedding(nn.Module): """ Módulo PyTorch para RoPE — mantém o cache de frequências como buffer. Uso: rope = RotaryPositionEmbedding(head_dim=64, max_seq_len=512) # Em forward: q = rope(q) # q: [batch, n_heads, seq_len, head_dim] k = rope(k) """ def __init__(self, head_dim: int, max_seq_len: int = 512, base: float = 10000.0): super().__init__() if head_dim % 2 != 0: raise ValueError(f"head_dim deve ser par; recebido {head_dim}") self.head_dim = head_dim self.max_seq_len = max_seq_len self.base = base # Pré-computar e registrar como buffer (move com .to(device)) cos_sin = precompute_rope_frequencies(head_dim, max_seq_len, base) self.register_buffer("cos_sin", cos_sin, persistent=False) def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor: """ Aplica RoPE a x. Args: x: [batch, n_heads, seq_len, head_dim] offset: deslocamento posicional (para geração autoregressiva com cache KV, onde a posição do novo token é len(cached) + i) Returns: x_rotated: mesmo shape que x """ if x.dim() != 4: raise ValueError( f"x deve ser [batch, n_heads, seq_len, head_dim]; recebido {x.shape}" ) T = x.size(2) if offset + T > self.max_seq_len: # Estender o cache dinamicamente self._extend_cache(offset + T) # Slice das posições relevantes cos_sin = self.cos_sin[offset:offset + T] return apply_rope(x, cos_sin) def _extend_cache(self, new_max: int) -> None: """Estende o cache de frequências para new_max posições.""" new_max_pow2 = 1 while new_max_pow2 < new_max: new_max_pow2 *= 2 new_cos_sin = precompute_rope_frequencies( self.head_dim, new_max_pow2, self.base, device=self.cos_sin.device, dtype=self.cos_sin.dtype ) self.cos_sin = new_cos_sin.to(self.cos_sin.device) self.max_seq_len = new_max_pow2 # Re-registrar como buffer (substitui o anterior) self.register_buffer("cos_sin", self.cos_sin, persistent=False) __all__ = [ "precompute_rope_frequencies", "apply_rope", "RotaryPositionEmbedding", ]