Create components.py

components.py

@@ -0,0 +1,237 @@
+"""
+Model components optimized for CPU training.
+
+Design rationale:
+- RMSNorm instead of LayerNorm: simpler, faster (no mean computation)
+- Rotary Position Embeddings (RoPE): no learned position embeddings needed,
+  saves parameters and generalizes better
+- LoRA-style low-rank linear layers: dramatically reduces parameter count
+  while maintaining expressiveness
+- All operations use float32 for CPU stability (no mixed precision)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple
+
+
+class RMSNorm(nn.Module):
+    """
+    Root Mean Square normalization.
+    
+    Why: ~30% faster than LayerNorm on CPU since it skips mean computation.
+    Empirically equivalent performance for transformers.
+    """
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * norm * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    """
+    Rotary Position Embedding (RoPE).
+    
+    Why:
+    - No learned parameters (saves memory)
+    - Relative position awareness without extra params
+    - Extrapolates better to unseen sequence lengths
+    - Computationally efficient on CPU (just sin/cos)
+    """
+    def __init__(self, dim: int, max_seq_len: int = 512, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        # Pre-compute for max_seq_len to avoid recomputation
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len: int):
+        t = torch.arange(seq_len, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer('cos_cached', emb.cos())
+        self.register_buffer('sin_cached', emb.sin())
+
+    def forward(self, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        if seq_len > self.cos_cached.size(0):
+            self._build_cache(seq_len)
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, 
+                          cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply rotary embeddings to queries and keys."""
+    # cos, sin: [seq_len, dim]
+    cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, dim]
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class LoRALinear(nn.Module):
+    """
+    Low-Rank Adaptation linear layer.
+    
+    Why: Instead of full d_in x d_out matrix, uses two smaller matrices:
+    d_in x rank + rank x d_out. For rank=16, d_in=d_out=256:
+    Full: 65,536 params
+    LoRA: 256*16 + 16*256 = 8,192 params (8x reduction!)
+    
+    Still maintains good expressiveness for the tasks we need.
+    """
+    def __init__(self, in_features: int, out_features: int, rank: int = 16, bias: bool = False):
+        super().__init__()
+        self.rank = rank
+        # If rank is large enough, just use full linear
+        if rank >= min(in_features, out_features) // 2:
+            self.use_lora = False
+            self.linear = nn.Linear(in_features, out_features, bias=bias)
+        else:
+            self.use_lora = True
+            self.down = nn.Linear(in_features, rank, bias=False)
+            self.up = nn.Linear(rank, out_features, bias=bias)
+            # Initialize to approximate identity-like behavior
+            nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
+            nn.init.zeros_(self.up.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.use_lora:
+            return self.up(self.down(x))
+        return self.linear(x)
+
+
+class GatedMLP(nn.Module):
+    """
+    SwiGLU-style gated MLP.
+    
+    Why: Gated activation functions consistently outperform standard ReLU/GELU
+    in transformers, especially at small scale. The gate provides a learned
+    "feature selection" mechanism.
+    
+    Uses LoRA projections to save parameters.
+    """
+    def __init__(self, d_model: int, d_ff: int, rank: int = 16, dropout: float = 0.05):
+        super().__init__()
+        self.gate_proj = LoRALinear(d_model, d_ff, rank=rank)
+        self.up_proj = LoRALinear(d_model, d_ff, rank=rank)
+        self.down_proj = LoRALinear(d_ff, d_model, rank=rank)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        return self.dropout(self.down_proj(gate * up))
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention with RoPE and optional Grouped Query Attention.
+    
+    Why these choices:
+    - Grouped Query Attention (GQA): shares KV heads, reducing memory and params
+      while maintaining quality. For 8 heads with 4 KV groups: 50% KV param reduction.
+    - Pre-computed causal mask: avoids recomputing each forward pass on CPU
+    - RoPE applied per-head: correct relative position encoding
+    """
+    def __init__(self, d_model: int, n_heads: int, rank: int = 16, 
+                 dropout: float = 0.05, max_seq_len: int = 512,
+                 n_kv_heads: Optional[int] = None):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads or n_heads
+        self.head_dim = d_model // n_heads
+        self.n_rep = n_heads // self.n_kv_heads  # repetition factor for GQA
+        
+        assert d_model % n_heads == 0
+        
+        self.q_proj = LoRALinear(d_model, d_model, rank=rank)
+        self.k_proj = LoRALinear(d_model, self.n_kv_heads * self.head_dim, rank=rank)
+        self.v_proj = LoRALinear(d_model, self.n_kv_heads * self.head_dim, rank=rank)
+        self.o_proj = LoRALinear(d_model, d_model, rank=rank)
+        
+        self.dropout = nn.Dropout(dropout)
+        self.rope = RotaryEmbedding(self.head_dim, max_seq_len)
+        
+        # Pre-compute causal mask
+        mask = torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1).bool()
+        self.register_buffer('causal_mask', mask)
+
+    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
+        """Repeat KV heads to match Q heads for GQA."""
+        if self.n_rep == 1:
+            return x
+        bs, n_kv, seq_len, head_dim = x.shape
+        x = x[:, :, None, :, :].expand(bs, n_kv, self.n_rep, seq_len, head_dim)
+        return x.reshape(bs, self.n_heads, seq_len, head_dim)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, T, C = x.shape
+        
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        
+        # Apply RoPE
+        cos, sin = self.rope(T)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        
+        # Expand KV for GQA
+        k = self._repeat_kv(k)
+        v = self._repeat_kv(v)
+        
+        # Attention
+        scale = math.sqrt(self.head_dim)
+        attn = torch.matmul(q, k.transpose(-2, -1)) / scale
+        
+        # Apply causal mask
+        causal = self.causal_mask[:T, :T].unsqueeze(0).unsqueeze(0)
+        attn = attn.masked_fill(causal, float('-inf'))
+        
+        if mask is not None:
+            # mask shape: [B, T] -> [B, 1, 1, T]
+            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+        
+        attn = F.softmax(attn, dim=-1)
+        attn = self.dropout(attn)
+        
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(out)
+
+
+class TransformerBlock(nn.Module):
+    """
+    Single transformer block with pre-norm architecture.
+    
+    Why pre-norm: More stable training, especially at small scale.
+    Gradient flow is better since residual path is unimpeded.
+    """
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, 
+                 rank: int = 16, dropout: float = 0.05, 
+                 max_seq_len: int = 512, n_kv_heads: Optional[int] = None):
+        super().__init__()
+        self.attn_norm = RMSNorm(d_model)
+        self.attn = MultiHeadAttention(d_model, n_heads, rank, dropout, max_seq_len, n_kv_heads)
+        self.ff_norm = RMSNorm(d_model)
+        self.ff = GatedMLP(d_model, d_ff, rank, dropout)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x), mask)
+        x = x + self.ff(self.ff_norm(x))
+        return x