transformer_components.py

"""Minimal transformer components for decoder LM training."""

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import torch.nn as nn
from torch import Tensor


@dataclass(frozen=True)
class PositionalInfo:
    pe_type: str
    embeddings: Optional[Tensor] = None
    rope_freqs: Optional[Tuple[Tensor, Tensor]] = None
    rel_embeddings: Optional[Tensor] = None
    apply_to_embeddings: bool = False


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x: Tensor) -> Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x: Tensor) -> Tensor:
        if x.dtype in (torch.float16, torch.bfloat16):
            output = self._norm(x.float()).type_as(x)
        else:
            output = self._norm(x)
        return output * self.weight


class FeedForward(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        dropout: float = 0.0,
        activation: str = "gelu",
        use_bias: bool = True,
    ):
        super().__init__()
        self.activation = activation
        self.use_swiglu = activation == "swiglu"
        self.dropout = nn.Dropout(dropout)
        if self.use_swiglu:
            self.w_gate = nn.Linear(input_dim, hidden_dim, bias=use_bias)
            self.w_up = nn.Linear(input_dim, hidden_dim, bias=use_bias)
            self.w_down = nn.Linear(hidden_dim, input_dim, bias=use_bias)
            self.activation_fn = nn.SiLU()
        else:
            self.linear1 = nn.Linear(input_dim, hidden_dim, bias=use_bias)
            self.linear2 = nn.Linear(hidden_dim, input_dim, bias=use_bias)
            if activation == "gelu":
                self.activation_fn = nn.GELU()
            elif activation == "relu":
                self.activation_fn = nn.ReLU()
            elif activation == "identity":
                self.activation_fn = nn.Identity()
            else:
                raise ValueError(f"Unsupported ffn activation: {activation}")
        self._init_weights()

    def _init_weights(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight, gain=1.0)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, x: Tensor) -> Tensor:
        if self.use_swiglu:
            gate = self.activation_fn(self.w_gate(x))
            up = self.w_up(x)
            x = gate * up
            x = self.dropout(x)
            return self.w_down(x)

        x = self.linear1(x)
        x = self.activation_fn(x)
        x = self.dropout(x)
        x = self.linear2(x)
        return x


