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:
@@ -0,0 +1,25 @@
|
||||
"""Experiment v1 configuration.
|
||||
|
||||
One experiment = one config file. Iteration 2 is a copy of this file with
|
||||
different numbers, not a code change.
|
||||
|
||||
v1 is deliberately small (~4.3M parameters): the goal is a complete, legible
|
||||
pipeline, not a strong model. It sits below TinyStories' coherence sweet spot
|
||||
(~30M params), so some incoherence in the output is expected.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelConfig:
|
||||
vocab_size: int = 4096 # matches the trained tokenizer
|
||||
context_length: int = 256 # tokens of history the model attends over
|
||||
n_layers: int = 4
|
||||
n_heads: int = 4
|
||||
d_model: int = 256 # residual-stream width
|
||||
d_ff: int = 1024 # MLP hidden width (4 x d_model)
|
||||
|
||||
|
||||
# The single source of truth other modules import.
|
||||
model = ModelConfig()
|
||||
+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
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Model tests: parameter count, output shapes, init loss, and causality.
|
||||
|
||||
All forward-pass only (no training). Run: `.venv/bin/python tests/test_model.py`
|
||||
"""
|
||||
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
sys.path.insert(0, str(ROOT / "configs"))
|
||||
|
||||
from model import GPT # noqa: E402
|
||||
from v1 import ModelConfig # noqa: E402
|
||||
|
||||
|
||||
def test_param_count_matches_hand_estimate():
|
||||
model = GPT(ModelConfig())
|
||||
# tok_emb 4096*256 + pos_emb 256*256
|
||||
# + 4 blocks * (qkv 256*768 + attn_proj 256*256 + 2 norms*256
|
||||
# + mlp_fc 256*1024 + mlp_proj 1024*256)
|
||||
# + final norm 256 (head is weight-tied, so no extra params)
|
||||
expected = (
|
||||
4096 * 256 + 256 * 256
|
||||
+ 4 * (256 * 768 + 256 * 256 + 2 * 256 + 256 * 1024 + 1024 * 256)
|
||||
+ 256
|
||||
)
|
||||
assert expected == 4_262_144
|
||||
assert model.num_params() == expected
|
||||
|
||||
|
||||
def test_weight_tying():
|
||||
model = GPT(ModelConfig())
|
||||
# Head and token embedding must be the very same parameter tensor.
|
||||
assert model.head.weight is model.tok_emb.weight
|
||||
|
||||
|
||||
def test_forward_shapes():
|
||||
cfg = ModelConfig()
|
||||
model = GPT(cfg)
|
||||
idx = torch.randint(0, cfg.vocab_size, (2, 16))
|
||||
logits, loss = model(idx)
|
||||
assert logits.shape == (2, 16, cfg.vocab_size)
|
||||
assert loss is None
|
||||
_, loss = model(idx, targets=idx)
|
||||
assert loss.ndim == 0 # scalar
|
||||
|
||||
|
||||
def test_init_loss_near_uniform():
|
||||
# A well-initialised model predicts ~uniformly at first, so the
|
||||
# cross-entropy loss should be close to ln(vocab_size). Targets are drawn
|
||||
# independently of the input: with weight tying, using the input as its own
|
||||
# target would leak (each position's residual already holds its embedding).
|
||||
cfg = ModelConfig()
|
||||
torch.manual_seed(0)
|
||||
model = GPT(cfg)
|
||||
idx = torch.randint(0, cfg.vocab_size, (8, 64))
|
||||
targets = torch.randint(0, cfg.vocab_size, (8, 64))
|
||||
_, loss = model(idx, targets=targets)
|
||||
assert abs(loss.item() - math.log(cfg.vocab_size)) < 0.5, loss.item()
|
||||
|
||||
|
||||
def test_causality():
|
||||
# Changing a token at position t must not affect logits at positions < t.
|
||||
cfg = ModelConfig()
|
||||
torch.manual_seed(0)
|
||||
model = GPT(cfg)
|
||||
model.eval()
|
||||
idx = torch.randint(0, cfg.vocab_size, (1, 32))
|
||||
with torch.no_grad():
|
||||
base, _ = model(idx)
|
||||
changed = idx.clone()
|
||||
changed[0, 20] = (changed[0, 20] + 1) % cfg.vocab_size
|
||||
after, _ = model(changed)
|
||||
# Positions before the edit are identical; the edited position differs.
|
||||
assert torch.allclose(base[0, :20], after[0, :20], atol=1e-5)
|
||||
assert not torch.allclose(base[0, 20], after[0, 20])
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||||
for t in tests:
|
||||
t()
|
||||
print(f"ok {t.__name__}")
|
||||
print(f"\n{len(tests)} passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user