A GPT from Scratch

import warnings
warnings.filterwarnings("ignore")
import math, random, time
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

random.seed(0); np.random.seed(0); torch.manual_seed(0)
torch.set_num_threads(max(1, torch.get_num_threads()))

sample_text = (
    "the model learns to predict the next token. the next token depends on the tokens before it. "
    "a token can be a word, a part of a word, or a single character. learning to predict the next "
    "word teaches the model grammar, facts and reasoning, because all of them help prediction. "
    "transformers use attention to let every token look at the tokens that came before it."
)

def bpe_train(text, n_merges):
    seq = list(text)                         # start from characters
    merges = []
    for _ in range(n_merges):
        pairs = Counter(zip(seq, seq[1:]))
        (a, b), count = pairs.most_common(1)[0]
        if count < 2:
            break
        merges.append((a, b))
        new, i = [], 0
        while i < len(seq):                  # replace every occurrence of the pair
            if i < len(seq) - 1 and seq[i] == a and seq[i + 1] == b:
                new.append(a + b); i += 2
            else:
                new.append(seq[i]); i += 1
        seq = new
    return merges, seq

def bpe_encode(text, merges):
    seq = list(text)
    for a, b in merges:                      # apply merges in the order they were learned
        new, i = [], 0
        while i < len(seq):
            if i < len(seq) - 1 and seq[i] == a and seq[i + 1] == b:
                new.append(a + b); i += 2
            else:
                new.append(seq[i]); i += 1
        seq = new
    return seq

merges, tokens = bpe_train(sample_text, 40)
print("first 15 merges:", [a + b for a, b in merges[:15]])
print(f"\n{len(sample_text)} characters -> {len(tokens)} tokens after {len(merges)} merges "
      f"({len(sample_text) / len(tokens):.1f} chars per token)\n")
print("encoding a new sentence:", bpe_encode("the next word predicts the model", merges))

COUNTRIES = {"france": "paris", "spain": "madrid", "italy": "rome", "germany": "berlin", "portugal": "lisbon",
             "japan": "tokyo", "egypt": "cairo", "peru": "lima", "chile": "santiago", "kenya": "nairobi",
             "canada": "ottawa", "norway": "oslo", "greece": "athens", "poland": "warsaw", "cuba": "havana",
             "india": "new delhi", "china": "beijing", "russia": "moscow", "mexico": "mexico city", "ireland": "dublin",
             "austria": "vienna", "sweden": "stockholm", "finland": "helsinki", "turkey": "ankara", "brazil": "brasilia",
             "argentina": "buenos aires", "colombia": "bogota", "costa rica": "san jose", "morocco": "rabat", "australia": "canberra"}

SUBJ = [("i", "yo", 0), ("you", "tú", 1), ("he", "él", 2), ("she", "ella", 2), ("we", "nosotros", 3), ("they", "ellos", 4)]
VERB = {"eat": ("eats", ["como", "comes", "come", "comemos", "comen"]),
        "see": ("sees", ["veo", "ves", "ve", "vemos", "ven"]),
        "want": ("wants", ["quiero", "quieres", "quiere", "queremos", "quieren"]),
        "buy": ("buys", ["compro", "compras", "compra", "compramos", "compran"]),
        "have": ("has", ["tengo", "tienes", "tiene", "tenemos", "tienen"])}
NOUN = [("apple", "manzana", "f"), ("book", "libro", "m"), ("car", "coche", "m"), ("house", "casa", "f"),
        ("dog", "perro", "m"), ("table", "mesa", "f"), ("cat", "gato", "m"), ("shirt", "camisa", "f")]
ADJ = [("red", "rojo", "roja"), ("new", "nuevo", "nueva"), ("old", "viejo", "vieja"), ("white", "blanco", "blanca"), (None, None, None)]

def translation_pairs():
    out = []
    for en_s, es_s, p in SUBJ:
        for v, (v3, conj) in VERB.items():
            for en_n, es_n, g in NOUN:
                for en_a, am, af in ADJ:
                    en = f"{en_s} {v3 if en_s in ('he', 'she') else v} the {en_a + ' ' if en_a else ''}{en_n}"
                    es = f"{es_s} {conj[p]} {'el' if g == 'm' else 'la'} {es_n}" + (f" {am if g == 'm' else af}" if en_a else "")
                    out.append((en, es))
    return out

