Phase 3: decoder-only transformer (~4.3M params)
- src/model.py: GPT written module by module — RMSNorm, explicit causal self-attention (scores/mask/softmax/weighted-sum), GELU MLP, pre-norm residual blocks, learned positions, weight-tied head. GPT-2 style init. - configs/v1.py: frozen ModelConfig dataclass (4L/256d/4h/256ctx/4096vocab) - tests/test_model.py: exact param count (4,262,144), forward shapes, init loss ~= ln(vocab), weight tying, causal-masking check (5 tests, forward-only) - Verified forward pass runs on MPS.
This commit is contained in:
+162
@@ -0,0 +1,162 @@
|
||||
"""A decoder-only (GPT-style) transformer, written module by module.
|
||||
|
||||
Architecture (pre-norm, as in modern GPTs):
|
||||
|
||||
tokens ─► token embedding + learned position embedding
|
||||
─► N x [ RMSNorm ─► causal self-attention ─► + residual
|
||||
RMSNorm ─► MLP ─► + residual ]
|
||||
─► final RMSNorm ─► linear head (weight-tied to token embedding)
|
||||
|
||||
Attention is implemented explicitly (scores → causal mask → softmax → weighted
|
||||
sum of values) rather than via a fused kernel, because seeing the mechanism is
|
||||
the point here. `torch.nn.functional.scaled_dot_product_attention` is the
|
||||
faster production alternative and an easy iteration-2 swap.
|
||||
|
||||
The model takes any config object exposing: vocab_size, context_length,
|
||||
n_layers, n_heads, d_model, d_ff.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
"""Root-mean-square layer norm: rescale each vector to unit RMS, then a
|
||||
learned per-channel gain. Cheaper than LayerNorm (no mean-centering, no bias)."""
|
||||
|
||||
def __init__(self, dim: int, eps: float = 1e-5) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
return x * rms * self.weight
|
||||
|
||||
|
||||
class CausalSelfAttention(nn.Module):
|
||||
"""Multi-head self-attention where each position may only attend to itself
|
||||
and earlier positions (the causal constraint that makes this a language model)."""
|
||||
|
||||
def __init__(self, cfg) -> None:
|
||||
super().__init__()
|
||||
assert cfg.d_model % cfg.n_heads == 0
|
||||
self.n_heads = cfg.n_heads
|
||||
self.d_head = cfg.d_model // cfg.n_heads
|
||||
# One matrix produces query, key and value for every head at once.
|
||||
self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
|
||||
self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
|
||||
# Lower-triangular mask, precomputed once. Not a parameter, so register
|
||||
# it as a buffer (moves with .to(device), excluded from grads).
|
||||
causal = torch.tril(torch.ones(cfg.context_length, cfg.context_length))
|
||||
self.register_buffer("mask", causal.view(1, 1, cfg.context_length, cfg.context_length))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
B, T, C = x.shape
|
||||
q, k, v = self.qkv(x).split(C, dim=2)
|
||||
# (B, T, C) -> (B, n_heads, T, d_head): split channels across heads.
|
||||
q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
||||
k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
||||
v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
||||
|
||||
# Attention scores, scaled so their variance doesn't grow with d_head.
|
||||
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head) # (B, nh, T, T)
|
||||
att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
|
||||
att = F.softmax(att, dim=-1)
|
||||
y = att @ v # (B, nh, T, d_head): each position = weighted sum of past values
|
||||
|
||||
y = y.transpose(1, 2).contiguous().view(B, T, C) # merge heads back
|
||||
return self.proj(y)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Position-wise feed-forward: expand, GELU, project back."""
|
||||
|
||||
def __init__(self, cfg) -> None:
|
||||
super().__init__()
|
||||
self.fc = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
|
||||
self.proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.proj(F.gelu(self.fc(x)))
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
"""One transformer block: attention then MLP, each with a pre-norm and a
|
||||
residual connection so gradients flow straight down the stack."""
|
||||
|
||||
def __init__(self, cfg) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = RMSNorm(cfg.d_model)
|
||||
self.attn = CausalSelfAttention(cfg)
|
||||
self.norm2 = RMSNorm(cfg.d_model)
|
||||
self.mlp = MLP(cfg)
|
||||
|
||||
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 GPT(nn.Module):
|
||||
def __init__(self, cfg) -> None:
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
|
||||
self.pos_emb = nn.Embedding(cfg.context_length, cfg.d_model)
|
||||
self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
|
||||
self.norm_f = RMSNorm(cfg.d_model)
|
||||
self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
|
||||
# Weight tying: the output head reuses the token-embedding matrix. Saves
|
||||
# ~1M params and ties "recognising a token" to "predicting it".
|
||||
self.head.weight = self.tok_emb.weight
|
||||
|
||||
self.apply(self._init_weights)
|
||||
# GPT-2 trick: shrink residual-projection init by 1/sqrt(2*n_layers) so
|
||||
# the residual stream doesn't blow up with depth.
|
||||
std = 0.02 / math.sqrt(2 * cfg.n_layers)
|
||||
for name, p in self.named_parameters():
|
||||
if name.endswith("proj.weight"):
|
||||
nn.init.normal_(p, mean=0.0, std=std)
|
||||
|
||||
@staticmethod
|
||||
def _init_weights(module: nn.Module) -> None:
|
||||
if isinstance(module, nn.Linear):
|
||||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
elif isinstance(module, nn.Embedding):
|
||||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||||
|
||||
def forward(
|
||||
self, idx: torch.Tensor, targets: torch.Tensor | None = None
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
B, T = idx.shape
|
||||
assert T <= self.cfg.context_length, "sequence longer than context_length"
|
||||
pos = torch.arange(T, device=idx.device)
|
||||
x = self.tok_emb(idx) + self.pos_emb(pos) # broadcast pos over batch
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
x = self.norm_f(x)
|
||||
logits = self.head(x) # (B, T, vocab_size)
|
||||
|
||||
loss = None
|
||||
if targets is not None:
|
||||
loss = F.cross_entropy(
|
||||
logits.view(-1, logits.size(-1)), targets.reshape(-1)
|
||||
)
|
||||
return logits, loss
|
||||
|
||||
def num_params(self, non_embedding: bool = False) -> int:
|
||||
"""Total trainable parameters. Weight tying means the token embedding /
|
||||
head matrix is counted once. `non_embedding` excludes the embedding tables."""
|
||||
n = sum(p.numel() for p in self.parameters())
|
||||
if non_embedding:
|
||||
n -= self.tok_emb.weight.numel()
|
||||
n -= self.pos_emb.weight.numel()
|
||||
return n
|
||||
Reference in New Issue
Block a user