""" Core Transformer layers. Implemented: - RMSNorm - Rotary Positional Embedding - Causal Self-Attention - SwiGLU MLP - Pre-Norm Transformer block """ import math from typing import Tuple import torch import torch.nn as nn import torch.nn.functional as F from src.model.config import ModelConfig class RMSNorm(nn.Module): """ RMSNorm from Zhang & Sennrich. Difference from LayerNorm: - Does not subtract mean. - Normalizes by root mean square. - Usually cheaper and stable for LLMs. Input: x: Tensor of shape [batch, seq, hidden] Output: Normalized tensor of same shape. """ def __init__(self, dim: int, eps: float): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: # Upcast to float32 for the variance computation: under AMP (bf16/fp16) # the mean-of-squares can lose precision / overflow, which destabilizes # normalization. LLaMA-style: normalize in fp32, cast back, then scale. orig_dtype = x.dtype x_fp32 = x.float() rms = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps) return self.weight * (x_fp32 * rms).to(orig_dtype) class RotaryEmbedding(nn.Module): """ Rotary Positional Embedding. RoPE rotates query/key vectors based on token position. It avoids learned absolute positional embeddings. Input: q/k tensor: [batch, heads, seq, head_dim] Output: position-rotated q/k tensor. """ def __init__(self, head_dim: int, theta: float = 10000.0): super().__init__() if head_dim % 2 != 0: raise ValueError("RoPE requires even head_dim.") inv_freq = 1.0 / ( theta ** (torch.arange(0, head_dim, 2).float() / head_dim) ) self.register_buffer("inv_freq", inv_freq, persistent=False) def get_cos_sin( self, seq_len: int, device: torch.device, dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Input: seq_len: Sequence length. device: CPU/GPU. dtype: Target dtype. Output: cos, sin tensors shaped [1, 1, seq, head_dim/2] """ positions = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.einsum("i,j->ij", positions, self.inv_freq.to(device)) cos = freqs.cos()[None, None, :, :].to(dtype) sin = freqs.sin()[None, None, :, :].to(dtype) return cos, sin @staticmethod def apply_rope( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: """ Input: x: [batch, heads, seq, head_dim] cos/sin: [1, 1, seq, head_dim/2] Output: RoPE-applied tensor with same shape. """ # Store original dtype orig_dtype = x.dtype # Upcast to float32 for precision stability during rotation x = x.float() cos = cos.float() sin = sin.float() x_even = x[..., 0::2] x_odd = x[..., 1::2] rotated_even = x_even * cos - x_odd * sin rotated_odd = x_even * sin + x_odd * cos out = torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2) return out.to(orig_dtype) def forward( self, q: torch.Tensor, k: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: seq_len = q.size(2) cos, sin = self.get_cos_sin(seq_len, q.device, q.dtype) q = self.apply_rope(q, cos, sin) k = self.apply_rope(k, cos, sin) return q, k class CausalSelfAttention(nn.Module): """ Multi-head causal self-attention. Uses PyTorch scaled_dot_product_attention. This still counts as your own model architecture; the CUDA kernel is optimized internally by PyTorch. Input: x: [batch, seq, hidden] Output: [batch, seq, hidden] """ def __init__(self, config: ModelConfig): super().__init__() if config.n_embd % config.n_head != 0: raise ValueError("n_embd must be divisible by n_head.") self.n_head = config.n_head self.head_dim = config.n_embd // config.n_head self.dropout = config.dropout self.q_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.k_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.v_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.rope = RotaryEmbedding( head_dim=self.head_dim, theta=config.rope_theta, ) def forward(self, x: torch.Tensor) -> torch.Tensor: batch, seq_len, hidden = x.shape q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) # [B, T, H] -> [B, heads, T, head_dim] q = q.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) k = k.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) v = v.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) q, k = self.rope(q, k) attn_output = F.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True, ) # [B, heads, T, head_dim] -> [B, T, H] attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(batch, seq_len, hidden) return self.out_proj(attn_output) class SwiGLUMLP(nn.Module): """ SwiGLU feed-forward network. Formula: FFN(x) = W_down( SiLU(W_gate x) * W_up x ) Used in modern LLMs because it gives better quality than plain GELU MLP. Input: x: [batch, seq, hidden] Output: [batch, seq, hidden] """ def __init__(self, config: ModelConfig): super().__init__() hidden_dim = int(2 * (4 * config.n_embd) / 3) if config.ffn_dim_multiplier is not None: hidden_dim = int(config.ffn_dim_multiplier * hidden_dim) # Round to hardware-friendly multiple. hidden_dim = config.multiple_of * math.ceil(hidden_dim / config.multiple_of) self.gate_proj = nn.Linear(config.n_embd, hidden_dim, bias=config.bias) self.up_proj = nn.Linear(config.n_embd, hidden_dim, bias=config.bias) self.down_proj = nn.Linear(hidden_dim, config.n_embd, bias=config.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj( F.silu(self.gate_proj(x)) * self.up_proj(x) ) class TransformerBlock(nn.Module): """ Pre-Norm Transformer block. Structure: x = x + Attention(RMSNorm(x)) x = x + MLP(RMSNorm(x)) Pre-norm is more stable for deep LMs. """ def __init__(self, config: ModelConfig): super().__init__() self.attn_norm = RMSNorm(config.n_embd, config.norm_eps) self.ffn_norm = RMSNorm(config.n_embd, config.norm_eps) self.attn = CausalSelfAttention(config) self.mlp = SwiGLUMLP(config) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.attn_norm(x)) x = x + self.mlp(self.ffn_norm(x)) return x