ALL_TR = translation_pairs()
random.shuffle(ALL_TR)
TR_TEST, TR_TRAIN = ALL_TR[:120], ALL_TR[120:]
ADD_ALL = [(a, b) for a in range(100) for b in range(100)]
random.shuffle(ADD_ALL)
ADD_TEST, ADD_TRAIN = ADD_ALL[:500], ADD_ALL[500:]

def rand_word():
    return "".join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(random.randint(3, 6)))

def example(task, split="train"):
    # returns (prompt, answer); a training line is prompt + answer
    if task == "add":
        a, b = random.choice(ADD_TRAIN if split == "train" else ADD_TEST)
        return f"add {a}+{b}=", f"{a + b}"
    if task == "rev":
        w = rand_word()
        return f"rev {w}=", w[::-1]
    if task == "sort":
        ds = [random.randint(0, 9) for _ in range(random.randint(3, 6))]
        return f"sort {' '.join(map(str, ds))}=", " ".join(map(str, sorted(ds)))
    if task == "es":
        en, es = random.choice(TR_TRAIN if split == "train" else TR_TEST)
        return f"es {en}=", es
    if task == "capital":
        c = random.choice(list(COUNTRIES))
        return f"capital {c}=", COUNTRIES[c]

TASKS = ["add", "rev", "sort", "es", "capital"]
WEIGHTS = [5, 1, 2, 1, 1]          # arithmetic is the hardest skill, so it gets more data
corpus = "\n".join("".join(example(random.choices(TASKS, WEIGHTS)[0])) for _ in range(60000)) + "\n"

print(f"corpus: {len(corpus):,} characters, {corpus.count(chr(10)):,} lines\n")
print(corpus[:400])

chars = sorted(set(corpus) | set("".join(e for pair in ALL_TR for e in pair)))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = torch.tensor(encode(corpus), dtype=torch.long)

print(f"vocabulary ({len(chars)} tokens): {''.join(chars)!r}")
print("encode('add 2+2=4') ->", encode("add 2+2=4"))

BLOCK = 64

def get_batch(B, source=None):
    source = data if source is None else source
    ix = torch.randint(len(source) - BLOCK - 1, (B,))
    x = torch.stack([source[i:i + BLOCK] for i in ix])
    y = torch.stack([source[i + 1:i + BLOCK + 1] for i in ix])
    return x, y

x, y = get_batch(1)
for t in range(12):
    print(f"context {decode(x[0, :t + 1].tolist())!r:22s} -> target {decode([y[0, t].item()])!r}")

class CausalSelfAttention(nn.Module):
    def __init__(self, d, n_head, block_size):
        super().__init__()
        self.h = n_head
        self.qkv = nn.Linear(d, 3 * d)            # Q, K and V projections in one matrix
        self.proj = nn.Linear(d, d)               # W_O
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size, dtype=torch.bool)))
        self.last_attn = None

    def forward(self, x, past=None, keep_attn=False):
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=2)
        q, k, v = (t.view(B, T, self.h, D // self.h).transpose(1, 2) for t in (q, k, v))   # (B, heads, T, dh)
        if past is not None:                       # KV cache: prepend keys/values of earlier tokens
            k, v = torch.cat([past[0], k], dim=2), torch.cat([past[1], v], dim=2)
        present = (k, v)
        if keep_attn:                              # explicit version, keeps the weights for plotting
            Tk = k.shape[2]
            att = (q @ k.transpose(-2, -1)) / math.sqrt(D // self.h)
            att = att.masked_fill(~self.mask[Tk - T:Tk, :Tk], float("-inf")).softmax(-1)
            self.last_attn = att.detach()
            y = att @ v
        else:                                      # fused kernel, same maths, faster
            y = F.scaled_dot_product_attention(q, k, v, is_causal=(past is None and T > 1))
        return self.proj(y.transpose(1, 2).reshape(B, T, D)), present

class Block(nn.Module):
    def __init__(self, d, n_head, block_size):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d), CausalSelfAttention(d, n_head, block_size)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
    def forward(self, x, past=None, keep_attn=False):
        a, present = self.attn(self.ln1(x), past, keep_attn)
        x = x + a
        return x + self.mlp(self.ln2(x)), present

class GPT(nn.Module):
    def __init__(self, vocab, block_size=64, d=128, n_layer=4, n_head=4):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab, d)
        self.pos_emb = nn.Embedding(block_size, d)
        self.blocks = nn.ModuleList(Block(d, n_head, block_size) for _ in range(n_layer))
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok_emb.weight                 # weight tying
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def forward(self, idx, targets=None, past=None, keep_attn=False):
        B, T = idx.shape
        start = 0 if past is None else past[0][0].shape[2]
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(start, start + T))
        presents = []
        for i, blk in enumerate(self.blocks):
            x, pr = blk(x, None if past is None else past[i], keep_attn)
            presents.append(pr)
        logits = self.head(self.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), targets.view(-1))
        return logits, loss, presents

