{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c024",
   "metadata": {},
   "source": [
    "# Step 5 \u2014 Assembling the GPT\n",
    "\n",
    "**Solution notebook.** Step 4's head, multiplied and wrapped into the\n",
    "complete architecture, with a parameter count verified **to the\n",
    "integer** against a formula derived on paper. No training \u2014 that is\n",
    "Step 6. Runs in seconds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "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": "c026",
   "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": "c027",
   "metadata": {},
   "source": [
    "## 5.1\u20135.3 The architecture\n",
    "\n",
    "Pre-norm blocks (Lecture 5 \u00a73): each sublayer reads a normalized copy\n",
    "of the stream and **adds** its output back, so the Jacobian is\n",
    "$I + \\partial F$ and gradients flow. All cross-position communication\n",
    "happens inside `Head`; everything else is position-wise, which is what\n",
    "makes the causality proof survive assembly."
   ]
  },
  {
   "cell_type": "code",
   "id": "c028",
   "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": "c029",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "model = GPT(V=65, T=64, d=128, H=4, L=4)\n",
    "print(model.__class__.__name__, '| blocks:', len(model.blocks))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c030",
   "metadata": {},
   "source": [
    "## 5.4 Verification 1 \u2014 the parameter count\n",
    "\n",
    "Derive on paper first (Lecture 5, Thm 5.1), then check. Getting this\n",
    "exact means you have accounted for every matrix in the model."
   ]
  },
  {
   "cell_type": "code",
   "id": "c031",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def count(mod):\n",
    "    return sum(p.numel() for p in mod.parameters())\n",
    "\n",
    "total = count(model)\n",
    "emb   = model.tok.weight.numel() + model.pos.weight.numel()\n",
    "attn  = sum(count(b.attn) for b in model.blocks)\n",
    "mlps  = sum(count(b.mlp)  for b in model.blocks)\n",
    "lns   = sum(count(b.ln1) + count(b.ln2) for b in model.blocks)\n",
    "final = count(model.lnf) + model.head.weight.numel()\n",
    "\n",
    "for name, v in [('embeddings (tok+pos)', emb), ('attention', attn),\n",
    "                ('MLPs', mlps), ('block LayerNorms', lns),\n",
    "                ('final LN + unembed', final)]:\n",
    "    print(f'  {name:<22}{v:>9,}  ({100*v/total:5.1f}%)')\n",
    "print(f'  {\"TOTAL\":<22}{total:>9,}')\n",
    "assert total == 816_640\n",
    "\n",
    "d, L = 128, 4\n",
    "print(f'\\n  12*L*d^2 rule of thumb: {12*L*d*d:,}')\n",
    "print(f'  actual block total    : {attn+mlps+lns:,}')\n",
    "print(f'  difference            : {attn+mlps+lns-12*L*d*d:,}'\n",
    "      f'  (= biases + LayerNorms, which the rule ignores)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c032",
   "metadata": {},
   "source": [
    "**The MLPs hold nearly twice the parameters of the attention layers**\n",
    "(64.5% against 32.2%) \u2014 worth remembering whenever someone calls this\n",
    "'an attention architecture'."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c033",
   "metadata": {},
   "source": [
    "## 5.4 Verification 2 \u2014 causality survives assembly\n",
    "\n",
    "Step 4's perturbation test, on the full model. It must still pass:\n",
    "residuals, LayerNorm, and MLPs are all position-wise, so the causal\n",
    "mask remains the *only* cross-position gate in the entire network."
   ]
  },
  {
   "cell_type": "code",
   "id": "c034",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "idx = ids[:32][None, :]\n",
    "with torch.no_grad():\n",
    "    o1, _ = model(idx)\n",
    "    idx2 = idx.clone(); idx2[0, 20] = (idx2[0, 20] + 7) % 65\n",
    "    o2, _ = model(idx2)\n",
    "assert torch.equal(o1[:, :20, :], o2[:, :20, :])\n",
    "assert not torch.allclose(o1[:, 20:, :], o2[:, 20:, :])\n",
    "print('causality end-to-end          OK  (positions 0-19 bitwise identical)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c035",
   "metadata": {},
   "source": [
    "## 5.4 Verification 3 \u2014 initial loss\n",
    "\n",
    "Should be just above $\\ln 65 = 4.1744$: the model starts knowing only\n",
    "the vocabulary size. It sits slightly *above* because PyTorch's default\n",
    "`nn.Linear` init gives the unembedding a std of $1/\\sqrt d$, spreading\n",
    "the logits a little; GPT-2's uniform std 0.02 would pull it to $\\ln V$\n",
    "almost exactly. A value near 8 means something is badly mis-scaled."
   ]
  },
  {
   "cell_type": "code",
   "id": "c036",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "xb, yb = ids[:64][None, :], ids[1:65][None, :]\n",
    "with torch.no_grad():\n",
    "    _, loss0 = model(xb, yb)\n",
    "print(f'initial loss on a real batch: {loss0.item():.4f}')\n",
    "print(f'ln V                        : {math.log(65):.4f}')\n",
    "assert 4.0 < loss0.item() < 4.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c037",
   "metadata": {},
   "source": [
    "### Smoke test: can it memorize one batch?\n",
    "\n",
    "50 steps on a single repeated batch. The loss should plummet \u2014 not\n",
    "because the model is good, but because it is enormously\n",
    "overparametrized relative to 4,096 tokens. This checks the plumbing\n",
    "(gradients reach every parameter) before we spend an hour on Step 6."
   ]
  },
  {
   "cell_type": "code",
   "id": "c038",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import copy\n",
    "probe = copy.deepcopy(model)\n",
    "opt = torch.optim.AdamW(probe.parameters(), lr=1e-3)\n",
    "xb, yb = ids[:64][None, :], ids[1:65][None, :]\n",
    "for i in range(51):\n",
    "    _, l = probe(xb, yb)\n",
    "    opt.zero_grad(set_to_none=True); l.backward(); opt.step()\n",
    "    if i % 25 == 0:\n",
    "        print(f'  step {i:>2}  loss {l.item():.4f}')\n",
    "assert l.item() < loss0.item() / 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c039",
   "metadata": {},
   "source": [
    "## 5.4 Verification 4 \u2014 the shape table\n",
    "\n",
    "Print the stream's shape after every stage. **This table is the\n",
    "architecture**; keep it beside the five formulas of Lecture 5 \u00a75."
   ]
  },
  {
   "cell_type": "code",
   "id": "c040",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    idx = ids[:64][None, :]\n",
    "    print(f'{\"token indices\":<28}{tuple(idx.shape)}')\n",
    "    x = model.tok(idx) + model.pos(torch.arange(64))\n",
    "    print(f'{\"after embedding\":<28}{tuple(x.shape)}')\n",
    "    for i, blk in enumerate(model.blocks):\n",
    "        x = blk(x)\n",
    "        print(f'{f\"after block {i}\":<28}{tuple(x.shape)}')\n",
    "    x = model.lnf(x)\n",
    "    print(f'{\"after final LayerNorm\":<28}{tuple(x.shape)}')\n",
    "    z = model.head(x)\n",
    "    print(f'{\"logits\":<28}{tuple(z.shape)}   <- (B, T, V)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c041",
   "metadata": {},
   "source": [
    "## The 'before' picture\n",
    "\n",
    "Samples from the untrained model: uniform gibberish over all 65\n",
    "characters. Save it \u2014 the contrast with Step 6 is the point."
   ]
  },
  {
   "cell_type": "code",
   "id": "c042",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "ctx = torch.tensor([[stoi['\\n']]])\n",
    "print(decode(model.generate(ctx, 300)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c043",
   "metadata": {},
   "source": [
    "## LayerNorm geometry (Lecture 5, Prop. 4.2)\n",
    "\n",
    "The normalization step is orthogonal projection onto $\\mathbf 1^\\perp$\n",
    "followed by radial projection onto the sphere of radius $\\sqrt d$ \u2014 so\n",
    "its image is a $(d-2)$-sphere, and sublayers see *direction*, not\n",
    "magnitude."
   ]
  },
  {
   "cell_type": "code",
   "id": "c044",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ln = nn.LayerNorm(128, elementwise_affine=False)\n",
    "# Use magnitudes in [2.5, 7.5] so every row's variance is >> eps=1e-5;\n",
    "# see the caveat below for why that matters.\n",
    "u = torch.randn(1000, 128) * (2.5 + 5 * torch.rand(1000, 1))\n",
    "y = ln(u)\n",
    "print(f'mean of output rows  : {y.mean(-1).abs().max():.2e}   (in 1-perp)')\n",
    "print(f'norms of output rows : {y.norm(dim=-1).min():.4f} .. '\n",
    "      f'{y.norm(dim=-1).max():.4f}   (sqrt(d) = {128**0.5:.4f})')\n",
    "assert torch.allclose(y.norm(dim=-1), torch.full((1000,), 128**0.5), atol=1e-3)\n",
    "assert torch.allclose(ln(3.7 * u), y, atol=1e-4)\n",
    "print('all rows on the sphere, and LN(cu) = LN(u)  OK')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c045",
   "metadata": {},
   "source": [
    "**A caveat the proof hides, and the code must not.** Proposition 4.2\n",
    "ignores the $\\varepsilon$ in $\\sqrt{\\sigma^2+\\varepsilon}$. That\n",
    "$\\varepsilon = 10^{-5}$ is what keeps LayerNorm defined when a row is\n",
    "constant ($\\sigma = 0$), but it means the sphere-radius and\n",
    "scale-invariance properties hold only *approximately*, and hold well\n",
    "precisely when $\\sigma^2 \\gg \\varepsilon$ \u2014 which is why the cell\n",
    "above deliberately keeps every row's magnitude bounded away from zero.\n",
    "Rescale a row to near-zero variance (try `u * 1e-3`) and both\n",
    "assertions fail: the normalized vector falls *short* of the sphere and\n",
    "$\\mathrm{LN}(cu) \\neq \\mathrm{LN}(u)$. In a trained model the residual\n",
    "stream has healthy variance, so this corner never bites \u2014 but it is the\n",
    "difference between the clean theorem and the floating-point object.\n",
    "\n",
    "\u2192 Continue with [Step 6](../project/step-6): train it."
   ]
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 }
}