The minimum code

"""llm/model.py — a teaching-scale decoder transformer."""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
 
@dataclass
class GPTConfig:
    vocab_size: int = 50257
    block_size: int = 64
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    ffn_mult: int = 4
    dropout: float = 0.0
    # Switches we'll use later in the book. Defaults preserve chapter-12 behavior.
    bias: bool = False               # bias on attention linears (ch.15 → True for GPT-2)
    tied_lm_head: bool = False       # share tok_emb ↔ head (ch.15 → True for GPT-2)
    gelu_approximate: str = "none"   # "none" or "tanh" (ch.15 → "tanh" for GPT-2)
 
class CausalSelfAttention(nn.Module):
    """Multi-head attention with a causal mask. Chapters 8-9."""
    def __init__(self, cfg):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0
        self.n_head = cfg.n_head
        self.n_embd = cfg.n_embd
        # combined Q, K, V projection for speed
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias)
        # causal mask: lower-triangular ones up to block_size
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(cfg.block_size, cfg.block_size)).view(1, 1, cfg.block_size, cfg.block_size),
        )
 
    def forward(self, x, past_kv=None):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(self.n_embd, dim=2)
        # split into heads: (B, T, n_head, head_dim) -> (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
 
        # KV cache: extend k, v with previously-computed keys/values.
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        present_kv = (k, v)
 
        # scaled dot-product attention with causal mask
        att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))
        if past_kv is None:
            # full prefix pass — apply causal mask
            att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        # else: T == 1, the new token attends to everything in the cache; no mask
        att = F.softmax(att, dim=-1)
        out = att @ v  # (B, n_head, T, head_dim)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out), present_kv
 
class FFN(nn.Module):
    """Per-token MLP with GELU. Chapter 10."""
    def __init__(self, cfg):
        super().__init__()
        self.fc1 = nn.Linear(cfg.n_embd, cfg.ffn_mult * cfg.n_embd)
        self.fc2 = nn.Linear(cfg.ffn_mult * cfg.n_embd, cfg.n_embd)
        self.approximate = cfg.gelu_approximate
 
    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x), approximate=self.approximate))
 
class Block(nn.Module):
    """One transformer block, pre-norm. Chapter 10."""
    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.n_embd)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.n_embd)
        self.ffn = FFN(cfg)
 
    def forward(self, x, past_kv=None):
        attn_out, present_kv = self.attn(self.ln1(x), past_kv=past_kv)
        x = x + attn_out
        x = x + self.ffn(self.ln2(x))
        return x, present_kv
 
class GPT(nn.Module):
    """Full decoder transformer. Chapters 8-10 + token + position embeddings."""
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        if cfg.tied_lm_head:
            self.head.weight = self.tok_emb.weight
 
    def forward(self, idx, targets=None, past_kvs=None):
        B, T = idx.shape
        # When caching, position embeddings must offset by the cached length.
        offset = 0 if past_kvs is None else past_kvs[0][0].size(2)
        tok = self.tok_emb(idx)
        pos = self.pos_emb(torch.arange(offset, offset + T, device=idx.device))
        x = tok + pos
        present_kvs = []
        past_kvs = past_kvs or [None] * len(self.blocks)
        for block, past_kv in zip(self.blocks, past_kvs):
            x, present_kv = block(x, past_kv=past_kv)
            present_kvs.append(present_kv)
        x = self.ln_f(x)
        logits = self.head(x)
        if targets is None:
            return logits, present_kvs
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

"""check_model.py — confirm the model's contract on a fake batch."""
import math
import torch
 
from llm.model import GPT, GPTConfig
 
 
# [1]
cfg = GPTConfig()
model = GPT(cfg)
 
# [2]
n_params = sum(p.numel() for p in model.parameters())
assert 13_500_000 < n_params < 14_500_000, f"expected ~14M params, got {n_params:,}"
 
# [3]
idx = torch.randint(0, cfg.vocab_size, (2, cfg.block_size))
logits, _ = model(idx)
expected_shape = (2, cfg.block_size, cfg.vocab_size)
assert tuple(logits.shape) == expected_shape, (
    f"expected output shape {expected_shape}, got {tuple(logits.shape)}"
)
 
# [4]
targets = torch.randint(0, cfg.vocab_size, (2, cfg.block_size))
_, loss = model(idx, targets=targets)
expected_loss = math.log(cfg.vocab_size)
assert abs(loss.item() - expected_loss) < 1.0, (
    f"expected loss near log(vocab) = {expected_loss:.2f}, got {loss.item():.2f}"
)
 
print(f"✓ {n_params:,} params")
print(f"✓ forward shape: {tuple(logits.shape)}")
print(f"✓ initial loss: {loss.item():.2f} (expected ~{expected_loss:.2f})")