torch.manual_seed(0)
model = GPT(len(chars), BLOCK, d=128, n_layer=4, n_head=4)
n_params = sum(p.numel() for p in model.parameters())
print(f"GPT with {n_params:,} parameters (GPT-2 small has 124 million, GPT-3 has 175 billion)")

@torch.no_grad()
def complete(model, prompt, max_new=40):
    # greedy decoding until a newline: the answer to a task prompt
    idx = torch.tensor([encode("\n" + prompt)])
    out = []
    for _ in range(max_new):
        logits, _, _ = model(idx[:, -BLOCK:])
        nxt = logits[0, -1].argmax().item()
        if itos[nxt] == "\n":
            break
        out.append(nxt)
        idx = torch.cat([idx, torch.tensor([[nxt]])], dim=1)
    return decode(out)

xb, yb = get_batch(64)
with torch.no_grad():
    _, loss0, _ = model(xb, yb)
print(f"initial loss {loss0.item():.3f}   vs   ln(vocab) = {math.log(len(chars)):.3f}")
print("untrained answer to 'add 2+2=':", repr(complete(model, "add 2+2=", 20)))

def evaluate(model, n=60, seed=1234):
    state = random.getstate(); random.seed(seed)       # same test prompts every time
    model.eval()
    res = {}
    for t in TASKS:
        ok = 0
        for _ in range(n):
            p, a = example(t, "test")
            ok += complete(model, p) == a
        res[t] = ok / n
    model.train(); random.setstate(state)
    return res

STEPS, BATCH, EVAL_EVERY = 2500, 32, 500
opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.1, betas=(0.9, 0.95))
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=2e-3, total_steps=STEPS, pct_start=0.05)

losses, skill_curve = [], {0: evaluate(model)}
t0 = time.time()
for step in range(1, STEPS + 1):
    xb, yb = get_batch(BATCH)
    _, loss, _ = model(xb, yb)
    opt.zero_grad(); loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step()
    losses.append(loss.item())
    if step % EVAL_EVERY == 0:
        skill_curve[step] = evaluate(model)
        accs = "  ".join(f"{t} {a:.0%}" for t, a in skill_curve[step].items())
        print(f"step {step:5d}  loss {np.mean(losses[-100:]):.3f}  |  {accs}  ({time.time() - t0:.0f}s)", flush=True)

fig, axes = plt.subplots(1, 2, figsize=(15, 4.5))
smooth = np.convolve(losses, np.ones(50) / 50, mode="valid")
axes[0].plot(smooth); axes[0].set_xlabel("step"); axes[0].set_ylabel("loss (50-step average)")
axes[0].set_title("Pre-training loss: next-character cross-entropy"); axes[0].grid(alpha=0.3)
steps_eval = list(skill_curve)
for t in TASKS:
    axes[1].plot(steps_eval, [skill_curve[s][t] for s in steps_eval], marker="o", label=t)
axes[1].set_xlabel("step"); axes[1].set_ylabel("held-out accuracy"); axes[1].set_ylim(-0.03, 1.03)
axes[1].set_title("Each skill is learned at its own pace"); axes[1].legend(); axes[1].grid(alpha=0.3)
plt.tight_layout(); plt.show()

model.eval()
random.seed(42)
for t in TASKS:
    print(f"--- {t} ---")
    for _ in range(3):
        p, a = example(t, "test")
        out = complete(model, p)
        print(f"  {p}{out:28s} {'✓' if out == a else '✗ expected ' + a}")

for country in ["france", "peru", "japan", "atlantis", "narnia", "germanyland", "nepal", "spainia"]:
    print(f"capital {country:12s} -> {complete(model, f'capital {country}='):14s}"
          f"{'(memorised)' if country in COUNTRIES else '(never seen: hallucinated)'}")

