"""Tiny CPU decoder-only language model for Lab B (lessons 1-3, 7-10). A minimal GPT-style decoder assembled from the lesson components: * token embeddings + learned absolute position embeddings (lessons 1, 2) * N pre-norm decoder blocks: RMSNorm -> causal self-attention -> residual, RMSNorm -> SwiGLU MLP -> residual (lesson 8) * final RMSNorm + output projection to vocab logits, tied to embeddings (lesson 9) * next-token cross-entropy with loss-side shift logits[:, :-1] vs labels[:, 1:] (lessons 3, 10) Everything runs on CPU with float32. No downloads, no API keys. Weight-shape convention note: * Mathematical descriptions (e.g. lessons 8-9) write a linear map as W [d, d_out] applied as x @ W. * PyTorch nn.Linear(d_in, d_out) stores its weight tensor as [d_out, d_in] (out_features x in_features) and applies it as x @ weight.T. The shapes in comments below follow the mathematical [d_in, d_out] convention; the stored nn.Parameter shape is the transpose. This matters when reading .weight.shape directly. """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F class RMSNorm(nn.Module): """Root-mean-square normalisation (lesson 8). Gain only, no bias.""" def __init__(self, d: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(d)) def forward(self, x: torch.Tensor) -> torch.Tensor: rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return x / rms * self.weight class CausalSelfAttention(nn.Module): """Manual scaled dot-product causal self-attention (lessons 1, 4-6). Projects Q/K/V/O independently, applies a lower-triangular causal mask so position t can attend only to positions 0..t, and merges heads back to d. """ def __init__(self, d: int, n_heads: int): super().__init__() assert d % n_heads == 0, "d must be divisible by n_heads" self.d = d self.n_heads = n_heads self.d_head = d // n_heads self.W_Q = nn.Linear(d, d, bias=True) # math [d, d], stored [d, d] (square) self.W_K = nn.Linear(d, d, bias=True) # math [d, d], stored [d, d] (square) self.W_V = nn.Linear(d, d, bias=True) # math [d, d], stored [d, d] (square) self.W_O = nn.Linear(d, d, bias=True) # math [d, d], stored [d, d] (square) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, d = x.shape Q = ( self.W_Q(x) .view(B, T, self.n_heads, self.d_head) .transpose(1, 2) ) # [B, h, T, d_head] K = ( self.W_K(x) .view(B, T, self.n_heads, self.d_head) .transpose(1, 2) ) # [B, h, T, d_head] V = ( self.W_V(x) .view(B, T, self.n_heads, self.d_head) .transpose(1, 2) ) # [B, h, T, d_head] scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5) # [B, h, T, T] causal_mask = torch.tril( torch.ones(T, T, device=x.device, dtype=torch.bool) ).view(1, 1, T, T) scores = scores.masked_fill(~causal_mask, float("-inf")) attn_weights = F.softmax(scores, dim=-1) # [B, h, T, T] attn_out = attn_weights @ V # [B, h, T, d_head] attn_out = ( attn_out.transpose(1, 2).contiguous().view(B, T, d) ) # [B, T, d] return self.W_O(attn_out) class SwiGLU(nn.Module): """Gated MLP (lesson 8). d -> 4d (gate+value) -> SiLU(gate)*value -> 2d -> d.""" def __init__(self, d: int, hidden_mult: int = 4): super().__init__() hidden = d * hidden_mult self.w = nn.Linear(d, hidden, bias=False) # math [d, 4d], stored [4d, d] self.w_out = nn.Linear(hidden // 2, d, bias=False) # math [2d, d], stored [d, 2d] def forward(self, x: torch.Tensor) -> torch.Tensor: x_proj = self.w(x) # [B, T, 4d] gate, val = x_proj.chunk(2, dim=-1) # each [B, T, 2d] return self.w_out(F.silu(gate) * val) # [B, T, d] class DecoderBlock(nn.Module): """Pre-norm decoder block (lesson 8): x + Attn(Norm(x)), x + MLP(Norm(x)).""" def __init__(self, d: int, n_heads: int): super().__init__() self.norm1 = RMSNorm(d) self.attn = CausalSelfAttention(d, n_heads) self.norm2 = RMSNorm(d) self.mlp = SwiGLU(d) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x class TinyDecoder(nn.Module): """Complete tiny decoder-only language model (lesson 9).""" def __init__( self, vocab_size: int, d: int, n_heads: int, n_blocks: int, max_position: int, tie_weights: bool = True, ): super().__init__() self.d = d self.vocab_size = vocab_size self.tie_weights = tie_weights self.tok_emb = nn.Embedding(vocab_size, d) # [V, d] self.pos_emb = nn.Embedding(max_position, d) # [max_position, d] self.blocks = nn.ModuleList( [DecoderBlock(d, n_heads) for _ in range(n_blocks)] ) self.final_norm = RMSNorm(d) if tie_weights: self.output_proj = None # reuse tok_emb.weight.T, 0 extra params else: self.output_proj = nn.Linear(d, vocab_size, bias=False) # math [d, V], stored [V, d] def forward( self, input_ids: torch.Tensor, positions: torch.Tensor | None = None ) -> torch.Tensor: B, T = input_ids.shape if positions is None: positions = ( torch.arange(T, device=input_ids.device) .unsqueeze(0) .expand(B, T) ) h = self.tok_emb(input_ids) + self.pos_emb(positions) # [B, T, d] for block in self.blocks: h = block(h) # [B, T, d] -> [B, T, d] h = self.final_norm(h) # [B, T, d] if self.tie_weights: logits = h @ self.tok_emb.weight.T # [B, T, d] x [d, V] -> [B, T, V] else: logits = self.output_proj(h) # [B, T, V] return logits def count_parameters(model: nn.Module) -> int: """Total trainable parameters.""" return sum(p.numel() for p in model.parameters()) def expected_parameter_count( vocab_size: int, d: int, n_heads: int, n_blocks: int, max_position: int, tie_weights: bool = True, ) -> int: """Hand-computed parameter count matching TinyDecoder architecture. Per decoder block (d=n_heads*d_head): * attention Q/K/V/O: 4*(d*d) + 4*d (biases on all four) * MLP SwiGLU: w [d, 4d]=4d*d + w_out [2d, d]=2d*d -> 6*d*d * RMSNorm x2: 2*d (gain only, no bias) * subtotal: 10*d*d + 6*d Plus: token embeddings V*d, position embeddings max_position*d, final norm d, output projection 0 (tied) or d*V (untied). """ tok = vocab_size * d pos = max_position * d per_block = 10 * d * d + 6 * d blocks = n_blocks * per_block final_norm = d out = 0 if tie_weights else d * vocab_size return tok + pos + blocks + final_norm + out