{
 "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",
   "metadata": {
    "id": "fast-transformer-case-study-making-a-transformer-fast"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import torch as d2l\n",
    "import json\n",
    "import os\n",
    "import pathlib\n",
    "import subprocess\n",
    "import sys\n",
    "import time\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.checkpoint as ckpt\n",
    "\n",
    "torch.set_float32_matmul_precision('high')\n",
    "# The profiler's CUDA instrumentation must not outlive its cell — without\n",
    "# this, every timing after the first profile pays a per-launch tax.\n",
    "os.environ['TEARDOWN_CUPTI'] = '1'"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "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=pytorch -->\n",
    "\n",
    "At the `num_hiddens=256` of that section, the model's matmuls are\n",
    "small enough that some techniques do not clear the measurement noise —\n",
    "worse, bf16 actively *hurts*, because width-256 matmuls are too small to\n",
    "feed the tensor cores and the casting overhead dominates (an instructive\n",
    "failure, pursued in the exercises). Widening to `num_hiddens=512` (about\n",
    "19M parameters, six blocks, char-level *Time Machine*, context 128) puts\n",
    "every rung comfortably above the noise floor. The metric is training\n",
    "throughput in **tokens per second**, measured with the discipline of\n",
    "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",
   "metadata": {
    "id": "fast-transformer-the-subject-1"
   },
   "source": [
    "device = d2l.try_gpu()\n",
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char')\n",
    "vocab_size = len(data.vocab)\n",
    "\n",
    "def make_gpt():\n",
    "    torch.manual_seed(0)\n",
    "    return d2l.GPT(vocab_size, num_hiddens=512, num_blks=6).to(device)\n",
    "\n",
    "model = make_gpt()\n",
    "print(f'{sum(p.numel() for p in model.parameters()) / 1e6:.1f}M parameters')\n",
    "\n",
    "def batches(dm):\n",
    "    while True:\n",
    "        for X, Y in dm.train_dataloader():\n",
    "            if X.shape[0] == dm.batch_size:   # constant shape: no retrace\n",
    "                yield X.to(device), Y.to(device)\n",
    "stream = batches(data)"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "18.9M parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-5 fw=pytorch -->\n",
    "\n",
    "Before any rung, the measurement itself must be trustworthy, and two traps\n",
    "had to be defused here — both worth naming, because each corrupted an\n",
    "early draft of this section by tens of percent. First, the dataset's\n",
    "ragged final batch (10,000 windows do not divide by 64) hands\n",
    "`torch.compile` a new shape once per epoch, and the retrace it triggers\n",
    "(that section) lands *inside* a timing window; the stream\n",
    "above drops the ragged tail so the compiler sees one shape. Second, the\n",
    "profiler's instrumentation is not free and, by default, does not fully\n",
    "detach when its `with` block ends — every measurement taken afterwards in\n",
    "the same process inherits a per-launch tax that hits small-batch\n",
    "configurations hardest; the `TEARDOWN_CUPTI` line in the imports cell\n",
    "makes profiling stop costing once it stops running. Both are the\n",
    "chapter's opening lesson in miniature: *the first thing to profile is\n",
    "your own experiment.* Beyond that, every measurement warms up long enough\n",
    "for one-time costs — allocator growth, autotuning, compilation caches,\n",
    "and whatever clock state the driver is in — to settle, and timed windows\n",
    "close with a device sync:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-the-subject-2"
   },
   "source": [
    "def throughput(step_fn, warmup=60, timed=100):\n",
    "    \"\"\"Tokens/s with warmup + device sync (see :numref:`sec_perf_model`).\"\"\"\n",
    "    for _ in range(warmup):\n",
    "        X, Y = next(stream); step_fn(X, Y)\n",
    "    torch.cuda.synchronize(); t0 = time.perf_counter(); n = 0\n",
    "    for _ in range(timed):\n",
    "        X, Y = next(stream); step_fn(X, Y); n += X.numel()\n",
    "    torch.cuda.synchronize()\n",
    "    return n / (time.perf_counter() - t0)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-6 fw=pytorch -->\n",
    "\n",
    "One note on the metric: `throughput` times the *whole* pipeline —\n",
    "`next(stream)` includes the `DataLoader` and the host-to-device copy — so\n",
    "these are end-to-end tokens per second, the number a training run actually\n",
    "delivers, not a kernels-only figure. (At this scale the input pipeline is\n",
    "under two percent of a step; we measured that too.)"
   ]
  },
  {
   "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=pytorch -->\n",
    "\n",
    "Eager execution, tf32-fair (matmul precision already `'high'`). Before\n",
    "optimizing, we *classify*: profile one step and read where the time goes."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rung-0-baseline-profiled"
   },
   "source": [
    "opt = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
    "\n",
    "def step_eager(X, Y):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    loss = F.cross_entropy(model(X).reshape(-1, vocab_size), Y.reshape(-1))\n",
    "    loss.backward(); opt.step()\n",
    "\n",
    "for _ in range(3):\n",
    "    step_eager(*next(stream))\n",
    "with torch.profiler.profile(activities=[\n",
    "        torch.profiler.ProfilerActivity.CPU,\n",
    "        torch.profiler.ProfilerActivity.CUDA], acc_events=True) as prof:\n",
    "    for _ in range(5):\n",
    "        step_eager(*next(stream))\n",
    "print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=6))\n",
    "tput0 = throughput(step_eager)\n",
    "print(f'R0 eager: {tput0:.0f} tokens/s')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n                                               aten::mm         9.84%      13.373ms        14.07%      19.133ms      41.146us      67.959ms        53.15%      67.959ms     146.148us           465  \n       autograd::engine::evaluate_function: MmBackward0         1.60%       2.176ms        15.00%      20.388ms     131.537us       0.000us         0.00%      45.438ms     293.149us           155  \n                                            MmBackward0         1.44%       1.963ms        13.40%      18.212ms     117.498us       0.000us         0.00%      45.438ms     293.149us           155  \n                                           aten::linear         0.29%     394.204us         6.43%       8.736ms      56.360us       0.000us         0.00%      22.521ms     145.295us           155  \n                                           aten::matmul         0.53%     715.844us         5.41%       7.357ms      47.462us       0.000us         0.00%      22.521ms     145.295us           155  \n                                              aten::mul         6.15%       8.366ms        10.25%      13.933ms      21.111us      16.700ms        13.06%      16.700ms      25.303us           660  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \nSelf CPU time total: 135.957ms\nSelf CUDA time total: 127.873ms\n\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R0 eager: 262957 tokens/s\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-9 fw=pytorch -->\n",
    "\n",
    "Read the profile for the model's character. Matmuls (the attention\n",
    "projections and the feed-forward blocks) carry about half of the device\n",
    "time — but only about half: the rest is a long tail of small elementwise\n",
    "and normalization kernels (note the *hundreds* of `aten::mul` calls per\n",
    "five steps), each launched separately, none individually visible. And the\n",
    "CPU-total column is nearly as large as the CUDA total — remembering\n",
    "that section's caveat that these totals overlap across\n",
    "threads and streams, that still says the dispatch thread is busy most of\n",
    "the step. The diagnosis, then: partly compute-bound, but with a fusible\n",
    "elementwise tail and non-trivial launch traffic — which predicts that\n",
    "compilation should pay something here, and tensor-core precision after\n",
    "it. The rungs will check both."
   ]
  },
  {
   "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=pytorch -->\n",
    "\n",
    "**Rung 1 — compile (that section).** The profile promised a\n",
    "fusible tail and busy dispatch; compilation is the matching fix. Before\n",
    "the timing, the correctness gate — the compiled model must produce the\n",
    "same logits as the eager one (they agree to about $10^{-2}$; the residue\n",
    "is tf32 matmuls re-associated by fusion, not a wrong answer). *Same\n",
    "answer first, then faster:*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-1"
   },
   "source": [
    "compiled = torch.compile(model)\n",
    "def step_compiled(X, Y):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    loss = F.cross_entropy(compiled(X).reshape(-1, vocab_size), Y.reshape(-1))\n",
    "    loss.backward(); opt.step()\n",
    "\n",
    "t0 = time.perf_counter(); step_compiled(*next(stream)); torch.cuda.synchronize()\n",
    "print(f'first (compiling) step: {time.perf_counter() - t0:.1f} s')\n",
    "X, _ = next(stream)\n",
    "assert torch.allclose(compiled(X), model(X), atol=1e-2, rtol=1e-3)\n",
    "tput1 = throughput(step_compiled)\n",
    "print(f'R1 compiled: {tput1:.0f} tokens/s ({tput1 / tput0:.2f}x)')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "first (compiling) step: 1.9 s\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R1 compiled: 345595 tokens/s (1.31x)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-12 fw=pytorch -->\n",
    "\n",
    "The first call pays about two seconds of compile time; steady state pays\n",
    "that back at roughly 1.3× — real money, from a model that is *not*\n",
    "overhead-bound in the that section sense. The gain is the\n",
    "profile's tail, cashed in: the elementwise chains between matmuls fuse\n",
    "into a handful of generated kernels, and the launch traffic drops with\n",
    "them. A re-profile confirms the bottleneck moved:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-5"
   },
   "source": [
    "with torch.profiler.profile(activities=[\n",
    "        torch.profiler.ProfilerActivity.CPU,\n",
    "        torch.profiler.ProfilerActivity.CUDA], acc_events=True) as prof:\n",
    "    for _ in range(5):\n",
    "        step_compiled(*next(stream))\n",
    "print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=5))"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \nautograd::engine::evaluate_function: CompiledFunctio...         0.34%     401.876us        45.09%      53.268ms      10.654ms       0.000us         0.00%      72.528ms      14.506ms             5  \n                               CompiledFunctionBackward         1.01%       1.198ms        44.75%      52.866ms      10.573ms       0.000us         0.00%      72.528ms      14.506ms             5  \n## Call CompiledFxGraph fdcdeefsm5ww4waci2yvsvbra44l...        14.62%      17.273ms        43.73%      51.668ms      10.334ms       0.000us         0.00%      72.528ms      14.506ms             5  \n                                               aten::mm        10.74%      12.691ms        15.78%      18.642ms      40.527us      68.973ms        59.32%      68.973ms     149.941us           460  \n                             Torch-Compiled Region: 0/0         0.47%     550.283us        12.21%      14.430ms       2.886ms       0.000us         0.00%      35.378ms       7.076ms             5  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \nSelf CPU time total: 118.147ms\nSelf CUDA time total: 116.272ms\n\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-13 fw=pytorch -->\n",
    "\n",
    "The elementwise tail has folded into `CompiledFunction` regions, and the\n",
    "matmuls' share of device time has *risen* — the step is now more\n",
    "tensor-core-bound than before, which is precisely the regime that\n",
    "precision attacks. So, precision next."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-2"
   },
   "source": [
    "def step_bf16(X, Y):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    with torch.autocast('cuda', dtype=torch.bfloat16):\n",
    "        loss = F.cross_entropy(compiled(X).reshape(-1, vocab_size), Y.reshape(-1))\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "tput2 = throughput(step_bf16)\n",
    "print(f'R2 +bf16: {tput2:.0f} tokens/s ({tput2 / tput1:.2f}x)')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R2 +bf16: 468280 tokens/s (1.35x)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-14 fw=pytorch -->\n",
    "\n",
    "**Rung 3 — raise the batch, raise the intensity\n",
    "(that section, that section).** A bigger\n",
    "per-device batch raises the matmuls' arithmetic intensity, climbing the\n",
    "roofline toward the compute roof — a rung in its own right, whatever the\n",
    "precision. Memory is its usual price, so we also run a control: a few\n",
    "fp32 steps at batch 512, to see what the batch *would* cost without bf16.\n",
    "The control shows fp32 at 512 still fits this 24 GB card — snugly — so\n",
    "bf16 did not unlock this batch; what it bought is *headroom*, roughly\n",
    "half the footprint, which becomes the difference between fitting and not\n",
    "the moment the model, context, or batch grows again:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-3"
   },
   "source": [
    "big = d2l.TimeMachine(batch_size=512, num_steps=128, tokenization='char')\n",
    "big_stream = batches(big)\n",
    "\n",
    "def throughput_big(step_fn, warmup=60, timed=50):\n",
    "    for _ in range(warmup):\n",
    "        step_fn(*next(big_stream))\n",
    "    torch.cuda.synchronize(); t0 = time.perf_counter(); n = 0\n",
    "    for _ in range(timed):\n",
    "        X, Y = next(big_stream); step_fn(X, Y); n += X.numel()\n",
    "    torch.cuda.synchronize()\n",
    "    return n / (time.perf_counter() - t0)\n",
    "\n",
    "torch.cuda.reset_peak_memory_stats()      # control: batch 512 in fp32\n",
    "for _ in range(3):\n",
    "    step_eager(*next(big_stream))\n",
    "mem_fp32 = torch.cuda.max_memory_allocated() / 2**30\n",
    "torch.cuda.reset_peak_memory_stats()\n",
    "tput3 = throughput_big(step_bf16)\n",
    "print(f'R3 +batch-up (512): {tput3:.0f} tokens/s ({tput3 / tput2:.2f}x), '\n",
    "      f'peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB '\n",
    "      f'(fp32 control: {mem_fp32:.1f} GiB)')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R3 +batch-up (512): 622604 tokens/s (1.33x), peak 8.6 GiB (fp32 control: 16.6 GiB)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-15 fw=pytorch -->\n",
    "\n",
    "Another rung, another re-profile. At batch 512 the average matmul runs\n",
    "several times longer per call — bigger tiles, better tensor-core\n",
    "utilization — and the launch-and-dispatch side has receded to a footnote:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-6"
   },
   "source": [
    "with torch.profiler.profile(activities=[\n",
    "        torch.profiler.ProfilerActivity.CPU,\n",
    "        torch.profiler.ProfilerActivity.CUDA], acc_events=True) as prof:\n",
    "    for _ in range(5):\n",
    "        step_bf16(*next(big_stream))\n",
    "print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=5))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \nautograd::engine::evaluate_function: CompiledFunctio...         0.09%     453.758us        14.88%      78.479ms      15.696ms       0.000us         0.00%     334.474ms      66.895ms             5  \n                               CompiledFunctionBackward         0.30%       1.575ms        14.79%      78.025ms      15.605ms       0.000us         0.00%     334.474ms      66.895ms             5  \n## Call CompiledFxGraph fqqnnoosg4xsoqgb7lugw72ak7fg...         5.19%      27.371ms        14.49%      76.450ms      15.290ms       0.000us         0.00%     334.474ms      66.895ms             5  \n                                               aten::mm         2.56%      13.512ms         3.82%      20.170ms      43.377us     236.886ms        45.31%     236.886ms     509.432us           465  \n                             Torch-Compiled Region: 0/2         0.12%     653.556us         4.51%      23.769ms       4.754ms       0.000us         0.00%     179.349ms      35.870ms             5  \n-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \nSelf CPU time total: 527.430ms\nSelf CUDA time total: 522.771ms\n\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-16 fw=pytorch -->\n",
    "\n",
    "**Rung 4 — activation checkpointing, and why it does *not* help here\n",
    "(that section).** This is a deliberate *negative* rung.\n",
    "Checkpointing trades compute for memory — but at this scale memory is not\n",
    "the binding constraint (the bf16 batch-512 step fits in about a third of\n",
    "the card). Recomputing activations therefore only *costs* time and saves\n",
    "memory we did not need. Measuring the loss is the lesson: knowing when a\n",
    "technique does not apply is as much the method as knowing when it does.\n",
    "The wrapper below mirrors `GPT.forward` exactly, checkpointing each\n",
    "block — and because a mirror can drift, we assert it is one before timing\n",
    "anything:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-4"
   },
   "source": [
    "class CheckpointedGPT(nn.Module):\n",
    "    \"\"\"ch. 11's GPT with its transformer blocks run under checkpointing.\n",
    "    Reuses the exact submodules; only the block loop differs.\"\"\"\n",
    "    def __init__(self, gpt):\n",
    "        super().__init__(); self.gpt = gpt\n",
    "    def forward(self, X):\n",
    "        m = self.gpt                              # mirror GPT.forward exactly,\n",
    "        H = m.token_emb(X)                        # but checkpoint each block\n",
    "        if m.pos == 'learned':                    # (default 'rope' adds no emb)\n",
    "            H = H + m.pos_emb(torch.arange(X.shape[1], device=X.device))\n",
    "        for blk in m.blks:\n",
    "            H = ckpt.checkpoint(blk, H, use_reentrant=False)\n",
    "        return F.linear(m.norm(H), m.token_emb.weight)\n",
    "\n",
    "ckpt_gpt = CheckpointedGPT(model)\n",
    "X, _ = next(big_stream)\n",
    "assert torch.allclose(ckpt_gpt(X), model(X), atol=1e-6)   # an exact mirror\n",
    "\n",
    "ckpt_model = torch.compile(ckpt_gpt)\n",
    "def step_ckpt(X, Y):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    with torch.autocast('cuda', dtype=torch.bfloat16):\n",
    "        loss = F.cross_entropy(ckpt_model(X).reshape(-1, vocab_size), Y.reshape(-1))\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "torch.cuda.reset_peak_memory_stats()\n",
    "tput4 = throughput_big(step_ckpt)\n",
    "print(f'R4 +checkpoint: {tput4:.0f} tokens/s ({tput4 / tput3:.2f}x), '\n",
    "      f'peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB')"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R4 +checkpoint: 545503 tokens/s (0.88x), peak 3.1 GiB\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-17 fw=pytorch -->\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 allreduce every\n",
    "step; the NCCL collective on our host-staged box sustains around five\n",
    "GB/s per device in that section's bytes-per-device convention\n",
    "(the cell prices with 4.5, toward the conservative end of its run-to-run\n",
    "range; NCCL's own \"busbw\" reads ~2 GB/s — the default-fallback figure:\n",
    "that section measures a five-fold-faster configured\n",
    "mode, and shows why these runs nonetheless keep the library's defaults), which\n",
    "puts the allreduce at\n",
    "roughly the same tens of milliseconds as the compute itself — a\n",
    "transformer's parameters are\n",
    "proportional to its compute, so unlike a convolutional ResNet it offers\n",
    "little extra compute to hide communication behind. Summing the two terms,\n",
    "the equation predicts a *weak* two-GPU gain: scarcely more than one\n",
    "single-GPU throughput at $k=2$, under half of linear at $k=4$. One\n",
    "refinement before trusting the measurement: the serial sum is a *floor* —\n",
    "DDP overlaps its bucketed allreduce under backward\n",
    "(the figure), so the measurement can land above the\n",
    "serial prediction, and how far above measures the overlap. We reuse\n",
    "13.6's sidecar idiom on this very model (eager, per-rank batch 64 — weak\n",
    "scaling, the convention of that section; each rank\n",
    "streams its own batches, so this is a throughput measurement, not a\n",
    "convergence claim):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-7"
   },
   "source": [
    "GPT_DDP = r'''\n",
    "import json, os, time, torch\n",
    "import torch.distributed as dist\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "from d2l import torch as d2l\n",
    "\n",
    "rank = int(os.environ[\"LOCAL_RANK\"])\n",
    "torch.cuda.set_device(rank); dist.init_process_group(\"nccl\")\n",
    "torch.set_float32_matmul_precision(\"high\")\n",
    "\n",
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char')\n",
    "torch.manual_seed(0)\n",
    "model = d2l.GPT(len(data.vocab), num_hiddens=512, num_blks=6).to(rank)\n",
    "model = DDP(model, device_ids=[rank])\n",
    "opt = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
    "torch.manual_seed(rank)                 # each rank streams its own batches\n",
    "\n",
    "def batches():\n",
    "    while True:\n",
    "        for X, Y in data.train_dataloader():\n",
    "            yield X.to(rank), Y.to(rank)\n",
    "stream = batches()\n",
    "\n",
    "def step(X, Y):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    loss = F.cross_entropy(model(X).reshape(-1, len(data.vocab)),\n",
    "                           Y.reshape(-1))\n",
    "    loss.backward(); opt.step()\n",
    "\n",
    "for _ in range(30):                     # warmup: clocks, caches, NCCL\n",
    "    step(*next(stream))\n",
    "torch.cuda.synchronize(); t0 = time.time(); n = 0\n",
    "for _ in range(50):\n",
    "    X, Y = next(stream); step(X, Y); n += X.numel()\n",
    "torch.cuda.synchronize()\n",
    "here = os.path.dirname(os.path.abspath(__file__))\n",
    "with open(os.path.join(here, f\"rank{rank}.json\"), \"w\") as f:\n",
    "    json.dump({\"tokens_per_s\": n / (time.time() - t0)}, f)\n",
    "dist.destroy_process_group()\n",
    "'''\n",
    "GPT_DIR = pathlib.Path('ddp_gpt')\n",
    "GPT_DIR.mkdir(exist_ok=True)\n",
    "_ = (GPT_DIR / 'train_gpt_ddp.py').write_text(GPT_DDP)\n",
    "\n",
    "def ddp_gpt_tokens(k):\n",
    "    for stale in GPT_DIR.glob('rank*.json'):\n",
    "        stale.unlink()\n",
    "    torchrun = str(pathlib.Path(sys.executable).parent / 'torchrun')\n",
    "    subprocess.run([torchrun, '--standalone', f'--nproc-per-node={k}',\n",
    "                    str(GPT_DIR / 'train_gpt_ddp.py')], check=True)\n",
    "    total = 0.0\n",
    "    for r in range(k):\n",
    "        with open(GPT_DIR / f'rank{r}.json') as f:\n",
    "            total += json.load(f)['tokens_per_s']\n",
    "    return total"
   ],
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-rungs-each-one-measured-8"
   },
   "source": [
    "n_bytes = 4 * sum(p.numel() for p in model.parameters())\n",
    "beta = 4.5e9   # NCCL allreduce, effective bytes/device/s on this box (13.5)\n",
    "t_comm, t_cmp = 2 * n_bytes / beta, 64 * 128 / tput0\n",
    "print(f'per step: t_comm ~ {1e3 * t_comm:.0f} ms vs '\n",
    "      f't_compute ~ {1e3 * t_cmp:.0f} ms')\n",
    "torch.cuda.empty_cache()               # hand the ranks the parent's cache\n",
    "for k in (k for k in (2, 4) if k <= d2l.num_gpus()):\n",
    "    floor = k * 64 * 128 / (t_cmp + t_comm)\n",
    "    meas = ddp_gpt_tokens(k)\n",
    "    print(f'R5 k={k}: no-overlap floor {floor / 1e3:.0f}k tokens/s, '\n",
    "          f'measured {meas / 1e3:.0f}k ({meas / tput0:.2f}x of one GPU)')"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per step: t_comm ~ 34 ms vs t_compute ~ 31 ms\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R5 k=2: no-overlap floor 253k tokens/s, measured 297k (1.13x of one GPU)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "R5 k=4: no-overlap floor 506k tokens/s, measured 455k (1.73x of one GPU)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-18 fw=pytorch -->\n",
    "\n",
    "The measurement lands where the model said it would — pinned to the\n",
    "floor, nowhere near the linear ceiling. At $k=2$ it sits somewhat above\n",
    "the floor: DDP's overlapped buckets buy back part of the communication\n",
    "time. At $k=4$ it slips a little *below* it: the floor was priced with\n",
    "the $\\beta$ we measured at two GPUs, and four ranks staging through the\n",
    "same host path do not quite sustain that constant. Either way, data\n",
    "parallelism on this fabric turns a second GPU into a modest fraction of\n",
    "one, and four GPUs into well under half of linear. Communication-hungry\n",
    "is the diagnosis, and it was written down *before* the launch. On an NVLink\n",
    "box the same accounting predicts near-linear scaling — the constant\n",
    "changes, the method does not. Note what this rung is *not*: it uses the\n",
    "eager baseline, matching the prediction's inputs; stacking DDP on the\n",
    "compiled-bf16-batch-512 configuration is the exercise, and thinking\n",
    "through why its efficiency comes out *higher* is the point of doing it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-19 fw=all -->\n",
    "\n",
    "## The Waterfall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-20 fw=pytorch -->\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 while its\n",
    "ratio to the first bar prices the whole stack."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fast-transformer-the-waterfall"
   },
   "source": [
    "rungs = ['R0\\neager', 'R1\\ncompile', 'R2\\n+bf16', 'R3\\n+batch', 'R4\\n+ckpt']\n",
    "tputs = [tput0, tput1, tput2, tput3, tput4]\n",
    "print(f'cumulative, R0 -> R3: {tputs[3] / tputs[0]:.2f}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[0]:.1f}x', ha='center', va='bottom')\n",
    "d2l.plt.show()"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cumulative, R0 -> R3: 2.37x\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 31748 characters]"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-21 fw=pytorch -->\n",
    "\n",
    "Read left to right, and let the measurements calibrate the intuition.\n",
    "Compilation paid about 1.3× — on a model that was *not* overhead-bound —\n",
    "by fusing the elementwise tail the profile exposed. Bf16 paid about 1.4×\n",
    "on top, the tensor cores earning their silicon now that the matmuls are\n",
    "wide enough to feed them. Raising the batch bought roughly another 1.3×\n",
    "by climbing the roofline. And checkpointing, the red bar, *cost* about\n",
    "a tenth of the throughput while cutting peak memory to around a third\n",
    "(about 9 GiB down to 3) — a negative rung for *speed*, because memory was\n",
    "never the binding constraint at this scale, even though it did exactly\n",
    "what it promised for memory. The cumulative single-GPU speedup, baseline\n",
    "to the batch-up rung, is the number the cell prints — about 2.4× in our\n",
    "runs — and every increment traces to a section of this chapter. That the\n",
    "pieces compose; that no single rung dominated but three modest ones\n",
    "multiplied; that a technique can *cost* time when its constraint does not\n",
    "bind; and that the largest errors we found while building this section\n",
    "were in the *measurements*, not the model — that is the whole lesson:\n",
    "**you cannot copy a recipe — you measure, classify, fix, and\n",
    "re-measure.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-22 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",
   "metadata": {
    "id": "fast-transformer-the-waterfall-1"
   },
   "source": [
    "learner = make_gpt()\n",
    "learner_c = torch.compile(learner)\n",
    "opt_l = torch.optim.AdamW(learner.parameters(), lr=1e-3)\n",
    "losses = []\n",
    "for _ in range(300):\n",
    "    X, Y = next(big_stream)\n",
    "    opt_l.zero_grad(set_to_none=True)\n",
    "    with torch.autocast('cuda', dtype=torch.bfloat16):\n",
    "        loss = F.cross_entropy(learner_c(X).reshape(-1, vocab_size),\n",
    "                               Y.reshape(-1))\n",
    "    loss.backward(); opt_l.step()\n",
    "    losses.append(loss.detach())       # no host read in the hot loop (13.1)\n",
    "losses = torch.stack(losses).cpu()\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'"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "smoothed loss: first 20 steps 3.01 -> last 20 steps 0.07\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-23 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-24 fw=pytorch -->\n",
    "\n",
    "* The chapter's method applied whole: profile the baseline, classify the\n",
    "  regime, apply one technique per rung, re-profile. On ch. 11's GPT\n",
    "  (width 512, ~19M params), compile paid about 1.3×, bf16 about 1.4× on\n",
    "  top, a bigger batch roughly another 1.3× — no single dominant win, but\n",
    "  three modest ones multiplying to about 2.4×, with the waterfall cell\n",
    "  printing the exact cumulative figure.\n",
    "* Every rung is attributed to a section, and the bottleneck is\n",
    "  re-profiled as it moves: after compilation the elementwise tail is\n",
    "  fused away and the matmuls' share rises, which is what makes precision\n",
    "  the next fix.\n",
    "* Measurement is a technique too, and it can fail: a ragged final batch\n",
    "  put a `torch.compile` retrace inside our timing window, and the\n",
    "  profiler's instrumentation outlived its cell until told otherwise.\n",
    "  Both had to be fixed before the waterfall meant anything.\n",
    "* Two rungs teach by *not* helping: bf16 is a negative rung at width 256\n",
    "  (matmuls too small for the tensor cores), and activation checkpointing\n",
    "  is a negative rung here (memory is not the binding constraint). Knowing\n",
    "  when a technique does not apply is the method too.\n",
    "* Data parallelism is predicted before it is measured: a transformer's\n",
    "  parameters scale with its compute, so DP on our host-staged fabric is\n",
    "  communication-hungry — the measured throughput lands at the cost\n",
    "  model's no-overlap floor (a shade above it at two GPUs, a shade below\n",
    "  at four, where the host-staged fabric no longer sustains its two-GPU\n",
    "  bandwidth), nowhere near the linear ceiling. The\n",
    "  prediction-then-measurement is the demonstration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=fast-transformer-md-25 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": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
