{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-1 fw=all -->\n",
    "\n",
    "# Case Study: Making a Transformer Fast\n",
    "\n",
    "This section teaches nothing new. That is the point. Six sections built a\n",
    "method — measure, classify the regime, apply the matching fix, re-measure —\n",
    "and a toolbox: the roofline and the profiler (that section),\n",
    "the hardware that sets the constants (that section), compilation\n",
    "(that section), precision and memory\n",
    "(that section), and data parallelism\n",
    "(that section, that section). Here we run the\n",
    "whole loop end to end on a real model — the GPT of that section — and\n",
    "take it down a *waterfall*: one rung per technique, each measured, each\n",
    "attributed to the section that taught it, re-profiling as the bottleneck\n",
    "moves. The result is a single plot that is the chapter's closing argument:\n",
    "the method is real, and it compounds.\n",
    "\n",
    "*Prerequisites: the entire chapter. The subject is* `d2l.GPT` *and*\n",
    "`d2l.TimeMachine` *from* that section*, reused verbatim.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:23:23.691381Z",
     "iopub.status.busy": "2026-07-21T07:23:23.691011Z",
     "iopub.status.idle": "2026-07-21T07:23:28.339469Z",
     "shell.execute_reply": "2026-07-21T07:23:28.337981Z"
    },
    "id": "fast-transformer-case-study-making-a-transformer-fast"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'0'"
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "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",
    "import os\n",
    "import time\n",
    "\n",
    "# XLA's NCCL buffer registration fails harmlessly on this P2P-less box;\n",
    "# opt out rather than let every collective print the warning (see 13.6).\n",
    "os.environ.setdefault('NCCL_LOCAL_REGISTER', '0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-2 fw=all -->\n",
    "\n",
    "## The Subject\n",
    "\n",
    "We take ch. 11's GPT at a width where the experiment can be read cleanly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-3 fw=jax -->\n",
    "\n",
    "At the `num_hiddens=256` of that section the matmuls are small\n",
    "enough that several rungs would sit inside the measurement noise;\n",
    "widening to `num_hiddens=512` (about 19M parameters, six blocks,\n",
    "char-level *Time Machine*, context 128) gives every real effect room to\n",
    "show. The metric is training throughput in **tokens per second**,\n",
    "measured with the discipline of that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-4 fw=all -->\n",
    "\n",
    "The case study runs in both frameworks, deliberately in parallel: the\n",
    "same model, the same data, the same metric, the same discipline. But the\n",
    "ladders diverge exactly where the frameworks do — what a \"compile\" rung\n",
    "even means, which fix the diagnosis demands next, how data parallelism is\n",
    "launched — and that divergence is itself part of the closing lesson: the\n",
    "*method* transfers; the recipe does not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:23:28.411684Z",
     "iopub.status.busy": "2026-07-21T07:23:28.410974Z",
     "iopub.status.idle": "2026-07-21T07:24:17.870646Z",
     "shell.execute_reply": "2026-07-21T07:24:17.869170Z"
    },
    "id": "fast-transformer-the-subject-1"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "18.9M parameters\n"
     ]
    }
   ],
   "source": [
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char')\n",
    "vocab_size = len(data.vocab)\n",
    "\n",
    "def make_gpt():\n",
    "    return d2l.GPT(vocab_size, num_hiddens=512, num_blks=6,\n",
    "                   rngs=nnx.Rngs(0))\n",
    "\n",
    "model = make_gpt()\n",
    "n_params = sum(p.size for p in jax.tree.leaves(nnx.state(model, nnx.Param)))\n",
    "print(f'{n_params / 1e6:.1f}M parameters')\n",
    "\n",
    "def batches(dm, batch_size=None, seed=0):\n",
    "    \"\"\"Shuffled minibatches, staged on device once (the corpus is 5 MB).\"\"\"\n",
    "    X, Y = jnp.asarray(dm.X[:dm.num_train]), jnp.asarray(dm.Y[:dm.num_train])\n",
    "    B = dm.batch_size if batch_size is None else batch_size\n",
    "    n = (len(X) // B) * B                 # whole batches only: one shape\n",
    "    key = jax.random.key(seed)\n",
    "    while True:\n",
    "        key, sub = jax.random.split(key)\n",
    "        perm = jax.random.permutation(sub, len(X))[:n]\n",
    "        for i in range(0, n, B):\n",
    "            idx = perm[i:i + B]\n",
    "            yield X[idx], Y[idx]\n",
    "stream = batches(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-5 fw=jax -->\n",
    "\n",
    "Before any rung, the measurement itself must be trustworthy, and the\n",
    "traps that corrupted early drafts of the PyTorch tab have JAX\n",
    "counterparts, only sharper. A changed batch shape here does not cost a\n",
    "mere retrace — it triggers a full XLA recompilation, seconds long,\n",
    "*inside* the timing window; the stream above rounds every epoch down to\n",
    "whole batches so a jitted step compiles once\n",
    "(that section's shape specialization). The stream also\n",
    "does something the PyTorch tab does not need to: it stages the corpus —\n",
    "all 5 MB of it — on the device once and shuffles *there*. The book's\n",
    "stock `tf.data` loader is perfectly adequate for every training loop so\n",
    "far, but it costs a millisecond-and-change of host work per batch, and\n",
    "an early draft of this section timed the rungs through it: invisible\n",
    "next to the un-jitted baseline, it silently taxed the fastest rungs\n",
    "below by roughly a quarter — the bottleneck had moved *off the GPU and\n",
    "into the input pipeline*, exactly the failure mode\n",
    "that section warned about. Finally, asynchronous dispatch\n",
    "means a window that does not close on `block_until_ready` measures how\n",
    "fast Python *enqueues* work, not how fast the device finishes it, and\n",
    "the first call of every freshly jitted function pays its compilation —\n",
    "so every warmup below is long enough to absorb one. Warm up until the\n",
    "one-time costs settle; close every timed window with a sync:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:24:17.874046Z",
     "iopub.status.busy": "2026-07-21T07:24:17.873864Z",
     "iopub.status.idle": "2026-07-21T07:24:17.879517Z",
     "shell.execute_reply": "2026-07-21T07:24:17.878164Z"
    },
    "id": "fast-transformer-the-subject-2"
   },
   "outputs": [],
   "source": [
    "def throughput(step_fn, stream, warmup=60, timed=100):\n",
    "    \"\"\"Tokens/s with warmup + sync (see :numref:`sec_perf_model`).\"\"\"\n",
    "    for _ in range(warmup):\n",
    "        loss = step_fn(*next(stream))\n",
    "    jax.block_until_ready(loss)\n",
    "    t0 = time.perf_counter(); n = 0\n",
    "    for _ in range(timed):\n",
    "        X, Y = next(stream); loss = step_fn(X, Y); n += X.size\n",
    "    jax.block_until_ready(loss)\n",
    "    return n / (time.perf_counter() - t0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-6 fw=jax -->\n",
    "\n",
    "One note on the metric: `throughput` times the *whole* pipeline —\n",
    "`next(stream)` includes the shuffle, the batch gather, and the dispatch\n",
    "of both — so these are end-to-end tokens per second, the number a\n",
    "training run actually delivers, not a kernels-only figure. One thing it\n",
    "deliberately does *not* do is feed the same batch twice: a benchmark\n",
    "that re-times one cached batch lets the GPU serve every read from cache\n",
    "and flatters the throughput by double-digit percent (we measured that\n",
    "too). Honest data is fresh data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:24:17.883069Z",
     "iopub.status.busy": "2026-07-21T07:24:17.882827Z",
     "iopub.status.idle": "2026-07-21T07:24:45.345857Z",
     "shell.execute_reply": "2026-07-21T07:24:45.343947Z"
    },
    "id": "fast-transformer-the-subject-3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "fp32 control @512 (clocked up front): 216489 tokens/s, plan 15.1 GiB\n"
     ]
    }
   ],
   "source": [
    "def fp32_step(model, optimizer, X, Y):\n",
    "    def loss_fn(model):\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            model(X).reshape(-1, vocab_size), Y.reshape(-1)).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "_m = make_gpt(); _o = nnx.Optimizer(_m, optax.adamw(1e-3), wrt=nnx.Param)\n",
    "_jit, _big = nnx.jit(fp32_step), batches(data, batch_size=512)\n",
    "_Xb, _Yb = next(_big)\n",
    "ctl512 = throughput(lambda X, Y: _jit(_m, _o, X, Y), _big,\n",
    "                    warmup=20, timed=20)                # fp32 @512, clean pool\n",
    "plan32_gib = (_jit.lower(_m, _o, _Xb, _Yb).compile()\n",
    "              .memory_analysis().temp_size_in_bytes / 2**30)\n",
    "del _m, _o, _jit, _big, _Xb, _Yb; jax.clear_caches()    # release before R0\n",
    "print(f'fp32 control @512 (clocked up front): {ctl512:.0f} tokens/s, '\n",
    "      f'plan {plan32_gib:.1f} GiB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-7 fw=all -->\n",
    "\n",
    "## Rung 0: Baseline, Profiled"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-8 fw=jax -->\n",
    "\n",
    "JAX has no serious eager baseline, and it would be dishonest to pretend\n",
    "otherwise. that section measured why: an un-jitted step\n",
    "dispatches every operation — the forward, the hundreds of intermediate\n",
    "gradients, the optimizer's whole tree of updates — one at a time, and\n",
    "re-traces the function on every call; eager JAX is a development surface,\n",
    "not a training mode. (It is also already tf32-fair: on this card JAX's\n",
    "default matmul precision is tensor-core tf32, the fair baseline\n",
    "that section established.) But a waterfall wants a\n",
    "floor, so we measure the un-jitted step once, briefly, to know exactly\n",
    "what the first rung is worth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:24:45.349418Z",
     "iopub.status.busy": "2026-07-21T07:24:45.349164Z",
     "iopub.status.idle": "2026-07-21T07:25:28.250319Z",
     "shell.execute_reply": "2026-07-21T07:25:28.248642Z"
    },
    "id": "fast-transformer-rung-0-baseline-profiled"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R0 un-jitted: 13287 tokens/s\n"
     ]
    }
   ],
   "source": [
    "optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)\n",
    "\n",
    "def train_step(model, optimizer, X, Y):\n",
    "    def loss_fn(model):\n",
    "        logits = model(X)\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits.reshape(-1, vocab_size), Y.reshape(-1)).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "tput0 = throughput(lambda X, Y: train_step(model, optimizer, X, Y),\n",
    "                   stream, warmup=3, timed=10)\n",
    "print(f'R0 un-jitted: {tput0:.0f} tokens/s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-9 fw=jax -->\n",
    "\n",
    "Read the number for what it is: a price list for a mistake, not a\n",
    "baseline. Every rung below reports its ratio against the *jitted* step,\n",
    "because that is the configuration a JAX practitioner actually starts\n",
    "from — the PyTorch tab's eager baseline is a real place to train; this\n",
    "one is not. Note also what we did *not* do here: profile the un-jitted\n",
    "step. There is nothing to classify in it — the diagnosis is \"Python\n",
    "dispatch, wall to wall\" by construction — so classification is deferred\n",
    "one rung, to the first program worth examining: the compiled one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-10 fw=all -->\n",
    "\n",
    "## Rungs, Each One Measured"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-11 fw=jax -->\n",
    "\n",
    "**Rung 1 — jit the whole step (that section).** The rung is\n",
    "one transformation: `nnx.jit` swallows the loss, the gradients, *and* the\n",
    "optimizer update — 13.3's whole-step lesson — and the training step\n",
    "becomes a single compiled program. This is the framework's table stakes,\n",
    "not an optimization trick, which is why the waterfall's honest ratios\n",
    "start here. The correctness gate first: the jitted forward must produce\n",
    "the same logits as the op-by-op one (they agree to a few times $10^{-4}$;\n",
    "the residue is tf32 arithmetic re-associated by fusion, not a wrong\n",
    "answer). *Same answer first, then faster:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:25:28.253592Z",
     "iopub.status.busy": "2026-07-21T07:25:28.253416Z",
     "iopub.status.idle": "2026-07-21T07:25:45.586055Z",
     "shell.execute_reply": "2026-07-21T07:25:45.584215Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-1"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "first (compiling) step: 7.0 s\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R1 jit: 253478 tokens/s (19x)\n"
     ]
    }
   ],
   "source": [
    "step_jit = nnx.jit(train_step)         # the rung is one transformation\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "jax.block_until_ready(step_jit(model, optimizer, *next(stream)))\n",
    "print(f'first (compiling) step: {time.perf_counter() - t0:.1f} s')\n",
    "fwd = nnx.jit(lambda m, X: m(X))\n",
    "X, Y = next(stream)\n",
    "assert jnp.allclose(fwd(model, X), model(X), atol=1e-2, rtol=1e-3)\n",
    "tput1 = throughput(lambda X, Y: step_jit(model, optimizer, X, Y), stream)\n",
    "print(f'R1 jit: {tput1:.0f} tokens/s ({tput1 / tput0:.0f}x)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-12 fw=jax -->\n",
    "\n",
    "The first call pays several seconds of XLA compile time; steady state\n",
    "repays that at roughly twenty-fold — the two-orders-of-magnitude shape\n",
    "that section measured on its toy step, compressed here\n",
    "because this model's kernels are chunky enough to partly hide the\n",
    "dispatch. Now there is a program worth classifying. The profile the\n",
    "PyTorch tab reads is one tool; the compiler's own accounting is the more\n",
    "natural one in JAX (that section used the same trick\n",
    "for memory): lower the step, compile it, and ask the compiled object\n",
    "what it plans to execute —"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:25:45.589544Z",
     "iopub.status.busy": "2026-07-21T07:25:45.589179Z",
     "iopub.status.idle": "2026-07-21T07:25:45.625312Z",
     "shell.execute_reply": "2026-07-21T07:25:45.623615Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-5"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "XLA counts 765 GFLOP per step; at 32.3 ms per step that is 24 TFLOP/s achieved\n"
     ]
    }
   ],
   "source": [
    "compiled = step_jit.lower(model, optimizer, X, Y).compile()\n",
    "flops = compiled.cost_analysis()['flops']\n",
    "t_step = 64 * 128 / tput1\n",
    "print(f'XLA counts {flops / 1e9:.0f} GFLOP per step; at {1e3 * t_step:.1f} '\n",
    "      f'ms per step that is {flops / t_step / 1e12:.0f} TFLOP/s achieved')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-13 fw=jax -->\n",
    "\n",
    "Put the two numbers against the table: the compiled step\n",
    "achieves about a third of the card's ~83 tf32 TFLOP/s. That\n",
    "classification carries real information, but read it carefully — it is\n",
    "*ambiguous*. Jit already cured R0's op-by-op dispatch disease, and the\n",
    "step is not at the compute roof either; the gap could be arithmetic\n",
    "that is slow (a mixed workload of matmuls, attention softmax,\n",
    "normalizations, and rope trigonometry idles the tensor cores part of\n",
    "the time), or it could be a fixed per-step cost that even a fused\n",
    "program still pays. The cheap experiment that separates the suspects is\n",
    "the format ladder's next rung: bf16 doubles the tensor-core roof and\n",
    "halves every activation byte — if arithmetic is the wall, the effect\n",
    "must show. So, precision next — as a *measurement* as much as an\n",
    "optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:25:45.628949Z",
     "iopub.status.busy": "2026-07-21T07:25:45.628580Z",
     "iopub.status.idle": "2026-07-21T07:25:57.282579Z",
     "shell.execute_reply": "2026-07-21T07:25:57.280541Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\"bf16\", naive cast: 249436 tokens/s (0.98x)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "the receipt: float32\n"
     ]
    }
   ],
   "source": [
    "def bf16_arrays(model):\n",
    "    \"\"\"Cast every floating array of a copy -- surely bf16 now?\"\"\"\n",
    "    graphdef, state = nnx.split(model)\n",
    "    state = jax.tree.map(lambda x: x.astype(jnp.bfloat16)\n",
    "                         if jnp.issubdtype(x.dtype, jnp.floating) else x,\n",
    "                         state)\n",
    "    return nnx.merge(graphdef, state)\n",
    "\n",
    "@nnx.jit\n",
    "def step_naive(model, optimizer, X, Y):\n",
    "    def loss_fn(model):\n",
    "        logits = bf16_arrays(model)(X)\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits.reshape(-1, vocab_size).astype(jnp.float32),\n",
    "            Y.reshape(-1)).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "naive = throughput(lambda X, Y: step_naive(model, optimizer, X, Y), stream)\n",
    "print(f'\"bf16\", naive cast: {naive:.0f} tokens/s ({naive / tput1:.2f}x)')\n",
    "print('the receipt:', bf16_arrays(model).token_emb(X).dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-14 fw=jax -->\n",
    "\n",
    "Flat — and the second print is the receipt. The embedding layer still\n",
    "emits fp32: flax modules record a *compute dtype* at construction\n",
    "(`Embed` remembers its `param_dtype`), and their `__call__` promotes\n",
    "inputs and parameters to it — so the bf16 table we so carefully installed\n",
    "was promoted straight back to fp32 at the first layer, and everything\n",
    "downstream followed. Our \"bf16 rung\" measured a wardrobe change, not a\n",
    "precision change. This is the JAX tab's version of the PyTorch tab's\n",
    "corrupted-measurement traps, and it is nastier for being *silent*: no\n",
    "error, no warning, a plausible number. The measurement — plus one dtype\n",
    "print — caught it. Threading bf16 for real takes two more moves. First,\n",
    "set the compute dtype the modules remember, not just the arrays. Second,\n",
    "one leak is in the model's own arithmetic: rope's fp32 trigonometry\n",
    "promotes the query and key back to fp32 on their way into attention —\n",
    "and the attention kernel, rightly, refuses mixed-precision inputs\n",
    "(loudly, this time) — so a two-line subclass casts them back. The fp32\n",
    "masters stay in the model and the optimizer; the gradients arrive fp32 on\n",
    "their own, because the transpose of a cast is a cast back:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:25:57.286762Z",
     "iopub.status.busy": "2026-07-21T07:25:57.286402Z",
     "iopub.status.idle": "2026-07-21T07:26:15.204046Z",
     "shell.execute_reply": "2026-07-21T07:26:15.202327Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-9"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "the receipt, again: bfloat16\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "planned temp: bf16 1149 MiB (fp32: 1989 MiB)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R2 +bf16, threaded: 242897 tokens/s (0.96x)\n"
     ]
    }
   ],
   "source": [
    "def to_bf16(model):\n",
    "    \"\"\"A bf16 compute copy: cast every floating array AND the compute\n",
    "    dtype each flax module remembers from construction.\"\"\"\n",
    "    copy = bf16_arrays(model)\n",
    "    for _, m in nnx.iter_graph(copy):\n",
    "        if isinstance(m, nnx.Module) and hasattr(m, 'dtype'):\n",
    "            m.dtype = jnp.bfloat16\n",
    "    return copy\n",
    "\n",
    "class Bf16GPT(d2l.GPT):\n",
    "    \"\"\"ch. 11's GPT with one dtype leak patched: rope's fp32 trigonometry\n",
    "    promotes q and k back to fp32 (and the attention kernel rightly\n",
    "    refuses mixed inputs), so cast them back after rope.\"\"\"\n",
    "    class CausalAttention(d2l.GPT.CausalAttention):\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",
    "            q, k = q.astype(v.dtype), k.astype(v.dtype)\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 make_bf16_gpt():\n",
    "    return Bf16GPT(vocab_size, num_hiddens=512, num_blks=6,\n",
    "                   rngs=nnx.Rngs(0))\n",
    "\n",
    "model16 = make_bf16_gpt()\n",
    "opt16 = nnx.Optimizer(model16, optax.adamw(1e-3), wrt=nnx.Param)\n",
    "\n",
    "@nnx.jit\n",
    "def step_bf16(model, optimizer, X, Y):\n",
    "    def loss_fn(model):\n",
    "        logits = to_bf16(model)(X)         # bf16 compute copy of the step;\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits.reshape(-1, vocab_size).astype(jnp.float32),  # fp32 loss\n",
    "            Y.reshape(-1)).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)   # grads arrive fp32\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "print('the receipt, again:', to_bf16(model16).token_emb(X).dtype)\n",
    "plan16 = step_bf16.lower(model16, opt16, X, Y).compile().memory_analysis()\n",
    "print(f'planned temp: bf16 {plan16.temp_size_in_bytes / 2**20:.0f} MiB '\n",
    "      f'(fp32: {compiled.memory_analysis().temp_size_in_bytes / 2**20:.0f} '\n",
    "      f'MiB)')\n",
    "tput2 = throughput(lambda X, Y: step_bf16(model16, opt16, X, Y), stream)\n",
    "assert all(p.dtype == jnp.float32                      # fp32 masters intact\n",
    "           for p in jax.tree.leaves(nnx.state(model16, nnx.Param)))\n",
    "print(f'R2 +bf16, threaded: {tput2:.0f} tokens/s ({tput2 / tput1:.2f}x)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-15 fw=jax -->\n",
    "\n",
    "The second surprise, and the better lesson: bf16 is now provably *real* —\n",
    "the receipt prints `bfloat16`, the compiler's planned temporaries drop by\n",
    "about two-fifths — and the throughput *still* barely moves, nowhere near\n",
    "the factor the format ladder promised. Diagnose before despairing. The\n",
    "device-side work genuinely shrank — the plan says so; what did not\n",
    "shrink is everything around it — the per-call Python that walks the\n",
    "module graph, the\n",
    "stream's dispatches, the launch path — a fixed per-step toll that the\n",
    "fp32 step was already brushing against. At batch 64 this step is\n",
    "*overhead-bound* in exactly that section's sense, and making\n",
    "the arithmetic faster cannot move a wall made of dispatch. (The tell,\n",
    "if you re-run this notebook: the ratio above wobbles from run to run —\n",
    "deterministic device work does not do that; host state does.) The\n",
    "matching fix is not more precision. It is *amortization*: make every\n",
    "step carry more tokens, so the fixed toll is divided by a bigger\n",
    "number. That is the next rung — where the bf16 investment, correct all\n",
    "along but bought where the wall wasn't, will finally cash."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:26:15.208878Z",
     "iopub.status.busy": "2026-07-21T07:26:15.208496Z",
     "iopub.status.idle": "2026-07-21T07:26:38.738264Z",
     "shell.execute_reply": "2026-07-21T07:26:38.736140Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R3 +batch-up (512): 434987 tokens/s (1.79x), 44 TFLOP/s achieved\nfp32 control at 512: 216489 tokens/s -> bf16 pays 2.01x here\nplanned temp: bf16 8.9 GiB (fp32 control: 15.1 GiB)\n"
     ]
    }
   ],
   "source": [
    "big_stream = batches(data, batch_size=512)   # same corpus, bigger gathers\n",
    "Xb, Yb = next(big_stream)\n",
    "\n",
    "# The bf16-512 step (~9 GiB) fits alongside the pool the ladder has already\n",
    "# grown; the fp32-512 control (~15 GiB) did not, which is why it was clocked\n",
    "# up front (`ctl512`, `plan32_gib`). Here we run bf16-512 and read its plan.\n",
    "tput3 = throughput(lambda X, Y: step_bf16(model16, opt16, X, Y),\n",
    "                   big_stream, warmup=30, timed=30)\n",
    "c16 = step_bf16.lower(model16, opt16, Xb, Yb).compile()\n",
    "plan16_gib = c16.memory_analysis().temp_size_in_bytes / 2**30\n",
    "tf512 = c16.cost_analysis()['flops'] / (512 * 128 / tput3) / 1e12\n",
    "print(f'R3 +batch-up (512): {tput3:.0f} tokens/s ({tput3 / tput2:.2f}x), '\n",
    "      f'{tf512:.0f} TFLOP/s achieved')\n",
    "print(f'fp32 control at 512: {ctl512:.0f} tokens/s '\n",
    "      f'-> bf16 pays {tput3 / ctl512:.2f}x here')\n",
    "print(f'planned temp: bf16 {plan16_gib:.1f} GiB '\n",
    "      f'(fp32 control: {plan32_gib:.1f} GiB)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-16 fw=jax -->\n",
    "\n",
    "Read the three printed lines together, because they close R2's case and\n",
    "then break a lazy assumption. The bf16 step at batch 512 delivers most\n",
    "of a factor of two over the stalled batch-64 rung, and the achieved\n",
    "TFLOP/s roughly doubles — the tensor cores are finally earning their\n",
    "silicon. But look at the fp32 control: at batch 512 it is *slower than\n",
    "fp32 at batch 64*. The batch rung, taken alone, is a *negative* rung\n",
    "for the fp32 program — its compute rate does not rise, holding around\n",
    "twenty TFLOP/s, no better than the batch-64 step — and the plan says\n",
    "why: roughly 15 GiB of temporaries,\n",
    "dominated by materialized fp32 attention buffers, and a plan that large\n",
    "is not just occupancy, it is *traffic* — the fp32 step crosses into the\n",
    "bandwidth regime just as the host toll recedes. So neither rung pays by\n",
    "itself here: bf16 at batch 64 hit the dispatch wall; batch-up at fp32\n",
    "hits the memory wall; *together* they pay, because bf16 halves the\n",
    "bytes exactly where the bigger batch needs it, and the bigger batch\n",
    "amortizes the dispatch exactly where bf16 needs it. Techniques\n",
    "interact. That is why the waterfall is cumulative, why every rung is\n",
    "re-measured on top of the last — and why a recipe of independent \"tips\"\n",
    "would have called both of these rungs wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:26:38.742364Z",
     "iopub.status.busy": "2026-07-21T07:26:38.742002Z",
     "iopub.status.idle": "2026-07-21T07:27:12.589129Z",
     "shell.execute_reply": "2026-07-21T07:27:12.586866Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-4"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R4 +checkpoint: 377077 tokens/s (0.87x), planned temp 1.8 GiB\n"
     ]
    }
   ],
   "source": [
    "def ckpt_forward(model, X):\n",
    "    \"\"\"ch. 11's GPT.__call__ with each block under jax.checkpoint.\n",
    "    Reuses the exact submodules; only the block loop differs.\"\"\"\n",
    "    H = model.token_emb(X)\n",
    "    if model.pos == 'learned':                # (default 'rope' adds no emb)\n",
    "        H = H + model.pos_emb(jnp.arange(X.shape[1]))\n",
    "    for blk in model.blks:\n",
    "        graphdef, state = nnx.split(blk)\n",
    "        H = jax.checkpoint(\n",
    "            lambda state, H: nnx.merge(graphdef, state)(H))(state, H)\n",
    "    return model.token_emb.attend(model.norm(H))\n",
    "\n",
    "half = to_bf16(model16)\n",
    "assert jnp.allclose(ckpt_forward(half, Xb), half(Xb), atol=1e-6)  # a mirror\n",
    "\n",
    "@nnx.jit\n",
    "def step_ckpt(model, optimizer, X, Y):\n",
    "    def loss_fn(model):\n",
    "        logits = ckpt_forward(to_bf16(model), X)\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            logits.reshape(-1, vocab_size).astype(jnp.float32),\n",
    "            Y.reshape(-1)).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "plan_ck = step_ckpt.lower(model16, opt16, Xb, Yb).compile().memory_analysis()\n",
    "tput4 = throughput(lambda X, Y: step_ckpt(model16, opt16, X, Y),\n",
    "                   big_stream, warmup=30, timed=30)\n",
    "print(f'R4 +checkpoint: {tput4:.0f} tokens/s ({tput4 / tput3:.2f}x), '\n",
    "      f'planned temp {plan_ck.temp_size_in_bytes / 2**30:.1f} GiB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-17 fw=jax -->\n",
    "\n",
    "**Rung 5 — data parallel across 2–4 GPUs\n",
    "(that section), predicted first.** Before measuring, we\n",
    "*predict*, using the accounting of that section. This GPT has\n",
    "about 19M parameters, so roughly 76 MB of gradients must be summed across\n",
    "devices every step — and that section measured the very\n",
    "collective JAX will run, `psum` through NCCL on this host-staged box, at\n",
    "about 4.5 GB/s per device in its bytes-per-device convention. That prices\n",
    "the allreduce at the same tens of milliseconds as the *entire* compute\n",
    "step: a transformer's parameters are proportional to its compute, so\n",
    "unlike a convolutional ResNet it offers little extra compute to hide\n",
    "communication behind. Summing the two terms, the equation predicts\n",
    "something blunt — at $k=2$, barely one single-GPU throughput: the second\n",
    "GPU buys nearly nothing; under half of linear at $k=4$. The mechanism,\n",
    "though, is 13.6's declarative recipe at its cleanest: build a mesh,\n",
    "replicate the state, shard the batch, and feed the **unchanged** jitted\n",
    "step — no launcher, no sidecar script, no rank files; one process, three\n",
    "annotations. (One version note: the mesh names its axis `Auto` — under\n",
    "this JAX version's explicit-sharding types the embedding's gather cannot\n",
    "infer an output layout, and `Auto` hands the whole layout problem back to\n",
    "GSPMD, which is exactly 13.6's semantics.) Per-device batch 64 — weak\n",
    "scaling, the convention of that section; fresh shuffled\n",
    "batches each step, so this is a throughput measurement, not a convergence\n",
    "claim:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:27:12.593393Z",
     "iopub.status.busy": "2026-07-21T07:27:12.593012Z",
     "iopub.status.idle": "2026-07-21T07:27:12.604854Z",
     "shell.execute_reply": "2026-07-21T07:27:12.603122Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-7"
   },
   "outputs": [],
   "source": [
    "def dp_setup(k, make_model):\n",
    "    \"\"\"13.6's recipe: replicate the state, shard the batch, jit unchanged.\"\"\"\n",
    "    mesh = jax.make_mesh((k,), ('data',),\n",
    "                         axis_types=(jax.sharding.AxisType.Auto,))\n",
    "    P = jax.sharding.PartitionSpec\n",
    "    repl = jax.sharding.NamedSharding(mesh, P())\n",
    "    shard = jax.sharding.NamedSharding(mesh, P('data'))\n",
    "    mdl = make_model()\n",
    "    opt = nnx.Optimizer(mdl, optax.adamw(1e-3), wrt=nnx.Param)\n",
    "    nnx.update(mdl, jax.device_put(nnx.state(mdl), repl))\n",
    "    nnx.update(opt, jax.device_put(nnx.state(opt), repl))\n",
    "    return mdl, opt, shard\n",
    "\n",
    "def dp_stream(src, k, shard):\n",
    "    \"\"\"Global batches: k per-device batches, split along the batch axis.\"\"\"\n",
    "    while True:\n",
    "        XY = [next(src) for _ in range(k)]\n",
    "        yield (jax.device_put(jnp.concatenate([x for x, _ in XY]), shard),\n",
    "               jax.device_put(jnp.concatenate([y for _, y in XY]), shard))\n",
    "\n",
    "def dp_tokens(k, step, make_model, src, warmup=20, timed=40):\n",
    "    mdl, opt, shard = dp_setup(k, make_model)\n",
    "    return throughput(lambda X, Y: step(mdl, opt, X, Y),\n",
    "                      dp_stream(src, k, shard), warmup=warmup, timed=timed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:27:12.608604Z",
     "iopub.status.busy": "2026-07-21T07:27:12.608262Z",
     "iopub.status.idle": "2026-07-21T07:27:37.347961Z",
     "shell.execute_reply": "2026-07-21T07:27:37.345923Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-8"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per step: t_comm ~ 34 ms vs t_compute ~ 32 ms\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R5 k=2: no-overlap floor 249k tokens/s, measured 199k (0.79x of one GPU)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R5 k=4: no-overlap floor 497k tokens/s, measured 297k (1.17x of one GPU)\n"
     ]
    }
   ],
   "source": [
    "n_bytes = 4 * n_params\n",
    "beta = 4.5e9   # psum, effective bytes/device/s on this box (13.5)\n",
    "t_comm, t_cmp = 2 * n_bytes / beta, 64 * 128 / tput1\n",
    "print(f'per step: t_comm ~ {1e3 * t_comm:.0f} ms vs '\n",
    "      f't_compute ~ {1e3 * t_cmp:.0f} ms')\n",
    "for k in (k for k in (2, 4) if k <= jax.local_device_count()):\n",
    "    floor = k * 64 * 128 / (t_cmp + t_comm)\n",
    "    meas = dp_tokens(k, step_jit, make_gpt, stream)\n",
    "    print(f'R5 k={k}: no-overlap floor {floor / 1e3:.0f}k tokens/s, '\n",
    "          f'measured {meas / 1e3:.0f}k ({meas / tput1:.2f}x of one GPU)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-18 fw=jax -->\n",
    "\n",
    "The measurement lands where the model said it would — at the floor,\n",
    "nowhere near the linear ceiling — and in our runs a shade *under* it. At\n",
    "$k=2$ the second GPU buys nothing at all; the total can even dip below\n",
    "one GPU's throughput, because the floor is a model and models omit\n",
    "terms: the global batch stages through one device before scattering,\n",
    "and the collective must still be scheduled. At $k=4$ the shortfall\n",
    "widens — the floor was priced with the $\\beta$ measured at two GPUs, and\n",
    "four ranks hammering the same host path do not sustain that constant\n",
    "(the PyTorch tab's sidecar sees the same sag). Nor does overlap rescue\n",
    "it: XLA emits overlap-capable collectives, but overlap needs compute to\n",
    "hide behind, and this rung has none to spare. Communication-hungry is\n",
    "the diagnosis, and it was written down *before* anything ran. On an\n",
    "NVLink box the same accounting predicts near-linear scaling — the\n",
    "constant changes, the method does not. And here the one-process,\n",
    "annotate-the-layout model pays a dividend the sidecar cannot: stacking\n",
    "data parallelism on the *full* ladder — bf16, batch 512 — is one more\n",
    "loop over the same helpers, so we run the measurement the PyTorch tab\n",
    "leaves as an exercise. Eight-fold more compute per step, the same 76 MB\n",
    "gradient bill: the equation promises much better efficiency —"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:27:37.352285Z",
     "iopub.status.busy": "2026-07-21T07:27:37.351928Z",
     "iopub.status.idle": "2026-07-21T07:28:17.623254Z",
     "shell.execute_reply": "2026-07-21T07:28:17.620980Z"
    },
    "id": "fast-transformer-rungs-each-one-measured-10"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "fast config, k=2: 513k tokens/s (59% weak-scaling efficiency)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "fast config, k=4: 995k tokens/s (57% weak-scaling efficiency)\n"
     ]
    }
   ],
   "source": [
    "for k in (k for k in (2, 4) if k <= jax.local_device_count()):\n",
    "    meas = dp_tokens(k, step_bf16, make_bf16_gpt, big_stream,\n",
    "                     warmup=15, timed=25)\n",
    "    print(f'fast config, k={k}: {meas / 1e3:.0f}k tokens/s '\n",
    "          f'({meas / (k * tput3):.0%} weak-scaling efficiency)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-19 fw=jax -->\n",
    "\n",
    "— and delivers it: the fast configuration scales at markedly higher\n",
    "efficiency than the batch-64 rung, for the reason the cost model states:\n",
    "$t_{\\text{compute}}$ grew eight-fold while $t_{\\text{comm}}$ stayed\n",
    "fixed. Still well short of linear — the fabric is what it is — but the\n",
    "*shape* of the improvement was predicted by the same two-term sum that\n",
    "predicted the failure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-20 fw=all -->\n",
    "\n",
    "## The Waterfall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-21 fw=jax -->\n",
    "\n",
    "Collect the single-GPU rungs into one plot — the chapter's closing image.\n",
    "The waterfall is *cumulative*: each bar inherits every choice to its\n",
    "left, so a bar's ratio to its neighbor isolates one technique. The\n",
    "labels are ratios to the *jitted* bar — the honest baseline here — with\n",
    "the un-jitted bar left in the plot as a reminder of what the dishonest\n",
    "one would look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:28:17.627837Z",
     "iopub.status.busy": "2026-07-21T07:28:17.627446Z",
     "iopub.status.idle": "2026-07-21T07:28:17.766676Z",
     "shell.execute_reply": "2026-07-21T07:28:17.764908Z"
    },
    "id": "fast-transformer-the-waterfall"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cumulative, R1 -> R3: 1.72x (and R0 -> R3: 33x)\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 30656 characters]"
     },
     "metadata": {}
    }
   ],
   "source": [
    "rungs = ['R0\\nun-jitted', 'R1\\njit', 'R2\\n+bf16', 'R3\\n+batch', 'R4\\n+ckpt']\n",
    "tputs = [tput0, tput1, tput2, tput3, tput4]\n",
    "print(f'cumulative, R1 -> R3: {tputs[3] / tputs[1]:.2f}x '\n",
    "      f'(and R0 -> R3: {tputs[3] / tputs[0]:.0f}x)')\n",
    "d2l.plt.figure(figsize=(6, 3.5))\n",
    "bars = d2l.plt.bar(rungs, [t / 1e3 for t in tputs],\n",
    "                   color=['#7f7f7f', '#1f77b4', '#1f77b4', '#2ca02c', '#d62728'])\n",
    "d2l.plt.ylabel('throughput (k tokens/s)')\n",
    "d2l.plt.title('Making a Transformer fast: cumulative rungs')\n",
    "for b, t in zip(bars, tputs):\n",
    "    d2l.plt.text(b.get_x() + b.get_width() / 2, t / 1e3,\n",
    "                 f'{t / tputs[1]:.2f}x', ha='center', va='bottom')\n",
    "d2l.plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-22 fw=jax -->\n",
    "\n",
    "Read left to right, and notice how different this waterfall's shape is\n",
    "from the PyTorch tab's — same model, same data, same card. Jit is the\n",
    "cliff: around twenty-fold, the rung that is not optional. Then the flat\n",
    "bar that taught the most: bf16, threaded until it was provably real, and\n",
    "still worth almost nothing at batch 64 — not because the tensor cores\n",
    "idled by fault of the format, but because the step's wall had moved to\n",
    "the fixed per-step dispatch work, where precision cannot reach. The\n",
    "batch rung turns that diagnosis into the ladder's second real win, most\n",
    "of a factor of two — an *interaction*, the fp32 control showed, since\n",
    "either rung alone loses on this stack. And\n",
    "checkpointing, the red bar, *cost* about a tenth of the throughput\n",
    "while cutting the planned temporaries several-fold — a negative rung for\n",
    "*speed*, because memory was never the binding constraint here. The\n",
    "cumulative figure the cell prints runs from the honest baseline: jitted\n",
    "to batch-up lands at about 1.7× (the print keeps the un-jitted\n",
    "strawman's thirty-odd-fold for the record; the plot's labels do not).\n",
    "Two tabs,\n",
    "two different bottleneck sequences, two different winning moves — and\n",
    "the same loop found both. That is the whole lesson: **you cannot copy a\n",
    "recipe — not even from the other tab of the same book — you measure,\n",
    "classify, fix, and re-measure.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-23 fw=all -->\n",
    "\n",
    "One closing check, because a waterfall of throughputs proves nothing\n",
    "about learning: train the fast configuration — compiled, bf16, batch\n",
    "512 — for a few hundred steps from a fresh initialization and watch the\n",
    "loss, smoothed over twenty steps, actually fall. (On this tiny corpus a\n",
    "few hundred big-batch steps revisit the text over a hundred times, so the\n",
    "loss falls very far — the assertion, not the final value, is the point.)\n",
    "Speed that breaks the model is not speed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-21T07:28:17.770408Z",
     "iopub.status.busy": "2026-07-21T07:28:17.770157Z",
     "iopub.status.idle": "2026-07-21T07:29:14.739547Z",
     "shell.execute_reply": "2026-07-21T07:29:14.737686Z"
    },
    "id": "fast-transformer-the-waterfall-1"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "smoothed loss: first 20 steps 2.84 -> last 20 steps 0.06\n"
     ]
    }
   ],
   "source": [
    "learner = make_bf16_gpt()\n",
    "opt_l = nnx.Optimizer(learner, optax.adamw(1e-3), wrt=nnx.Param)\n",
    "losses = []\n",
    "for _ in range(300):\n",
    "    X, Y = next(big_stream)\n",
    "    losses.append(step_bf16(learner, opt_l, X, Y))  # no host read (13.1)\n",
    "losses = jnp.stack(losses)\n",
    "first, last = losses[:20].mean(), losses[-20:].mean()\n",
    "print(f'smoothed loss: first 20 steps {first:.2f} -> last 20 steps {last:.2f}')\n",
    "assert last < first, 'speed that breaks the model is not speed'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-24 fw=all -->\n",
    "\n",
    "## The Lore, and the Ladder Beyond\n",
    "\n",
    "The competitive edge of this method is a spectator sport. The\n",
    "*modded-nanoGPT* speedrun tracks the wall-clock time to train a small GPT\n",
    "to a target loss, and its record holders read like this chapter's table of\n",
    "contents stacked to the ceiling: compiled and fused kernels\n",
    "(that section), block-sparse and FlashAttention variants\n",
    "(that section), a better optimizer\n",
    "(that section), aggressive low precision down to fp8 — each record a\n",
    "new rung on a waterfall like ours. fp8 tensor cores exist on our Ada GPUs,\n",
    "but the recipe for *training* in fp8 (its own loss-scaling and master-copy\n",
    "discipline) is deferred to the Language Models part, along with the\n",
    "multi-node parallelism where these techniques are stacked at frontier\n",
    "scale.\n",
    "\n",
    "That is where the chapter ends, seven sections deep into a single idea.\n",
    "The machine has two budgets, arithmetic and bandwidth; a program lives in\n",
    "one of three regimes; and the way to make anything fast is always the same\n",
    "loop — **measure, classify, fix, re-measure** — now demonstrated on a real\n",
    "model, one rung at a time.\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-25 fw=jax -->\n",
    "\n",
    "* The chapter's method applied whole: measure the baseline, classify\n",
    "  with the compiler's own accounting, one technique per rung,\n",
    "  re-measure. On ch. 11's GPT (width 512, ~19M params), jit is worth\n",
    "  around twenty-fold over the un-jitted strawman; bf16 at batch 64 then\n",
    "  measures almost *nothing* (the dispatch wall), the fp32 control shows\n",
    "  batch-up alone is *negative* (the memory wall) — and together the two\n",
    "  rungs pay most of a factor of two. Rungs interact; ladders are\n",
    "  measured cumulatively. The waterfall prints the figure from the\n",
    "  jitted baseline.\n",
    "* Precision in JAX is explicit, and *layered*: casting the arrays is not\n",
    "  enough — flax modules remember a compute dtype from construction, and\n",
    "  the model's own fp32 arithmetic (rope's trigonometry) promotes values\n",
    "  straight back, which the attention kernel refuses loudly. The naive\n",
    "  cast measured exactly like fp32; one dtype print exposed it. The\n",
    "  masters, the gradients, and the Adam state stay fp32 by construction.\n",
    "* Measurement is a technique too, and it can fail: the stock `tf.data`\n",
    "  loader silently taxed the fastest rungs (fixed by staging the 5 MB\n",
    "  corpus on device), and re-timing one cached batch flatters throughput\n",
    "  by double-digit percent. Both had to be fixed before the waterfall\n",
    "  meant anything.\n",
    "* Rungs that do not help teach the most: bf16 at batch 64 is a *null*\n",
    "  rung here — the wall was per-step dispatch, which precision cannot\n",
    "  touch — and checkpointing costs time to save memory the compiler's\n",
    "  plan said we did not need. Knowing when a technique does not apply is\n",
    "  the method too.\n",
    "* Data parallelism is predicted before it is measured: with $\\beta$ from\n",
    "  13.5's measured `psum`, the equation prices the $k=2$ gain at\n",
    "  roughly nothing — the measurement lands at or below that floor — and\n",
    "  $k=4$ falls well short of the floor as the host-staged fabric sags.\n",
    "  Declarative sharding then makes rerunning the *fast* configuration a\n",
    "  loop rather than a launcher, and its recovered efficiency is the cost\n",
    "  model's grown compute term, verified."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-26 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Reproduce the negative bf16 rung: run the R0→R1→R2 sequence at\n",
    "   `num_hiddens=256` and confirm bf16 is *slower* than compile-alone. Then\n",
    "   profile both and explain, in terms of the figure, why the\n",
    "   width-256 matmuls do not benefit from the tensor cores.\n",
    "1. Add a rung of your own — `channels_last` memory format, pinned-memory\n",
    "   dataloading, or `sdpa_kernel` backend pinning — measure it, and slot it\n",
    "   into the waterfall. Which regime does it attack, and does it clear the\n",
    "   noise floor?\n",
    "1. Apply the whole ladder to a *different* model: ch. 12's Mamba capstone\n",
    "   (that section) or ch. 11's ViT. Which rungs change sign, and\n",
    "   why? (Hint: a model's arithmetic intensity decides which fixes pay.)\n",
    "1. Take a rung that helped here — say, batch-up — and construct a model\n",
    "   (deeper, or with a much larger vocabulary) where it *hurts*. Explain\n",
    "   the reversal with the memory anatomy of that section.\n",
    "1. Rung 5 parallelized the *eager* model. Stack DDP on the full ladder\n",
    "   instead — compiled, bf16, per-rank batch 512 — and measure tokens/s at\n",
    "   $k = 2$ and $k = 4$. Efficiency comes out markedly higher than the\n",
    "   eager rung's. Explain why with that section's cost model:\n",
    "   which term grew, and which stayed fixed?\n",
    "1. Our R1 measurement was once corrupted by a retrace inside the timing\n",
    "   window. Remove the ragged-batch filter from `batches`, re-run R1, and\n",
    "   find the retrace two ways: in the throughput number, and via\n",
    "   `torch._dynamo` logging. How many steps per epoch does the offending\n",
    "   batch appear in, and why does a longer timed window dilute but not\n",
    "   remove the damage?\n",
    "1. (JAX) The JAX tab never touched the attention kernel: `d2l.GPT`'s\n",
    "   default lowering materializes the $T \\times T$ score matrix, and the\n",
    "   fused-kernel fix of that section\n",
    "   (`implementation='cudnn'`) is one argument away — but it accepts only\n",
    "   half-precision inputs, so it needs R2's threading. Subclass\n",
    "   `CausalAttention` to name the kernel, verify parity against the\n",
    "   default lowering, and re-measure the ladder at context 128 and again\n",
    "   at context 1024, reading both tokens/s and the compiler's planned temp\n",
    "   bytes. Where does the fused kernel start to pay, and why does context\n",
    "   length matter more than batch size?\n",
    "1. (JAX) Rung 5 sharded only the batch. Following 13.6's punchline,\n",
    "   change the `PartitionSpec` so the *parameters* are sharded across the\n",
    "   mesh as well (the FSDP pattern), leaving the step function untouched.\n",
    "   Read the compiled HLO for the collectives GSPMD inserts in place of\n",
    "   the allreduce, and measure tokens/s at $k = 2$ and $k = 4$. At 19M\n",
    "   parameters this buys no speed — which term of\n",
    "   that section's memory anatomy would have to grow\n",
    "   before it pays?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l (JAX)",
   "language": "python",
   "name": "d2l-jax"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