class PositionalEncoding(nn.Module):
    def __init__(
        self,
        embedding_dim: int,
        pe_type: str,
        max_len: int,
        theta: float = 10000.0,
        max_rel_pos: Optional[int] = None,
        init_range: float = 0.15,
    ):
        super().__init__()
        if embedding_dim <= 0:
            raise ValueError(f"embedding_dim must be positive, got {embedding_dim}")
        if max_len <= 0:
            raise ValueError(f"max_len must be positive, got {max_len}")
        if pe_type not in {"sinusoidal", "learned", "relative", "rope", "none"}:
            raise ValueError(f"Unsupported pe_type: {pe_type}")
        if theta <= 0.0:
            raise ValueError(f"theta must be positive, got {theta}")
        if init_range <= 0.0:
            raise ValueError(f"init_range must be positive, got {init_range}")

        self.embedding_dim = embedding_dim
        self.pe_type = pe_type
        self.max_len = max_len

        if pe_type == "sinusoidal":
            if embedding_dim % 2 != 0:
                raise ValueError(f"Sinusoidal encoding requires even embedding_dim, got {embedding_dim}")
            pe = self._precompute_sinusoidal_embeddings(
                embedding_dim,
                max_len,
                device=torch.empty(0).device,
            )
            self.register_buffer("sinusoidal_embeddings", pe, persistent=False)
            self._sinusoidal_uninitialized = True
        elif pe_type == "learned":
            self.position_embeddings = nn.Embedding(max_len, embedding_dim)
            nn.init.uniform_(self.position_embeddings.weight, -init_range, init_range)
        elif pe_type == "relative":
            self.max_relative_position = max_rel_pos if max_rel_pos is not None else max_len // 2
            if self.max_relative_position <= 0:
                raise ValueError(
                    f"max_relative_position must be positive, got {self.max_relative_position}"
                )
            num_embeddings = 2 * self.max_relative_position + 1
            self.rel_pos_embeddings_table = nn.Embedding(num_embeddings, embedding_dim)
            nn.init.uniform_(self.rel_pos_embeddings_table.weight, -init_range, init_range)
        elif pe_type == "rope":
            if embedding_dim % 2 != 0:
                raise ValueError(f"RoPE requires even embedding_dim, got {embedding_dim}")
            self.theta = theta
            cos, sin = self._precompute_rope_freqs(embedding_dim, max_len, theta)
            self.register_buffer("rope_cos", cos, persistent=False)
            self.register_buffer("rope_sin", sin, persistent=False)
            self._rope_uninitialized = True

    def get_positional_info(self, seq_len: int, device: torch.device) -> PositionalInfo:
        if self.pe_type == "none":
            return PositionalInfo(pe_type="none")
        if self.pe_type == "sinusoidal":
            if seq_len > self.sinusoidal_embeddings.size(0):
                raise ValueError(
                    f"seq_len {seq_len} exceeds precomputed sinusoidal length {self.sinusoidal_embeddings.size(0)}"
                )
            if (
                getattr(self, "_sinusoidal_uninitialized", False)
                or self.sinusoidal_embeddings.device.type == "meta"
            ):
                embeddings = self._precompute_sinusoidal_embeddings(
                    self.embedding_dim,
                    self.max_len,
                    device=device,
                )
                self.register_buffer("sinusoidal_embeddings", embeddings, persistent=False)
                self._sinusoidal_uninitialized = False
            return PositionalInfo(
                pe_type="sinusoidal",
                embeddings=self.sinusoidal_embeddings[:seq_len].to(device),
                apply_to_embeddings=True,
            )
        if self.pe_type == "learned":
            if seq_len > self.max_len:
                raise ValueError(
                    f"seq_len {seq_len} exceeds maximum learned position length {self.max_len}"
                )
            positions = torch.arange(seq_len, dtype=torch.long, device=device)
            return PositionalInfo(
                pe_type="learned",
                embeddings=self.position_embeddings(positions),
                apply_to_embeddings=True,
            )
        if self.pe_type == "relative":
            return PositionalInfo(
                pe_type="relative",
                rel_embeddings=self._generate_relative_embeddings(seq_len, device),
                apply_to_embeddings=False,
            )
        if seq_len > self.rope_cos.size(0):
            # Dynamically extend RoPE buffer (standard practice for extrapolation)
            cos, sin = self._precompute_rope_freqs(self.embedding_dim, seq_len, self.theta)
            self.register_buffer("rope_cos", cos, persistent=False)
            self.register_buffer("rope_sin", sin, persistent=False)
            self._rope_uninitialized = False
        if getattr(self, "_rope_uninitialized", True) or torch.isnan(self.rope_cos[0, 0]):
            cos, sin = self._precompute_rope_freqs(self.embedding_dim, self.max_len, self.theta)
            self.rope_cos.copy_(cos)
            self.rope_sin.copy_(sin)
            self._rope_uninitialized = False
        return PositionalInfo(
            pe_type="rope",
            rope_freqs=(
                self.rope_cos[:seq_len].to(device),
                self.rope_sin[:seq_len].to(device),
            ),
            apply_to_embeddings=False,
        )

    @staticmethod
    def _rope_inv_freq(dim: int, theta: float, device: torch.device | str) -> Tensor:
        return 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))

    @staticmethod
    def _precompute_sinusoidal_embeddings(
        dim: int,
        end: int,
        device: torch.device | str,
    ) -> Tensor:
        pe = torch.zeros(end, dim, device=device)
        position = torch.arange(0, end, dtype=torch.float, device=device).unsqueeze(1)
        div_term = PositionalEncoding._rope_inv_freq(dim, 10000.0, device=device)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe

    @staticmethod
    def _precompute_rope_freqs(dim: int, end: int, theta: float) -> Tuple[Tensor, Tensor]:
        freqs = PositionalEncoding._rope_inv_freq(dim, theta, device="cpu")
        positions = torch.arange(end, device="cpu")
        freqs = torch.outer(positions, freqs).float()
        return torch.cos(freqs), torch.sin(freqs)

    def _generate_relative_embeddings(self, seq_len: int, device: torch.device) -> Tensor:
        range_q = torch.arange(seq_len, device=device)
        range_k = torch.arange(seq_len, device=device)
        distance_mat = range_k[None, :] - range_q[:, None]
        distance_mat_clipped = torch.clamp(
            distance_mat,
            -self.max_relative_position,
            self.max_relative_position,
        )
        final_indices = distance_mat_clipped + self.max_relative_position
        return self.rel_pos_embeddings_table(final_indices.long())


def reshape_for_broadcast(freqs_cis: Tensor, x: Tensor) -> Tensor:
    ndim = x.ndim
    if ndim < 2:
        raise ValueError(f"Input tensor x must have at least 2 dimensions, got {ndim}")
    if x.shape[-2] != freqs_cis.shape[0] or x.shape[-1] != freqs_cis.shape[-1]:
        raise ValueError(
            f"Shape mismatch for RoPE broadcasting: freqs_cis {freqs_cis.shape}, x {x.shape}"
        )
    shape = [1] * ndim
    shape[-2] = x.shape[-2]
    shape[-1] = x.shape[-1]
    return freqs_cis.view(shape)


def apply_rotary_emb(xq: Tensor, xk: Tensor, freqs_cos: Tensor, freqs_sin: Tensor) -> Tuple[Tensor, Tensor]:
    if xq.shape[-1] % 2 != 0:
        raise ValueError(f"Query feature dimension must be even for RoPE, got {xq.shape[-1]}")
    if xk.shape[-1] % 2 != 0:
        raise ValueError(f"Key feature dimension must be even for RoPE, got {xk.shape[-1]}")

    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)

    freqs_cos_q = reshape_for_broadcast(freqs_cos[: xq.shape[-2]], xq_r)
    freqs_sin_q = reshape_for_broadcast(freqs_sin[: xq.shape[-2]], xq_r)
    freqs_cos_k = reshape_for_broadcast(freqs_cos[: xk.shape[-2]], xk_r)
    freqs_sin_k = reshape_for_broadcast(freqs_sin[: xk.shape[-2]], xk_r)

    xq_out_r = xq_r * freqs_cos_q - xq_i * freqs_sin_q
    xq_out_i = xq_r * freqs_sin_q + xq_i * freqs_cos_q
    xk_out_r = xk_r * freqs_cos_k - xk_i * freqs_sin_k
    xk_out_i = xk_r * freqs_sin_k + xk_i * freqs_cos_k

    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2)
    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2)

    return xq_out.type_as(xq), xk_out.type_as(xk)