probes = [("add 123+4=", "127", "3-digit number (training used 0..99)"),
          ("rev abcdefghij=", "jihgfedcba", "10-letter word (training used 3..6)"),
          ("sort 9 8 7 6 5 4 3 2 1=", "1 2 3 4 5 6 7 8 9", "9 digits (training used 3..6)"),
          ("es the cat eats the red apple=", "?", "noun as subject (training subjects were pronouns)")]
for p, a, why in probes:
    out = complete(model, p)
    print(f"{p:32s} -> {out:22s} expected {a:20s} [{why}]")

@torch.no_grad()
def next_char_probs(prompt):
    logits, _, _ = model(torch.tensor([encode("\n" + prompt)]))
    return logits[0, -1].softmax(-1)

fig, axes = plt.subplots(1, 2, figsize=(15, 3.5))
for ax, prompt in zip(axes, ["add 23+58=", "\n"]):
    p = next_char_probs(prompt if prompt != "\n" else "")
    top = p.topk(12)
    labels = [repr(itos[i])[1:-1] if itos[i] != " " else "␣" for i in top.indices.tolist()]
    ax.bar(labels, top.values.numpy(), color="tab:blue")
    ax.set_title(f"next character after {prompt!r}" if prompt != "\n" else "first character of a new line (which task?)")
    ax.set_ylim(0, 1)
plt.tight_layout(); plt.show()

@torch.no_grad()
def generate(model, prompt="\n", max_new=200, temperature=1.0, top_k=None, top_p=None, seed=0):
    g = torch.Generator().manual_seed(seed)
    idx = torch.tensor([encode(prompt)])
    for _ in range(max_new):
        logits, _, _ = model(idx[:, -BLOCK:])
        logits = logits[0, -1]
        if temperature == 0:
            nxt = logits.argmax().view(1)
        else:
            logits = logits / temperature
            if top_k is not None:
                kth = logits.topk(top_k).values[-1]
                logits = logits.masked_fill(logits < kth, float("-inf"))
            probs = logits.softmax(-1)
            if top_p is not None:
                sorted_p, order = probs.sort(descending=True)
                keep = sorted_p.cumsum(0) - sorted_p < top_p        # smallest set reaching top_p
                probs = torch.zeros_like(probs).scatter(0, order[keep], sorted_p[keep])
                probs = probs / probs.sum()
            nxt = torch.multinomial(probs, 1, generator=g)
        idx = torch.cat([idx, nxt.view(1, 1)], dim=1)
    return decode(idx[0].tolist())

for label, kw in [("greedy (T=0)", dict(temperature=0)),
                  ("T=0.5", dict(temperature=0.5)),
                  ("T=1.0, top-p 0.9", dict(temperature=1.0, top_p=0.9)),
                  ("T=2.0 (too hot)", dict(temperature=2.0))]:
    text = generate(model, "\n", max_new=150, **kw)
    print(f"===== {label} =====\n{text.strip()}\n")

prompt = "add 47+38="
answer = complete(model, prompt)
full = "\n" + prompt + answer
idx = torch.tensor([encode(full)])
with torch.no_grad():
    model(idx, keep_attn=True)
labels = [c if c != "\n" else "⏎" for c in full]
q_rows = list(range(len(full) - len(answer) - 1, len(full) - 1))   # positions that predict each answer digit
print(f"{prompt}{answer}")

L, H = len(model.blocks), model.blocks[0].attn.h
fig, axes = plt.subplots(L, H, figsize=(4.2 * H, 2.5 * L))
for l in range(L):
    for h in range(H):
        A = model.blocks[l].attn.last_attn[0, h][q_rows].numpy()
        ax = axes[l, h]
        ax.imshow(A, cmap="magma", vmin=0, vmax=1, aspect="auto")
        ax.set_xticks(range(len(labels))); ax.set_xticklabels(labels, fontsize=8)
        ax.set_yticks(range(len(q_rows))); ax.set_yticklabels([f"→{answer[i]}" for i in range(len(answer))], fontsize=8)
        ax.set_title(f"layer {l} head {h}", fontsize=9)
fig.suptitle("Attention from the positions that write each answer digit (rows) to earlier characters (cols)")
plt.tight_layout(); plt.show()

@torch.no_grad()
def generate_greedy_nocache(model, prompt, n):
    idx = torch.tensor([encode(prompt)])
    for _ in range(n):
        logits, _, _ = model(idx)                       # full recompute every step
        idx = torch.cat([idx, logits[:, -1].argmax(-1, keepdim=True)], dim=1)
    return decode(idx[0].tolist())

