{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-1 fw=all -->\n",
    "\n",
    "# A GPT from Scratch\n",
    "\n",
    "One block, stacked, is most of a language model. This section adds the\n",
    "little that remains — an embedding, positions, a causal mask, an output\n",
    "head — and packages the result as a single `GPT` class whose constructor\n",
    "arguments span the design space: the block flags of\n",
    "that section plus a positional scheme. That\n",
    "one-class-many-configurations shape is the spine of this chapter. The same\n",
    "class, with the modern flags, trains from scratch on *The Time Machine* in\n",
    "about a minute; with the 2019 flags, it accepts the released weights of\n",
    "GPT-2 [@Radford.Wu.Child.ea.2019] and completes English sentences.\n",
    "Between those two demonstrations sit the practical crafts this section\n",
    "teaches: reading a loss curve for what it actually shows, breaking a model\n",
    "with a normalization flag, and sampling from a trained distribution."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-a-gpt-from-scratch"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "from flax import nnx\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import optax\n",
    "from safetensors import numpy as safetensors_numpy\n",
    "import tiktoken\n",
    "from tiktoken.load import data_gym_to_mergeable_bpe_ranks"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-2 fw=all -->\n",
    "\n",
    "## From Blocks to a Language Model\n",
    "\n",
    "Four ingredients turn a stack of transformer blocks into a GPT-style\n",
    "language model.\n",
    "\n",
    "**A token embedding and a tied head.** A lookup table maps token ids into\n",
    "the residual stream; after the final block (and one last normalization,\n",
    "pre-norm convention), the *same* table, transposed, maps the stream back to\n",
    "logits over the vocabulary. Tying the two matrices saves a\n",
    "$\\textrm{vocab} \\times d$ parameter block and consistently helps at small\n",
    "scale [@Press.Wolf.2017]; GPT-2 ties them too, which will spare us a\n",
    "tensor when we load its weights.\n",
    "\n",
    "**A causal mask.** A language model predicts token $t+1$ from tokens $1,\n",
    "\\ldots, t$, so position $t$ must not attend forward. In\n",
    "that section we got causality by handing\n",
    "`d2l.MultiHeadAttention` a per-query valid length. Here we build the causal\n",
    "variant natively, for two reasons. First, efficiency: the fused attention\n",
    "kernels of that section\n",
    "(`scaled_dot_product_attention` in PyTorch, `dot_product_attention` in\n",
    "JAX) take the causal mask as a flag and never materialize the $n \\times n$\n",
    "score matrix. Second, positions.\n",
    "\n",
    "**Positions, inside attention or outside.** The `pos` argument selects\n",
    "between the two schemes that that section found in\n",
    "deployed models: `'learned'` adds a trained position vector to the\n",
    "embedding — GPT-2's scheme, capped at the table's length — and `'rope'`\n",
    "rotates queries and keys inside every attention head, which is what\n",
    "essentially every current model does. RoPE lives where queries and keys\n",
    "are made, so the attention module implements it itself (the same\n",
    "`_rope` we gave `TinyCharLM`); nothing else in the model knows positions\n",
    "exist.\n",
    "\n",
    "**The blocks themselves** are `d2l.TransformerBlock`, unchanged: the\n",
    "causal attention drops in through the `attn_factory` hook, exactly the\n",
    "seam it was designed to be."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-from-blocks-to-a-language-model-1"
   },
   "source": [
    "class GPT(nnx.Module):\n",
    "    \"\"\"Decoder-only transformer language model built from configurable\n",
    "    blocks.\"\"\"\n",
    "\n",
    "    class CausalAttention(nnx.Module):\n",
    "        \"\"\"Multi-head causal self-attention, optionally rotary.\"\"\"\n",
    "        def __init__(self, num_hiddens, num_heads, bias=False, rope=False,\n",
    "                     rngs=None):\n",
    "            rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "            self.num_heads, self.rope = num_heads, rope\n",
    "            self.W_qkv = nnx.Linear(num_hiddens, 3 * num_hiddens,\n",
    "                                    use_bias=bias, rngs=rngs)\n",
    "            self.W_o = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,\n",
    "                                  rngs=rngs)\n",
    "\n",
    "        def _rope(self, x):\n",
    "            # x: (batch, num_steps, num_heads, head_dim)\n",
    "            d = x.shape[-1]\n",
    "            pos = jnp.arange(x.shape[1], dtype=jnp.float32)\n",
    "            inv_freq = 10000.0 ** (-jnp.arange(0, d, 2) / d)\n",
    "            theta = pos[:, None] * inv_freq[None, :]\n",
    "            cos = jnp.cos(theta)[:, None, :]  # broadcast over heads\n",
    "            sin = jnp.sin(theta)[:, None, :]\n",
    "            x1, x2 = x[..., 0::2], x[..., 1::2]\n",
    "            return jnp.stack([x1 * cos - x2 * sin,\n",
    "                              x1 * sin + x2 * cos], -1).reshape(x.shape)\n",
    "\n",
    "        def __call__(self, X, *_):\n",
    "            B, T, D = X.shape\n",
    "            q, k, v = jnp.split(self.W_qkv(X), 3, axis=-1)\n",
    "            q, k, v = (u.reshape(B, T, self.num_heads, -1)\n",
    "                       for u in (q, k, v))\n",
    "            if self.rope:\n",
    "                q, k = self._rope(q), self._rope(k)\n",
    "            Y = jax.nn.dot_product_attention(q, k, v, is_causal=True)\n",
    "            return self.W_o(Y.reshape(B, T, D)), None\n",
    "\n",
    "    def __init__(self, vocab_size, num_hiddens=256, num_heads=8, num_blks=6,\n",
    "                 max_len=1024, pos='rope', norm='rms', act='swiglu',\n",
    "                 pre_norm=True, bias=False, dropout=0, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.pos, self.max_len = pos, max_len\n",
    "        init = nnx.initializers.normal(0.02)\n",
    "        self.token_emb = nnx.Embed(vocab_size, num_hiddens,\n",
    "                                   embedding_init=init, rngs=rngs)\n",
    "        if pos == 'learned':\n",
    "            self.pos_emb = nnx.Embed(max_len, num_hiddens,\n",
    "                                     embedding_init=init, rngs=rngs)\n",
    "        attn = lambda rngs: self.CausalAttention(\n",
    "            num_hiddens, num_heads, bias, rope=(pos == 'rope'), rngs=rngs)\n",
    "        self.blks = nnx.List([\n",
    "            d2l.TransformerBlock(num_hiddens, num_heads, dropout, norm, act,\n",
    "                                 pre_norm, bias, attn_factory=attn,\n",
    "                                 rngs=rngs)\n",
    "            for _ in range(num_blks)])\n",
    "        self.norm = (nnx.RMSNorm if norm == 'rms'\n",
    "                     else nnx.LayerNorm)(num_hiddens, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        H = self.token_emb(X)\n",
    "        if self.pos == 'learned':\n",
    "            H = H + self.pos_emb(jnp.arange(X.shape[1]))\n",
    "        for blk in self.blks:\n",
    "            H = blk(H)\n",
    "        return self.token_emb.attend(self.norm(H))"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-3 fw=all -->\n",
    "\n",
    "A shape check, and the census of the default configuration:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-from-blocks-to-a-language-model-2"
   },
   "source": [
    "model = GPT(vocab_size=28)\n",
    "X = jnp.zeros((2, 16), dtype=jnp.int32)\n",
    "d2l.check_shape(model(X), (2, 16, 28))\n",
    "n = sum(p.size for p in jax.tree.leaves(nnx.state(model, nnx.Param)))\n",
    "print(f'{n / 1e6:.2f}M parameters')"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "4.73M parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-4 fw=all -->\n",
    "\n",
    "One decision above deserves a defense: the model reads *characters*, not\n",
    "subwords. Deployed language models tokenize into subwords — byte-pair\n",
    "encoding is the most common scheme — and that section built\n",
    "a full BPE tokenizer; we will even reuse its GPT-2 pattern verbatim when we\n",
    "load GPT-2 below. But BPE earns\n",
    "its keep on gigabytes; on a 180 KB novel a subword vocabulary would leave\n",
    "each token type a handful of training examples. Characters keep the\n",
    "statistics dense, the vocabulary trivial, and the comparison fair against\n",
    "the character-level models of previous chapters.\n",
    "\n",
    "## Training the Modern Configuration\n",
    "\n",
    "We train the default configuration — pre-norm, RMSNorm, SwiGLU, RoPE:\n",
    "the same flag settings you would find in a Llama or Qwen checkpoint\n",
    "[@touvron2023llama] — on the character-level Time Machine corpus,\n",
    "with one concession to its size: `dropout=0.1`, a regularizer that\n",
    "frontier models dropped precisely because their corpora outweigh their\n",
    "parameters, which is the opposite of our situation here. We watch the\n",
    "validation loss as we go rather than admiring the training curve alone."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-training-the-modern-configuration"
   },
   "source": [
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char',\n",
    "                       num_train=100000, num_val=3000)\n",
    "\n",
    "@nnx.jit\n",
    "def batch_loss(model, X, Y):\n",
    "    logits = model(X)\n",
    "    return optax.softmax_cross_entropy_with_integer_labels(\n",
    "        logits.reshape(-1, logits.shape[-1]), Y.reshape(-1)).mean()\n",
    "\n",
    "def val_loss(model, data):\n",
    "    \"\"\"Mean per-token validation loss (weighted by batch size: the final\n",
    "    batch may be smaller).\"\"\"\n",
    "    model.eval()\n",
    "    total = count = 0\n",
    "    for X, Y in data.val_dataloader():\n",
    "        Y = jnp.asarray(Y)\n",
    "        total += float(batch_loss(model, jnp.asarray(X), Y)) * Y.size\n",
    "        count += Y.size\n",
    "    model.train()\n",
    "    return total / count\n",
    "\n",
    "model = GPT(len(data.vocab), dropout=0.1, rngs=nnx.Rngs(0))\n",
    "optimizer = nnx.Optimizer(model, optax.adamw(1e-3, weight_decay=0.0),\n",
    "                          wrt=nnx.Param)\n",
    "steps, train_curve, val_curve = [], [], []\n",
    "for chunk in range(10):\n",
    "    losses = d2l.train_lm(model, data, optimizer, 200)\n",
    "    steps.append(200 * (chunk + 1))\n",
    "    train_curve.append(sum(losses[-50:]) / 50)\n",
    "    val_curve.append(val_loss(model, data))\n",
    "d2l.plot(steps, [train_curve, val_curve], 'step', 'loss',\n",
    "         legend=['train', 'validation'])\n",
    "print(f'final: train {train_curve[-1]:.2f}, '\n",
    "      f'validation {val_curve[-1]:.2f}; '\n",
    "      f'best validation {min(val_curve):.2f} at step '\n",
    "      f'{steps[val_curve.index(min(val_curve))]}')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "final: train 0.15, validation 3.26; best validation 1.47 at step 200\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"240.588352pt\" height=\"183.35625pt\" viewBox=\"0 0 240.588352 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T15:16:43.398353</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 240.588352 183.35625 \nL 240.588352 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 34.240625 145.8 \nL 229.540625 145.8 \nL 229.540625 7.2 \nL 34.240625 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path d=\"M 72.708807 145.8 \nL 72.708807 7.2 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"m4fe3d32905\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m4fe3d32905\" x=\"72.708807\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 500 -->\n      <g transform=\"translate(63.165057 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-35\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path d=\"M 122.026989 145.8 \nL 122.026989 7.2 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m4fe3d32905\" x=\"122.026989\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 1000 -->\n      <g transform=\"translate(109.301989 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(190.869141 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path d=\"M 171.34517 145.8 \nL 171.34517 7.2 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m4fe3d32905\" x=\"171.34517\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 1500 -->\n      <g transform=\"translate(158.62017 160.398438) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(190.869141 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path d=\"M 220.663352 145.8 \nL 220.663352 7.2 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m4fe3d32905\" x=\"220.663352\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 2000 -->\n      <g transform=\"translate(207.938352 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(190.869141 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_5\">\n     <!-- step -->\n     <g transform=\"translate(121.075 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-73\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(52.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(91.308594 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(152.832031 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_9\">\n      <path d=\"M 34.240625 145.404365 \nL 229.540625 145.404365 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <defs>\n       <path id=\"m78c99243a6\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m78c99243a6\" x=\"34.240625\" y=\"145.404365\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 0 -->\n      <g transform=\"translate(20.878125 149.203584) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_11\">\n      <path d=\"M 34.240625 105.53029 \nL 229.540625 105.53029 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#m78c99243a6\" x=\"34.240625\" y=\"105.53029\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 1 -->\n      <g transform=\"translate(20.878125 109.329508) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_13\">\n      <path d=\"M 34.240625 65.656214 \nL 229.540625 65.656214 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#m78c99243a6\" x=\"34.240625\" y=\"65.656214\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 2 -->\n      <g transform=\"translate(20.878125 69.455433) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_15\">\n      <path d=\"M 34.240625 25.782139 \nL 229.540625 25.782139 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#m78c99243a6\" x=\"34.240625\" y=\"25.782139\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 3 -->\n      <g transform=\"translate(20.878125 29.581358) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-33\" d=\"M 2597 2516 \nQ 3050 2419 3304 2112 \nQ 3559 1806 3559 1356 \nQ 3559 666 3084 287 \nQ 2609 -91 1734 -91 \nQ 1441 -91 1130 -33 \nQ 819 25 488 141 \nL 488 750 \nQ 750 597 1062 519 \nQ 1375 441 1716 441 \nQ 2309 441 2620 675 \nQ 2931 909 2931 1356 \nQ 2931 1769 2642 2001 \nQ 2353 2234 1838 2234 \nL 1294 2234 \nL 1294 2753 \nL 1863 2753 \nQ 2328 2753 2575 2939 \nQ 2822 3125 2822 3475 \nQ 2822 3834 2567 4026 \nQ 2313 4219 1838 4219 \nQ 1578 4219 1281 4162 \nQ 984 4106 628 3988 \nL 628 4550 \nQ 988 4650 1302 4700 \nQ 1616 4750 1894 4750 \nQ 2613 4750 3031 4423 \nQ 3450 4097 3450 3541 \nQ 3450 3153 3228 2886 \nQ 3006 2619 2597 2516 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-33\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_10\">\n     <!-- loss -->\n     <g transform=\"translate(14.798438 86.157813) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-6c\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(27.783203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(88.964844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(141.064453 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_17\">\n    <path d=\"M 43.117898 90.2833 \nL 62.84517 108.73387 \nL 82.572443 123.689647 \nL 102.299716 132.880091 \nL 122.026989 136.366572 \nL 141.754261 137.903965 \nL 161.481534 138.513244 \nL 181.208807 138.939239 \nL 200.93608 139.268592 \nL 220.663352 139.5 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_18\">\n    <path d=\"M 43.117898 86.668844 \nL 62.84517 80.955127 \nL 82.572443 63.158146 \nL 102.299716 42.129833 \nL 122.026989 33.311544 \nL 141.754261 24.227097 \nL 161.481534 18.666264 \nL 181.208807 15.742292 \nL 200.93608 13.5 \nL 220.663352 15.481628 \n\" clip-path=\"url(#p89c97c159d)\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 34.240625 145.8 \nL 34.240625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 229.540625 145.8 \nL 229.540625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 34.240625 145.8 \nL 229.540625 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 34.240625 7.2 \nL 229.540625 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 141.30625 92.678125 \nL 222.540625 92.678125 \nQ 224.540625 92.678125 224.540625 90.678125 \nL 224.540625 62.321875 \nQ 224.540625 60.321875 222.540625 60.321875 \nL 141.30625 60.321875 \nQ 139.30625 60.321875 139.30625 62.321875 \nL 139.30625 90.678125 \nQ 139.30625 92.678125 141.30625 92.678125 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_19\">\n     <path d=\"M 143.30625 68.420313 \nL 153.30625 68.420313 \nL 163.30625 68.420313 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_11\">\n     <!-- train -->\n     <g transform=\"translate(171.30625 71.920313) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6e\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-74\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(39.208984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(80.322266 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(141.601562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(169.384766 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_20\">\n     <path d=\"M 143.30625 83.098438 \nL 153.30625 83.098438 \nL 163.30625 83.098438 \n\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_12\">\n     <!-- validation -->\n     <g transform=\"translate(171.30625 86.598438) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-76\" d=\"M 191 3500 \nL 800 3500 \nL 1894 563 \nL 2988 3500 \nL 3597 3500 \nL 2284 0 \nL 1503 0 \nL 191 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-64\" d=\"M 2906 2969 \nL 2906 4863 \nL 3481 4863 \nL 3481 0 \nL 2906 0 \nL 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nz\nM 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-76\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(59.179688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(120.458984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(148.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(176.025391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(239.501953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(300.78125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(339.990234 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(367.773438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(428.955078 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p89c97c159d\">\n   <rect x=\"34.240625\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-5 fw=all -->\n",
    "\n",
    "Read the two curves separately, because they tell different stories. The\n",
    "training loss falls smoothly toward a few tenths of a nat — the\n",
    "architecture works, the optimizer works. The validation loss bottoms out\n",
    "around $1.5$ nats within the first few hundred steps and then *rises* while\n",
    "training keeps improving. No bug: our run feeds the model roughly 16\n",
    "million tokens drawn from a corpus of one hundred thousand characters,\n",
    "about 160 passes over the book, with 4.7 million parameters to spend —\n",
    "dozens of parameters per unique character of text. Past the first epochs,\n",
    "gradient descent has nothing left to learn from this book except the book\n",
    "itself. The cure is not fewer steps or a bigger dropout but\n",
    "*more data*; how loss actually scales when data and parameters grow\n",
    "together is the business of this chapter's closing section on scaling laws\n",
    "[@kaplan2020scaling; @hoffmann2022training].\n",
    "\n",
    "It is worth locating this run on the cost map. At roughly $6ND$\n",
    "floating-point operations for training a model of $N$ parameters on $D$\n",
    "tokens, our minute of GPU time spent about $5 \\times 10^{14}$ FLOPs. The\n",
    "124M-parameter GPT-2 we load below — same class, about 26 times our\n",
    "parameter count, a corpus five orders of magnitude larger — cost on the\n",
    "order of $7 \\times 10^{18}$ by the same estimate, taking WebText at\n",
    "roughly ten billion tokens (the 1.5B-parameter GPT-2 XL lands near\n",
    "$10^{20}$), and frontier runs sit around $10^{25}$ or beyond. Across\n",
    "those ten orders of magnitude the model definition barely changes —\n",
    "block, mask, embedding, head — but everything around it does: data\n",
    "pipelines, custom kernels, and the parallelism of\n",
    "that section. The block abstraction is the part that\n",
    "transfers, which is why it is worth learning on a novella.\n",
    "\n",
    "### Breaking It with One Flag\n",
    "\n",
    "that section predicted, from initialization statistics\n",
    "alone, that the post-LN arrangement starves its attention layers' query\n",
    "and key projections of gradient. Now that we own a trainable model, we\n",
    "can watch the prediction come true. Same model, same data, same 800\n",
    "steps at learning rate\n",
    "$3 \\times 10^{-3}$ — three times the rate above, still comfortable for\n",
    "pre-norm — with `pre_norm` flipped:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-breaking-it-with-one-flag"
   },
   "source": [
    "for pre_norm in (True, False):\n",
    "    ablation = GPT(len(data.vocab), pre_norm=pre_norm, rngs=nnx.Rngs(0))\n",
    "    optimizer = nnx.Optimizer(ablation, optax.adamw(3e-3, weight_decay=0.0),\n",
    "                              wrt=nnx.Param)\n",
    "    losses = d2l.train_lm(ablation, data, optimizer, 800)\n",
    "    print(f'pre_norm={pre_norm}: loss at step 100/200/400/800: ' +\n",
    "          '/'.join(f'{sum(losses[k-50:k]) / 50:.2f}'\n",
    "                   for k in (100, 200, 400, 800)))"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "pre_norm=True: loss at step 100/200/400/800: 1.94/1.48/1.10/0.54\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "pre_norm=False: loss at step 100/200/400/800: 2.83/2.83/2.83/2.83\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-6 fw=all -->\n",
    "\n",
    "The pre-norm model trains as if nothing happened. The post-norm model does\n",
    "not diverge — no NaNs, no explosion — it does something more telling: it\n",
    "sticks at about $2.8$ nats and never leaves. That number is exactly the\n",
    "unigram entropy of this text ($2.83$ nats: predict letter frequencies,\n",
    "ignore all context), and a language model pinned there is a model whose\n",
    "attention has learned nothing — the training-time realization of the\n",
    "starved query/key gradients we measured at initialization. At the gentler\n",
    "learning rate of the previous run, post-LN does train, trailing early\n",
    "(this is what learning-rate warmup was invented for\n",
    "[@xiong2020layer]); at the rate pre-norm shrugs off, it fails\n",
    "outright. GPT-2's quiet move of the normalization, ahead of most of the\n",
    "field, is one reason its 48-block variant was trainable at all in 2019\n",
    "[@Radford.Wu.Child.ea.2019].\n",
    "\n",
    "## Sampling from the Model\n",
    "\n",
    "A language model outputs a distribution; text comes from *decoding* it.\n",
    "that section covered the search view — greedy, beam,\n",
    "sampling-with-temperature — for sequence models in general. For open-ended\n",
    "generation the workhorse is temperature-plus-truncation: divide the logits\n",
    "by a temperature $\\tau$, optionally keep only the $k$ most probable tokens\n",
    "[@Fan.Lewis.Dauphin.2018] (or the smallest set covering probability\n",
    "$p$, *nucleus sampling* [@Holtzman.Buys.Du.ea.2020] — an exercise),\n",
    "and sample. We add `generate` to the class. It is deliberately naive:\n",
    "every new token reruns the full forward pass over the whole history,\n",
    "which is quadratic work per token — measuring and then eliminating that\n",
    "waste is the entire next section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-sampling-from-the-model-1"
   },
   "source": [
    "@d2l.add_to_class(GPT)\n",
    "def generate(self, prefix, num_tokens, seed=0, temperature=1.0,\n",
    "             top_k=None):\n",
    "    \"\"\"Sample a continuation of the token-id list prefix.\"\"\"\n",
    "    total = len(prefix) + num_tokens\n",
    "    buf = jnp.zeros((1, total), dtype=jnp.int32)\n",
    "    buf = buf.at[0, :len(prefix)].set(jnp.asarray(prefix))\n",
    "\n",
    "    @nnx.jit\n",
    "    def logits_at(model, buf, t):\n",
    "        return model(buf)[0, t - 1]\n",
    "\n",
    "    ids, key = list(prefix), jax.random.key(seed)\n",
    "    for t in range(len(prefix), total):\n",
    "        logits = logits_at(self, buf, t) / temperature\n",
    "        if top_k is not None:\n",
    "            logits = jnp.where(logits < jnp.sort(logits)[-top_k],\n",
    "                               -jnp.inf, logits)\n",
    "        key, sub = jax.random.split(key)\n",
    "        ids.append(int(jax.random.categorical(sub, logits)))\n",
    "        buf = buf.at[0, t].set(ids[-1])\n",
    "    return ids"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-7 fw=jax -->\n",
    "\n",
    "One JAX-specific choice: rather than growing the input by one token per\n",
    "step — which would trigger a fresh XLA compilation at every new length —\n",
    "we allocate a buffer of the final size once and overwrite it left to\n",
    "right. The causal mask makes the not-yet-written positions invisible to\n",
    "every query before them, so the logits at position $t-1$ are exact, and\n",
    "the jitted forward compiles exactly once. Fixed shapes are how real JAX\n",
    "serving systems work too, as the next section shows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-8 fw=all -->\n",
    "\n",
    "Let's hear what 4.7 million parameters trained for a minute sound like:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-sampling-from-the-model-2"
   },
   "source": [
    "model.eval()\n",
    "prefix = data.vocab[list('the time traveller ')]\n",
    "for temperature, top_k in ((1.0, None), (0.7, 8), (2.0, None)):\n",
    "    out = model.generate(prefix, 120, seed=0, temperature=temperature,\n",
    "                         top_k=top_k)\n",
    "    print(f'T={temperature}, top_k={top_k}: '\n",
    "          + repr(''.join(data.vocab.to_tokens(out))))"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=1.0, top_k=None: 'the time traveller holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=0.7, top_k=8: 'the time traveller holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=2.0, top_k=None: 'the time traveller holding the lamp aloft i intend to explore time is that plain i was never more serious in my life none of us quite knew '\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-9 fw=all -->\n",
    "\n",
    "The samples are fluent pseudo-Wells — and much of that fluency appears\n",
    "memorized rather than composed: the overfitting that the validation curve\n",
    "reported, made audible. Note how little the temperature changes the early\n",
    "tokens. A memorizing model is a *confident* model, and dividing\n",
    "near-one-hot logits\n",
    "by $0.7$ or $2$ barely moves them; the knobs only start to matter once\n",
    "the model is genuinely uncertain. We need a model whose uncertainty is\n",
    "real — so let's load one.\n",
    "\n",
    "## Loading GPT-2\n",
    "\n",
    "GPT-2 is the type specimen of this section's class: 12 blocks, 768 hidden\n",
    "units, 12 heads, LayerNorm (with biases), GELU, learned positions, tied\n",
    "head — and pre-norm, years before that was consensus. Its 124M-parameter\n",
    "release is a constructor call away; all we owe it is its exact flags, plus\n",
    "its weights and its tokenizer, both of which we pin by content hash in\n",
    "`d2l.DATA_HUB` so that one download serves every later run and framework."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-loading-gpt-2-1"
   },
   "source": [
    "HF_URL = 'https://huggingface.co/openai-community/gpt2/resolve/main/'\n",
    "d2l.DATA_HUB['gpt2-weights'] = (\n",
    "    HF_URL + 'model.safetensors',\n",
    "    '89a76996d7c6ee89b86618a265483aab73e61d50')\n",
    "d2l.DATA_HUB['gpt2-merges'] = (\n",
    "    HF_URL + 'merges.txt', '396d4d8ec90cb02f4d56e049e0e4add868bcd943')\n",
    "d2l.DATA_HUB['gpt2-encoder'] = (\n",
    "    HF_URL + 'vocab.json', 'f0223209235343bc067d7da838328bced8085ae1')\n",
    "\n",
    "enc = tiktoken.Encoding(\n",
    "    'gpt2', explicit_n_vocab=50257,\n",
    "    pat_str=d2l.BPETokenizer.GPT2_PATTERN,\n",
    "    mergeable_ranks=data_gym_to_mergeable_bpe_ranks(\n",
    "        d2l.download('gpt2-merges', '../data/gpt2'),\n",
    "        d2l.download('gpt2-encoder', '../data/gpt2')),\n",
    "    special_tokens={'<|endoftext|>': 50256})\n",
    "ids = enc.encode('Attention is all you need.')\n",
    "print(ids, [enc.decode([i]) for i in ids])"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[8086, 1463, 318, 477, 345, 761, 13] ['Att', 'ention', ' is', ' all', ' you', ' need', '.']\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-10 fw=all -->\n",
    "\n",
    "The tokenizer cell is the BPE machinery of that section\n",
    "meeting the real artifact: GPT-2's released merge list and vocabulary,\n",
    "interpreted with the *same* pre-tokenization pattern our own\n",
    "`d2l.BPETokenizer` uses, assembled into `tiktoken`'s fast encoder. No\n",
    "model library, no configuration framework — two data files and a regular\n",
    "expression.\n",
    "\n",
    "Now the weights. The checkpoint stores one tensor per parameter under\n",
    "names like `h.3.attn.c_attn.weight`; our job is a dictionary mapping those\n",
    "names onto our modules. One historical trap: GPT-2's original code\n",
    "implemented linear layers as a `Conv1D` class that keeps weights as\n",
    "$(\\textrm{in}, \\textrm{out})$ — the transpose of `nn.Linear`'s\n",
    "$(\\textrm{out}, \\textrm{in})$ — so every 2-D weight must be transposed on\n",
    "the way in. The fused `c_attn` matrix concatenates queries, keys, and\n",
    "values in the same order as our `W_qkv`, and the token embedding doubles\n",
    "as the output head in both models, so neither needs special handling."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-loading-gpt-2-2"
   },
   "source": [
    "gpt2 = GPT(vocab_size=50257, num_hiddens=768, num_heads=12, num_blks=12,\n",
    "           max_len=1024, pos='learned', norm='layer', act='gelu',\n",
    "           pre_norm=True, bias=True)  # the GPT-2 (124M) configuration\n",
    "weights = safetensors_numpy.load_file(\n",
    "    d2l.download('gpt2-weights', '../data/gpt2'))\n",
    "\n",
    "gpt2.token_emb.embedding[...] = jnp.asarray(weights['wte.weight'])\n",
    "gpt2.pos_emb.embedding[...] = jnp.asarray(weights['wpe.weight'])\n",
    "gpt2.norm.scale[...] = jnp.asarray(weights['ln_f.weight'])\n",
    "gpt2.norm.bias[...] = jnp.asarray(weights['ln_f.bias'])\n",
    "for i, blk in enumerate(gpt2.blks):\n",
    "    for norm, key in ((blk.norm1, f'h.{i}.ln_1'),\n",
    "                      (blk.norm2, f'h.{i}.ln_2')):\n",
    "        norm.scale[...] = jnp.asarray(weights[key + '.weight'])\n",
    "        norm.bias[...] = jnp.asarray(weights[key + '.bias'])\n",
    "    for linear, key in ((blk.attention.W_qkv, f'h.{i}.attn.c_attn'),\n",
    "                        (blk.attention.W_o, f'h.{i}.attn.c_proj'),\n",
    "                        (blk.ffn.W_1, f'h.{i}.mlp.c_fc'),\n",
    "                        (blk.ffn.W_2, f'h.{i}.mlp.c_proj')):\n",
    "        # Conv1D stores (in, out) -- exactly nnx.Linear's kernel layout\n",
    "        linear.kernel[...] = jnp.asarray(weights[key + '.weight'])\n",
    "        linear.bias[...] = jnp.asarray(weights[key + '.bias'])\n",
    "\n",
    "gpt2.eval()\n",
    "n = sum(p.size for p in jax.tree.leaves(nnx.state(gpt2, nnx.Param)))\n",
    "print(f'{n / 1e6:.1f}M parameters')"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "124.4M parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-11 fw=jax -->\n",
    "\n",
    "Loading takes no torch and no transposes here: the safetensors file opens\n",
    "directly into NumPy arrays, and the Conv1D $(\\textrm{in}, \\textrm{out})$\n",
    "layout that PyTorch users must transpose happens to be exactly the layout\n",
    "`nnx.Linear` keeps its kernels in."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-12 fw=all -->\n",
    "\n",
    "Did it work? Weight-loading bugs are notorious for failing silently — a\n",
    "transposed matrix still multiplies. Two checks, one quantitative and one\n",
    "you can read. First, the model should assign natural English a\n",
    "respectable probability:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-loading-gpt-2-3"
   },
   "source": [
    "text = (\"The Time Machine, by H. G. Wells. The Time Traveller was \"\n",
    "        \"expounding a recondite matter to us.\")\n",
    "x = jnp.asarray(enc.encode(text))[None]\n",
    "loss = optax.softmax_cross_entropy_with_integer_labels(\n",
    "    gpt2(x)[0, :-1], x[0, 1:]).mean()\n",
    "print(f'per-token loss {loss:.2f}, perplexity {jnp.exp(loss):.0f}')"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per-token loss 3.89, perplexity 49\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-13 fw=all -->\n",
    "\n",
    "A per-token perplexity around 50 over a 50,257-way vocabulary — against\n",
    "50,257 for uniform guessing — says the plumbing is right. Second,\n",
    "the readable check: greedy decoding of a stock prompt reproduces GPT-2's\n",
    "well-documented continuation, and sampled continuations respond to the\n",
    "temperature and top-$k$ knobs the way the char model could not:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gpt-loading-gpt-2-4"
   },
   "source": [
    "out = gpt2.generate(enc.encode('Alan Turing theorized that computers '\n",
    "                               'would one day become'), 16, top_k=1)\n",
    "print(enc.decode(out))\n",
    "out = gpt2.generate(enc.encode('The secret of a good deep learning '\n",
    "                               'textbook is'), 40, seed=0, temperature=0.8,\n",
    "                    top_k=50)\n",
    "print(enc.decode(out))"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Alan Turing theorized that computers would one day become the most powerful machines on the planet.\n\nThe computer is a computer that\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "The secret of a good deep learning textbook is to make sure that you have the right skills set to get started.\n\nA good deep learning textbook is to make sure you have the right skills set to get started.\n\nI used to\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=gpt-md-14 fw=all -->\n",
    "\n",
    "Take stock of what just happened: a class we wrote in two notebook\n",
    "sections, with the right five flags, runs a model that in 2019 was\n",
    "considered too dangerous to release. The gap between our minute of\n",
    "training and GPT-2 was never architectural — it is five orders of\n",
    "magnitude of data, four of compute, and the engineering to spend them. The\n",
    "rest of this chapter dissects exactly that gap: making generation cheap\n",
    "(the KV cache, next), other ways of wiring the same block\n",
    "(encoders and cross-attention), scaling its FFN sideways\n",
    "(mixture of experts), and the laws that say what all those FLOPs buy\n",
    "(scaling laws) — where the constructor calls for Llama, Qwen, and\n",
    "DeepSeek appear as rows of a table, every one of them an argument list\n",
    "for this class.\n",
    "\n",
    "## Summary\n",
    "\n",
    "A GPT is a token embedding, a stack of causally masked transformer\n",
    "blocks, a final norm, and an output head tied to the embedding; positions\n",
    "enter either as a learned table added at the bottom (GPT-2) or as\n",
    "rotations applied inside every attention head (RoPE, the modern default).\n",
    "Our `GPT` class takes the block's flags plus `pos`, and its causal\n",
    "attention drops into `d2l.TransformerBlock` through the same\n",
    "`attn_factory` hook later sections use for cache-friendly attention.\n",
    "Trained from scratch on a 180 KB novel, the modern configuration reaches\n",
    "its best validation loss within a few hundred steps and then memorizes —\n",
    "the val/train gap, and the apparently memorized stretches in its samples,\n",
    "both read straight off a 30-to-1 parameter-to-data ratio, and argue for\n",
    "the scaling laws that close the chapter. Flipping `pre_norm=False` at a\n",
    "learning rate the pre-norm model tolerates pins training at the unigram\n",
    "entropy: the initialization-time gradient starvation of the previous\n",
    "section, realized as a model that never learns to use context. Sampling\n",
    "is temperature plus truncation, implemented once and reused by the real\n",
    "GPT-2: the released 124M checkpoint loads into our class with a\n",
    "name-mapping dictionary (transposing the Conv1D layout in PyTorch;\n",
    "adopting it unchanged in JAX), its tokenizer reassembles from two pinned\n",
    "data files and the BPE pattern of ch. 8's tokenizer, and the loaded model\n",
    "passes both a perplexity check and a readable one — completing English\n",
    "prompts the way the history books say it should.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Account for every one of GPT-2's 124.4M parameters using the block\n",
    "   census of that section plus the embeddings: 12\n",
    "   blocks of roughly $12d^2$ at $d = 768$ (plus biases), a\n",
    "   $50257 \\times 768$ token table, a $1024 \\times 768$ position table,\n",
    "   and the final norm. What fraction sits in the embeddings, and why does\n",
    "   weight tying make the \"124M parameters\" figure slightly generous as a\n",
    "   measure of the transformer proper?\n",
    "2. Instantiate the GPT-2 configuration with one flag changed to\n",
    "   `norm='rms'` and attempt the weight-loading cell. Which line fails,\n",
    "   and why must it fail rather than load silently? List every tensor in\n",
    "   the checkpoint that would have no destination (and every module\n",
    "   parameter with no source) under the modern flags\n",
    "   `pos='rope', act='swiglu'`.\n",
    "3. Add nucleus (top-$p$) sampling [@Holtzman.Buys.Du.ea.2020] to\n",
    "   `generate`: keep the smallest set of tokens whose probabilities sum to\n",
    "   at least $p$. On the GPT-2 model, fix the seed and compare\n",
    "   continuations under top-$k = 50$ and $p = 0.9$ for a prompt where the\n",
    "   model is confident and one where it is not. Which truncation adapts\n",
    "   its cutoff to the model's uncertainty, and how?\n",
    "4. Show that as $\\tau \\to 0$, sampling with temperature $\\tau$ becomes\n",
    "   greedy decoding, and as $\\tau \\to \\infty$ it converges to the uniform\n",
    "   distribution over the vocabulary. Then measure the empirical entropy\n",
    "   of 200-character samples from the char model at\n",
    "   $\\tau \\in \\{0.5, 1, 2, 4\\}$ and plot it against these two limits.\n",
    "5. Our char model trained with RoPE at context 128 but `max_len` allows\n",
    "   evaluation at 512. Measure its validation loss at contexts 128, 256,\n",
    "   and 512 (build the longer-context dataloaders as in\n",
    "   that section). Does the failure of naive RoPE\n",
    "   extrapolation that section demonstrated reappear inside a full\n",
    "   transformer? Apply position interpolation — scale every angle by\n",
    "   $128/512$ — and report what changes.\n",
    "6. Compare GPT-2 and our char model on the same text on equal footing:\n",
    "   convert each model's mean loss on a held-out Time Machine passage to\n",
    "   *bits per character* (GPT-2's per-token loss must be divided by the\n",
    "   average characters per token of its tokenization). Which model wins,\n",
    "   by how much, and what does each side of the comparison pay for — a\n",
    "   28-symbol character vocabulary against 50,257 subwords, and 40 GB of\n",
    "   WebText against none?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
