{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# Step 6 \u2014 Training the Baby GPT\n",
    "\n",
    "**Solution notebook.** The one with the long cell: the 5,000-step\n",
    "training run takes ~1 hour on a 16-thread CPU and a few minutes on a\n",
    "Colab T4 GPU. Every other cell is seconds. The reference numbers quoted\n",
    "in the prose are from the run that produced the course's official\n",
    "checkpoint values; sampling and initialization are seeded, so your\n",
    "numbers should land within noise of them (minibatch order differs\n",
    "across hardware, so expect the last digits to move).\n",
    "\n",
    "At the end this notebook saves `step6_model.pt`, which Steps 7 and 8\n",
    "load."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c002",
   "metadata": {},
   "source": [
    "## Setup: data and tokenizer (Step 1's solution, reproduced)\n",
    "\n",
    "Every notebook in this series is self-contained: it re-creates what it\n",
    "needs from earlier steps in one compact cell, so you can run it top to\n",
    "bottom without opening the others. On Colab, uncomment the download."
   ]
  },
  {
   "cell_type": "code",
   "id": "c003",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "import torch\n",
    "\n",
    "with open('input.txt') as f:\n",
    "    text = f.read()\n",
    "vocab = sorted(set(text))\n",
    "V = len(vocab)\n",
    "stoi = {ch: i for i, ch in enumerate(vocab)}\n",
    "itos = {i: ch for i, ch in enumerate(vocab)}\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda ids: ''.join(itos[i] for i in ids)\n",
    "ids = torch.tensor(encode(text), dtype=torch.long)\n",
    "\n",
    "assert len(text) == 1_115_394 and V == 65\n",
    "print(f'{len(text):,} characters, vocab {V}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c004",
   "metadata": {},
   "source": [
    "## The model (Step 5's solution, reproduced)\n",
    "\n",
    "Identical to the Step 5 notebook, including the $1/\\sqrt{2L}$ scaling\n",
    "of the residual-writing projections \u2014 Lecture 6 \u00a71 derives why: the\n",
    "stream receives $2L$ writes, and variances add."
   ]
  },
  {
   "cell_type": "code",
   "id": "c005",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Head(nn.Module):\n",
    "    def __init__(self, d, dh, T):\n",
    "        super().__init__()\n",
    "        self.key   = nn.Linear(d, dh, bias=False)\n",
    "        self.query = nn.Linear(d, dh, bias=False)\n",
    "        self.value = nn.Linear(d, dh, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(T, T)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, _ = x.shape\n",
    "        q, k, v = self.query(x), self.key(x), self.value(x)\n",
    "        s = q @ k.transpose(-2, -1) / k.shape[-1]**0.5\n",
    "        s = s.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        A = F.softmax(s, dim=-1)\n",
    "        self.A = A.detach()          # stashed for Step 8 interpretability\n",
    "        return A @ v\n",
    "\n",
    "class MHA(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList(Head(d, d // H, T) for _ in range(H))\n",
    "        self.proj = nn.Linear(d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.fc, self.proj = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(F.gelu(self.fc(x)))\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.ln1, self.attn = nn.LayerNorm(d), MHA(d, H, T)\n",
    "        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        return x + self.mlp(self.ln2(x))\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    def __init__(self, V, T, d, H, L):\n",
    "        super().__init__()\n",
    "        self.T = T\n",
    "        self.tok = nn.Embedding(V, d)\n",
    "        self.pos = nn.Embedding(T, d)\n",
    "        self.blocks = nn.ModuleList(Block(d, H, T) for _ in range(L))\n",
    "        self.lnf = nn.LayerNorm(d)\n",
    "        self.head = nn.Linear(d, V, bias=False)\n",
    "        for blk in self.blocks:   # GPT-2 residual scaling (Lecture 6 \u00a71)\n",
    "            nn.init.normal_(blk.attn.proj.weight, std=0.02 / math.sqrt(2 * L))\n",
    "            nn.init.normal_(blk.mlp.proj.weight,  std=0.02 / math.sqrt(2 * L))\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        x = self.tok(idx) + self.pos(torch.arange(T))\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "        logits = self.head(self.lnf(x))\n",
    "        loss = None if targets is None else F.cross_entropy(\n",
    "            logits.view(-1, logits.size(-1)), targets.reshape(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, idx, n_new, temperature=1.0):\n",
    "        for _ in range(n_new):\n",
    "            logits, _ = self(idx[:, -self.T:])\n",
    "            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)\n",
    "            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "model = GPT(V=65, T=64, d=128, H=4, L=4)\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "print(f'{n_params:,} parameters')\n",
    "assert n_params == 816_640    # Step 5's checkpoint, to the integer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 6.1 Data pipeline\n",
    "\n",
    "Train/val split, then a batch sampler. `y` is `x` shifted by one:\n",
    "**every position of every sequence is a training example**, so one\n",
    "batch of shape (64, 64) contains 4,096 next-token problems. The causal\n",
    "mask (Lecture 4, Prop. 5.1) is what makes computing them all in one\n",
    "forward pass legitimate \u2014 position $i$'s prediction provably never saw\n",
    "positions $> i$."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "n_split = int(0.9 * len(ids))\n",
    "train_ids, val_ids = ids[:n_split], ids[n_split:]\n",
    "B, T = 64, 64\n",
    "\n",
    "def get_batch(split):\n",
    "    data = train_ids if split == 'train' else val_ids\n",
    "    ix = torch.randint(len(data) - T - 1, (B,))\n",
    "    x = torch.stack([data[i:i+T] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+T+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "assert torch.equal(xb[:, 1:], yb[:, :-1])   # y really is x shifted\n",
    "print(xb.shape, yb.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 6.2 Evaluation done right\n",
    "\n",
    "One minibatch loss is a noisy estimate (Lecture 3, Prop. 5.2 \u2014 variance\n",
    "$\\propto 1/B$). We average many batches, in eval mode, under\n",
    "`no_grad`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss(iters=100):\n",
    "    model.eval()\n",
    "    out = {}\n",
    "    for split in ('train', 'val'):\n",
    "        losses = [model(*get_batch(split))[1].item() for _ in range(iters)]\n",
    "        out[split] = sum(losses) / len(losses)\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 6.3 The training loop\n",
    "\n",
    "The three ingredients, each earned in lecture:\n",
    "\n",
    "- **AdamW** (Lecture 6 \u00a72): diagonal preconditioning so rare-feature\n",
    "  coordinates still move \u2014 the fix for exactly the pathology measured in\n",
    "  Step 2 \u2014 with *decoupled* weight decay.\n",
    "- **Warmup + cosine decay** (Lecture 6 \u00a73), set by hand on the\n",
    "  optimizer's param group so nothing is hidden.\n",
    "- **Gradient clipping** (Lecture 6 \u00a74) at norm 1.0 \u2014 and we log the\n",
    "  pre-clip norm every evaluation, because *whether the safeguard ever\n",
    "  fires* is data worth having. (Spoiler from the reference run: it\n",
    "  never does; the norm stays in 0.30\u20130.41. Instrumenting a safeguard to\n",
    "  discover it never fires beats assuming it was load-bearing.)"
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "MAX_STEPS, WARMUP = 5000, 100\n",
    "LR_MAX, LR_MIN = 3e-3, 3e-4\n",
    "\n",
    "def lr_at(t):\n",
    "    if t < WARMUP:\n",
    "        return LR_MAX * t / WARMUP\n",
    "    r = (t - WARMUP) / (MAX_STEPS - WARMUP)\n",
    "    return LR_MIN + 0.5 * (1 + math.cos(math.pi * r)) * (LR_MAX - LR_MIN)\n",
    "\n",
    "opt = torch.optim.AdamW(model.parameters(), lr=LR_MAX,\n",
    "                        weight_decay=0.1, betas=(0.9, 0.99))"
   ]
  },
  {
   "cell_type": "code",
   "id": "c013",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import time\n",
    "history = []          # (step, train, val, grad_norm, lr)\n",
    "t0, gnorm = time.time(), float('nan')\n",
    "\n",
    "for step in range(MAX_STEPS + 1):\n",
    "    if step % 500 == 0 or step == MAX_STEPS:\n",
    "        e = estimate_loss(50 if step < MAX_STEPS else 200)\n",
    "        history.append((step, e['train'], e['val'], gnorm, lr_at(step)))\n",
    "        print(f\"{step:>5}  train {e['train']:.4f}  val {e['val']:.4f}\"\n",
    "              f\"  ({e['val']/math.log(2):.4f} bits/char)\"\n",
    "              f\"  gnorm {gnorm:.3f}  lr {lr_at(step):.2e}\"\n",
    "              f\"  {time.time()-t0:.0f}s\", flush=True)\n",
    "    if step == MAX_STEPS:\n",
    "        break\n",
    "    for g in opt.param_groups:\n",
    "        g['lr'] = lr_at(step)\n",
    "    _, loss = model(*get_batch('train'))\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    gnorm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0))\n",
    "    opt.step()\n",
    "\n",
    "torch.save(model.state_dict(), 'step6_model.pt')\n",
    "print('saved step6_model.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c014",
   "metadata": {},
   "source": [
    "Reference trajectory (validation): 4.304 \u2192 1.813 \u2192 1.661 \u2192 1.609 \u2192\n",
    "1.568 \u2192 1.538 \u2192 1.521 \u2192 1.520 \u2192 1.519 \u2192 **1.509** \u2192 1.518. Final:\n",
    "**1.5178 nats = 2.190 bits/char**. Your bits/char table gains its\n",
    "biggest single improvement: 3.54 (bigram) \u2192 2.86 (MLP) \u2192 **2.19**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 6.4 Read the curves\n",
    "\n",
    "Four plots, and the prose you attach to them matters more than the\n",
    "plots."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "steps, tr, va, gn, lrs = zip(*history)\n",
    "fig, ax = plt.subplots(2, 2, figsize=(11, 7))\n",
    "ax[0,0].plot(steps, tr, label='train'); ax[0,0].plot(steps, va, label='val')\n",
    "ax[0,0].axhline(math.log(65), ls=':', c='gray', label='ln V')\n",
    "ax[0,0].set_title('loss'); ax[0,0].legend()\n",
    "ax[0,1].plot(steps, [v-t for t,v in zip(tr,va)], c='crimson')\n",
    "ax[0,1].set_title('generalization gap (val \u2212 train)')\n",
    "ax[1,0].plot(steps[1:], gn[1:], c='darkorange')\n",
    "ax[1,0].axhline(1.0, ls=':', c='gray', label='clip threshold')\n",
    "ax[1,0].set_title('pre-clip gradient norm'); ax[1,0].legend()\n",
    "ax[1,1].plot(steps, lrs, c='seagreen'); ax[1,1].set_title('learning rate')\n",
    "for a in ax.flat: a.set_xlabel('step')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "What the reference run shows, panel by panel:\n",
    "\n",
    "1. **Loss.** Starts at $\\ln 65$ (the model knows only the vocabulary\n",
    "   size), drops to 1.81 within 500 steps, then grinds. Roughly linear\n",
    "   in $\\log(\\text{step})$ \u2014 a within-run shadow of Lecture 7's\n",
    "   scaling laws.\n",
    "2. **Gap.** Widens monotonically, 0.158 \u2192 0.352. At ~0.8 parameters\n",
    "   per training character the model can memorize, and is beginning to.\n",
    "   Also look closely at val alone: it stops improving near step 3000\n",
    "   and its **minimum is at 4500, not the end** \u2014 the last 40% of the\n",
    "   run bought ~nothing, and the final model is not the best model.\n",
    "   Real pipelines checkpoint on validation for exactly this reason.\n",
    "3. **Gradient norm.** Never approaches the clip threshold \u2014 clipping\n",
    "   was pure insurance here. But it *rises* late (0.298 \u2192 0.414) while\n",
    "   the learning rate falls tenfold: the iterate is settling into\n",
    "   *sharper* curvature as steps shrink. That is edge-of-stability\n",
    "   behaviour (Lecture 6 \u00a75.4), visible on a laptop.\n",
    "4. **Learning rate.** The warmup spike is invisible at this resolution\n",
    "   (100 steps); the cosine decay is not."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c018",
   "metadata": {},
   "source": [
    "## Samples\n",
    "\n",
    "The payoff cell. Reference output includes real, correctly spelled\n",
    "speaker names \u2014 `WARWICK:`, `Citizen:`, `Montague`, `Barnardine`,\n",
    "`Volsces` all occur in the corpus (99, 98, 46, 15, 18 times). Compare\n",
    "Step 3's `LOET:`, a name-shaped fake \u2014 three characters of context\n",
    "cannot hold an identity, 64 can. That difference is attention, visible\n",
    "in the output."
   ]
  },
  {
   "cell_type": "code",
   "id": "c019",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "ctx = torch.tensor([[stoi['\\n']]])\n",
    "print(decode(model.generate(ctx, 500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c020",
   "metadata": {},
   "source": [
    "## 6.5 One controlled experiment\n",
    "\n",
    "The reference choice: depth $L \\in \\{1, 2, 4, 8\\}$ at fixed width.\n",
    "This cell takes several times longer than the main run \u2014 leave it for\n",
    "a GPU session, or shorten `MAX_STEPS`. (Left unexecuted here; pool\n",
    "your results with the class for Step 7's scaling plot.)\n",
    "\n",
    "```python\n",
    "for L in (1, 2, 4, 8):\n",
    "    torch.manual_seed(1337)\n",
    "    model = GPT(V=65, T=64, d=128, H=4, L=L)\n",
    "    # ... identical training loop ...\n",
    "```\n",
    "\n",
    "Whatever knob you choose, hold everything else fixed, tabulate\n",
    "(val loss, params, wallclock), and write three sentences on whether\n",
    "the trend is monotone \u2014 Step 3's $k$-sweep should have taught you not\n",
    "to assume it is.\n",
    "\n",
    "\u2192 Continue with [Step 7](../project/step-7): a real tokenizer and a\n",
    "real sampler for this trained model."
   ]
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 }
}