@torch.no_grad()
def generate_greedy_cache(model, prompt, n):
    idx = torch.tensor([encode(prompt)])
    logits, _, past = model(idx)                        # process the prompt once
    out = idx
    for _ in range(n):
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        out = torch.cat([out, nxt], dim=1)
        logits, _, past = model(nxt, past=past)         # only the new token goes through the model
    return decode(out[0].tolist())

prompt, n_new = "\nadd 12+30=42\nrev hello=olleh\n", BLOCK - 32
t0 = time.time(); a = generate_greedy_nocache(model, prompt, n_new); t_no = time.time() - t0
t0 = time.time(); b = generate_greedy_cache(model, prompt, n_new); t_kv = time.time() - t0
print("identical output:", a == b)
print(f"without cache {t_no * 1000:.0f} ms, with KV cache {t_kv * 1000:.0f} ms ({t_no / t_kv:.1f}x faster)")
with torch.no_grad():
    _, _, past = model(torch.tensor([encode(prompt)]))
print(f"cache: {len(past)} layers x (keys, values), each of shape {tuple(past[0][0].shape)} = (batch, heads, tokens so far, head_dim)")

def new_task_example(split):
    en, es = random.choice(TR_TRAIN if split == "train" else TR_TEST)
    return f"en {es}=", en

random.seed(7)
ft_corpus = "\n".join("".join(new_task_example("train")) for _ in range(5000)) + "\n"
ft_data = torch.tensor(encode(ft_corpus), dtype=torch.long)
pretrained_state = {k: v.clone() for k, v in model.state_dict().items()}

def finetune(start_from_pretrained, steps=200):
    torch.manual_seed(0)
    m = GPT(len(chars), BLOCK, d=128, n_layer=4, n_head=4)
    if start_from_pretrained:
        m.load_state_dict(pretrained_state)
    o = torch.optim.AdamW(m.parameters(), lr=5e-4 if start_from_pretrained else 2e-3)
    curve = []
    for s in range(steps):
        xb, yb = get_batch(32, ft_data)
        _, loss, _ = m(xb, yb)
        o.zero_grad(); loss.backward(); o.step()
        curve.append(loss.item())
    m.eval()
    random.seed(99)
    acc = np.mean([complete(m, p) == a for p, a in (new_task_example("test") for _ in range(100))])
    return m, curve, acc

t0 = time.time()
ft_model, curve_pre, acc_pre = finetune(True)
_, curve_scratch, acc_scratch = finetune(False)
print(f"fine-tuned both in {time.time() - t0:.0f}s")
print(f"held-out ES→EN accuracy:  from pre-trained {acc_pre:.0%}   from scratch {acc_scratch:.0%}")

plt.figure(figsize=(8, 3.8))
plt.plot(np.convolve(curve_pre, np.ones(10) / 10, "valid"), label="starting from pre-trained GPT")
plt.plot(np.convolve(curve_scratch, np.ones(10) / 10, "valid"), label="starting from random weights")
plt.xlabel("fine-tuning step"); plt.ylabel("loss"); plt.legend(); plt.grid(alpha=0.3)
plt.title("Same data, same steps: pre-training gives a head start"); plt.show()

random.seed(5)
for _ in range(4):
    p, a = new_task_example("test")
    print(f"  {p}{complete(ft_model, p):24s} expected: {a}")

after = evaluate(ft_model)
print(f"{'task':8s} {'before fine-tune':>17s} {'after':>8s}")
for t in TASKS:
    print(f"{t:8s} {skill_curve[STEPS][t]:17.0%} {after[t]:8.0%}")

tokens_seen = STEPS * BATCH * BLOCK
ours = 6 * n_params * tokens_seen
rows = [("this notebook", n_params, tokens_seen),
        ("GPT-2 XL", 1.5e9, 1e10),
        ("GPT-3", 175e9, 300e9),
        ("Chinchilla", 70e9, 1.4e12)]
print(f"{'model':14s} {'params':>10s} {'tokens':>10s} {'tokens/param':>13s} {'train FLOPs (6ND)':>18s} {'× ours':>10s}")
for name, N, D in rows:
    print(f"{name:14s} {N:10.2e} {D:10.2e} {D / N:13.1f} {6 * N * D:18.2e} {6 * N * D / ours:10.1e}")