{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-1 fw=all -->\n",
    "\n",
    "# Memory and Precision\n",
    "\n",
    "So far this chapter has fought for *time*: fewer memory round trips, fewer\n",
    "kernel launches, more arithmetic per byte. This section fights for *space*.\n",
    "Sooner or later a model stops fitting: the optimizer states, the\n",
    "activations held for the backward pass, and a large enough batch together\n",
    "overflow the card, and training dies with an out-of-memory error long\n",
    "before it is slow. The good news is that memory is the most *legible* of\n",
    "the resources — you can write down exactly where every byte of a training\n",
    "step goes, measure it against that prediction, and then spend three\n",
    "well-understood techniques to buy headroom: lower-precision formats, which\n",
    "also buy speed; activation checkpointing, which trades compute for memory;\n",
    "and gradient accumulation, which trades *time* for the effect of a bigger\n",
    "batch. Together these are the \"it doesn't fit\" rung of the ladder, and\n",
    "they are what stands between a 24 GB card and a model that wants 30.\n",
    "\n",
    "*Prerequisites: the format ladder and tensor-core requirement of*\n",
    "that section*; the* `d2l.Benchmark` *timer of*\n",
    "that section*; the AdamW optimizer states of*\n",
    "that section*. This section finally discharges the \"gradient\n",
    "accumulation, returned to in the performance chapter\" promise of*\n",
    "that section*.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-memory-and-precision"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import torch as d2l\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "torch.set_float32_matmul_precision('high')"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-2 fw=all -->\n",
    "\n",
    "## The Memory Anatomy of a Training Step\n",
    "\n",
    "A training step holds four kinds of tensor in memory, and knowing their\n",
    "sizes lets you predict the peak before you run anything\n",
    "(the figure). Let the model have $P$ parameters.\n",
    "\n",
    "* **Parameters** — $P$ values. In mixed-precision training these are kept\n",
    "  in fp32 (4 bytes) as the master copy: $4P$ bytes.\n",
    "* **Gradients** — one per parameter, produced during the backward pass:\n",
    "  another $4P$ bytes.\n",
    "* **Optimizer states** — Adam keeps a first and second moment per\n",
    "  parameter (that section), $2 \\times 4P = 8P$ bytes; plain SGD\n",
    "  keeps none.\n",
    "* **Activations** — every intermediate result the forward pass produces\n",
    "  that the backward pass will need. Unlike the first three, this scales\n",
    "  with the *batch* and the *depth*: roughly\n",
    "  $\\textrm{batch} \\times \\textrm{sequence} \\times \\textrm{width} \\times\n",
    "  \\textrm{depth}$, and it is usually the term that blows up.\n",
    "\n",
    "![Where a training step's memory goes, over the step timeline. Parameters\n",
    "and optimizer states are constant; activations grow through the forward\n",
    "pass and are released through the backward pass; gradients fill in as\n",
    "they are computed. Peak memory lands early in the backward pass, when\n",
    "activations are still full.](../img/mdl-perf-memory-anatomy.svg)\n",
    "\n",
    "Add the constant terms: training a model with Adam costs about\n",
    "$4P + 4P + 8P = 16P$ bytes *before a single activation* — sixteen bytes\n",
    "per parameter, the number to carry in your head. (Mixed precision, below,\n",
    "nudges this to roughly 18–20 by keeping some tensors in two formats; it\n",
    "still saves memory overall, because activations, the dominant term,\n",
    "halve.) The shape of the figure is the other half of the\n",
    "lesson: activations *accumulate* through the forward pass and are *freed*\n",
    "as the backward pass consumes them, so peak memory occurs early in\n",
    "backward — which is exactly the moment activation checkpointing will\n",
    "target.\n",
    "\n",
    "## Measuring Memory\n",
    "\n",
    "The frameworks let you check the anatomy against reality, and — as with\n",
    "timing — the two do it in characteristically different ways. PyTorch keeps\n",
    "running *counters* you query at runtime; JAX lets the compiler *plan*\n",
    "memory ahead of time and reports the plan. Start with PyTorch's counters,\n",
    "the numbers the anatomy table predicts:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-measuring-memory-1"
   },
   "source": [
    "gpu = d2l.try_gpu()\n",
    "net = nn.Sequential(*[nn.Linear(2048, 2048) for _ in range(6)]).to(gpu)\n",
    "X = torch.randn(256, 2048, device=gpu)\n",
    "opt = torch.optim.Adam(net.parameters())\n",
    "\n",
    "torch.cuda.reset_peak_memory_stats()\n",
    "opt.zero_grad(set_to_none=True)\n",
    "net(X).sum().backward()\n",
    "opt.step()\n",
    "P = sum(p.numel() for p in net.parameters())\n",
    "print(f'{P/1e6:.1f}M params  '\n",
    "      f'predicted states+grads+params: {16 * P / 1e6:.0f} MB')\n",
    "print(f'measured peak: {torch.cuda.max_memory_allocated() / 1e6:.0f} MB '\n",
    "      f'allocated, {torch.cuda.max_memory_reserved() / 1e6:.0f} MB reserved')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "25.2M params  predicted states+grads+params: 403 MB\nmeasured peak: 524 MB allocated, 547 MB reserved\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-3 fw=all -->\n",
    "\n",
    "The measured peak overshoots the $16P$ prediction by roughly 30%, and the\n",
    "gap has an anatomy of its own: this stack's activations (about 13 MB),\n",
    "autograd bookkeeping and cuBLAS workspaces, temporaries that the fused\n",
    "optimizer update allocates (about the size of one more gradient set), and\n",
    "the caching allocator's block rounding. The second printed number is\n",
    "larger again: `max_memory_allocated` counts bytes inside live tensors,\n",
    "while `max_memory_reserved` is the caching allocator's high-water mark —\n",
    "the distinction that section dissects. So read $16P$ as this\n",
    "*configuration's* floor — fp32 parameters under Adam — rather than as the\n",
    "whole anatomy; and whether activations dominate the total, as they do in\n",
    "the transformer regime, is workload-dependent: batch, sequence length,\n",
    "and depth decide.\n",
    "\n",
    "For a finer picture, PyTorch can record every allocation and free as a\n",
    "*snapshot*. Rather than embed its interactive viewer, we reconstruct the\n",
    "one plot that matters — allocated bytes over the step — directly from the\n",
    "recorded trace, and see the sawtooth of the figure in\n",
    "real data:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-measuring-memory-2"
   },
   "source": [
    "torch.cuda.memory._record_memory_history(max_entries=100000)\n",
    "for _ in range(3):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    net(X).sum().backward()\n",
    "    opt.step()\n",
    "snapshot = torch.cuda.memory._snapshot()\n",
    "torch.cuda.memory._record_memory_history(enabled=None)\n",
    "\n",
    "# Reconstruct cumulative allocated bytes from the allocation trace.\n",
    "cum, series = 0, []\n",
    "for ev in snapshot['device_traces'][0]:\n",
    "    if ev['action'] == 'alloc':\n",
    "        cum += ev['size']\n",
    "    elif ev['action'] == 'free_completed':\n",
    "        cum -= ev['size']\n",
    "    series.append(cum / 1e6)\n",
    "d2l.plot(list(range(len(series))), [series],\n",
    "         'allocation event', 'allocated (MB)')"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"257.845313pt\" height=\"183.35625pt\" viewBox=\"0 0 257.845313 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-20T20:17:12.176580</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 257.845313 183.35625 \nL 257.845313 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 55.345312 145.8 \nL 250.645312 145.8 \nL 250.645312 7.2 \nL 55.345312 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path d=\"M 64.222585 145.8 \nL 64.222585 7.2 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"m11e24b347d\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m11e24b347d\" x=\"64.222585\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0 -->\n      <g transform=\"translate(61.041335 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path d=\"M 112.468633 145.8 \nL 112.468633 7.2 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m11e24b347d\" x=\"112.468633\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 100 -->\n      <g transform=\"translate(102.924883 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path d=\"M 160.71468 145.8 \nL 160.71468 7.2 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m11e24b347d\" x=\"160.71468\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 200 -->\n      <g transform=\"translate(151.17093 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path d=\"M 208.960728 145.8 \nL 208.960728 7.2 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m11e24b347d\" x=\"208.960728\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 300 -->\n      <g transform=\"translate(199.416978 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-33\" d=\"M 2597 2516 \nQ 3050 2419 3304 2112 \nQ 3559 1806 3559 1356 \nQ 3559 666 3084 287 \nQ 2609 -91 1734 -91 \nQ 1441 -91 1130 -33 \nQ 819 25 488 141 \nL 488 750 \nQ 750 597 1062 519 \nQ 1375 441 1716 441 \nQ 2309 441 2620 675 \nQ 2931 909 2931 1356 \nQ 2931 1769 2642 2001 \nQ 2353 2234 1838 2234 \nL 1294 2234 \nL 1294 2753 \nL 1863 2753 \nQ 2328 2753 2575 2939 \nQ 2822 3125 2822 3475 \nQ 2822 3834 2567 4026 \nQ 2313 4219 1838 4219 \nQ 1578 4219 1281 4162 \nQ 984 4106 628 3988 \nL 628 4550 \nQ 988 4650 1302 4700 \nQ 1616 4750 1894 4750 \nQ 2613 4750 3031 4423 \nQ 3450 4097 3450 3541 \nQ 3450 3153 3228 2886 \nQ 3006 2619 2597 2516 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-33\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_5\">\n     <!-- allocation event -->\n     <g transform=\"translate(112.872656 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-63\" d=\"M 3122 3366 \nL 3122 2828 \nQ 2878 2963 2633 3030 \nQ 2388 3097 2138 3097 \nQ 1578 3097 1268 2742 \nQ 959 2388 959 1747 \nQ 959 1106 1268 751 \nQ 1578 397 2138 397 \nQ 2388 397 2633 464 \nQ 2878 531 3122 666 \nL 3122 134 \nQ 2881 22 2623 -34 \nQ 2366 -91 2075 -91 \nQ 1284 -91 818 406 \nQ 353 903 353 1747 \nQ 353 2603 823 3093 \nQ 1294 3584 2113 3584 \nQ 2378 3584 2631 3529 \nQ 2884 3475 3122 3366 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6e\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-76\" d=\"M 191 3500 \nL 800 3500 \nL 1894 563 \nL 2988 3500 \nL 3597 3500 \nL 2284 0 \nL 1503 0 \nL 191 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-61\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(61.279297 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(89.0625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(116.845703 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(178.027344 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(233.007812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(294.287109 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(333.496094 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(361.279297 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(422.460938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(485.839844 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(517.626953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-76\" transform=\"translate(579.150391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(638.330078 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(699.853516 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(763.232422 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_9\">\n      <path d=\"M 55.345312 139.054333 \nL 250.645312 139.054333 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <defs>\n       <path id=\"m91e15e73b1\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m91e15e73b1\" x=\"55.345312\" y=\"139.054333\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- −100 -->\n      <g transform=\"translate(20.878125 142.853552) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2212\" d=\"M 678 2272 \nL 4684 2272 \nL 4684 1741 \nL 678 1741 \nL 678 2272 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-2212\"/>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(83.789062 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(147.412109 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(211.035156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_11\">\n      <path d=\"M 55.345312 107.777166 \nL 250.645312 107.777166 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#m91e15e73b1\" x=\"55.345312\" y=\"107.777166\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- −50 -->\n      <g transform=\"translate(27.240625 111.576385) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-2212\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(83.789062 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(147.412109 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_13\">\n      <path d=\"M 55.345312 76.5 \nL 250.645312 76.5 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#m91e15e73b1\" x=\"55.345312\" y=\"76.5\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 0 -->\n      <g transform=\"translate(41.982812 80.299219) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_15\">\n      <path d=\"M 55.345312 45.222834 \nL 250.645312 45.222834 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#m91e15e73b1\" x=\"55.345312\" y=\"45.222834\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 50 -->\n      <g transform=\"translate(35.620312 49.022052) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-35\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_17\">\n      <path d=\"M 55.345312 13.945667 \nL 250.645312 13.945667 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_18\">\n      <g>\n       <use xlink:href=\"#m91e15e73b1\" x=\"55.345312\" y=\"13.945667\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 100 -->\n      <g transform=\"translate(29.257812 17.744886) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(127.246094 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_11\">\n     <!-- allocated (MB) -->\n     <g transform=\"translate(14.798437 112.660938) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-64\" d=\"M 2906 2969 \nL 2906 4863 \nL 3481 4863 \nL 3481 0 \nL 2906 0 \nL 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nz\nM 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-28\" d=\"M 1984 4856 \nQ 1566 4138 1362 3434 \nQ 1159 2731 1159 2009 \nQ 1159 1288 1364 580 \nQ 1569 -128 1984 -844 \nL 1484 -844 \nQ 1016 -109 783 600 \nQ 550 1309 550 2009 \nQ 550 2706 781 3412 \nQ 1013 4119 1484 4856 \nL 1984 4856 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-4d\" d=\"M 628 4666 \nL 1569 4666 \nL 2759 1491 \nL 3956 4666 \nL 4897 4666 \nL 4897 0 \nL 4281 0 \nL 4281 4097 \nL 3078 897 \nL 2444 897 \nL 1241 4097 \nL 1241 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-42\" d=\"M 1259 2228 \nL 1259 519 \nL 2272 519 \nQ 2781 519 3026 730 \nQ 3272 941 3272 1375 \nQ 3272 1813 3026 2020 \nQ 2781 2228 2272 2228 \nL 1259 2228 \nz\nM 1259 4147 \nL 1259 2741 \nL 2194 2741 \nQ 2656 2741 2882 2914 \nQ 3109 3088 3109 3444 \nQ 3109 3797 2882 3972 \nQ 2656 4147 2194 4147 \nL 1259 4147 \nz\nM 628 4666 \nL 2241 4666 \nQ 2963 4666 3353 4366 \nQ 3744 4066 3744 3513 \nQ 3744 3084 3544 2831 \nQ 3344 2578 2956 2516 \nQ 3422 2416 3680 2098 \nQ 3938 1781 3938 1306 \nQ 3938 681 3513 340 \nQ 3088 0 2303 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-29\" d=\"M 513 4856 \nL 1013 4856 \nQ 1481 4119 1714 3412 \nQ 1947 2706 1947 2009 \nQ 1947 1309 1714 600 \nQ 1481 -109 1013 -844 \nL 513 -844 \nQ 928 -128 1133 580 \nQ 1338 1288 1338 2009 \nQ 1338 2731 1133 3434 \nQ 928 4138 513 4856 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-61\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(61.279297 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(89.0625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(116.845703 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(178.027344 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(233.007812 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(294.287109 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(333.496094 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(395.019531 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(458.496094 0)\"/>\n      <use xlink:href=\"#DejaVuSans-28\" transform=\"translate(490.283203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-4d\" transform=\"translate(529.296875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-42\" transform=\"translate(615.576172 0)\"/>\n      <use xlink:href=\"#DejaVuSans-29\" transform=\"translate(684.179688 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_19\">\n    <path d=\"M 64.222585 76.5 \nL 64.705046 86.994876 \nL 66.152427 87 \nL 66.634888 97.494876 \nL 68.082269 97.5 \nL 68.564729 107.994876 \nL 70.012111 108 \nL 70.494571 118.494876 \nL 71.941953 118.5 \nL 72.424413 128.994876 \nL 73.871795 129 \nL 74.354255 139.494876 \nL 75.319176 139.5 \nL 78.213939 131.628843 \nL 82.073623 131.628841 \nL 82.556083 132.9407 \nL 83.038544 132.940698 \nL 84.003465 130.316979 \nL 84.485925 130.316979 \nL 84.968386 131.628838 \nL 85.450846 121.133963 \nL 85.933307 119.822103 \nL 86.415767 119.822103 \nL 86.898228 121.133963 \nL 87.863148 121.128838 \nL 88.345609 122.440698 \nL 88.828069 121.128838 \nL 89.31053 110.633963 \nL 90.275451 110.628838 \nL 90.757911 111.940698 \nL 91.240372 111.940698 \nL 91.722832 113.252557 \nL 92.205293 111.940698 \nL 92.687753 101.445822 \nL 93.652674 101.440698 \nL 94.135135 102.752557 \nL 94.617595 102.752557 \nL 95.100056 104.064417 \nL 95.582516 102.752557 \nL 96.064977 92.257682 \nL 97.029897 92.252557 \nL 97.512358 93.564417 \nL 97.994818 93.564417 \nL 98.477279 94.876276 \nL 98.959739 93.564417 \nL 99.4422 83.069541 \nL 100.407121 83.064417 \nL 100.889581 84.376276 \nL 101.372042 84.376276 \nL 101.854502 85.688136 \nL 102.336963 75.19326 \nL 103.301884 75.188136 \nL 103.784344 76.499995 \nL 105.714186 76.5 \nL 106.196646 66.005124 \nL 106.679107 66 \nL 107.161567 55.505124 \nL 107.644028 55.5 \nL 108.126488 45.005124 \nL 108.608949 45 \nL 109.091409 34.505124 \nL 109.57387 34.5 \nL 110.05633 24.005124 \nL 110.538791 24 \nL 111.021251 13.505124 \nL 112.951093 13.505124 \nL 113.433554 24 \nL 114.880935 24.005124 \nL 115.363396 34.5 \nL 116.810777 34.505124 \nL 117.293237 45 \nL 118.740619 45.005124 \nL 119.223079 55.5 \nL 120.670461 55.505124 \nL 121.152921 66 \nL 122.600303 66.005124 \nL 123.082763 76.5 \nL 123.565224 76.5 \nL 124.047684 86.994876 \nL 125.495065 87 \nL 125.977526 97.494876 \nL 127.424907 97.5 \nL 127.907368 107.994876 \nL 129.354749 108 \nL 129.83721 118.494876 \nL 131.284591 118.5 \nL 131.767052 128.994876 \nL 133.214433 129 \nL 133.696894 139.494876 \nL 134.661814 139.5 \nL 137.556577 131.628843 \nL 141.416261 131.628841 \nL 141.898722 132.9407 \nL 142.381182 132.940698 \nL 143.346103 130.316979 \nL 143.828563 130.316979 \nL 144.311024 131.628838 \nL 144.793484 121.133963 \nL 145.275945 119.822103 \nL 145.758405 119.822103 \nL 146.240866 121.133963 \nL 147.205787 121.128838 \nL 147.688247 122.440698 \nL 148.170708 121.128838 \nL 148.653168 110.633963 \nL 149.618089 110.628838 \nL 150.10055 111.940698 \nL 150.58301 111.940698 \nL 151.065471 113.252557 \nL 151.547931 111.940698 \nL 152.030392 101.445822 \nL 152.995313 101.440698 \nL 153.477773 102.752557 \nL 153.960233 102.752557 \nL 154.442694 104.064417 \nL 154.925154 102.752557 \nL 155.407615 92.257682 \nL 156.372536 92.252557 \nL 156.854996 93.564417 \nL 157.337457 93.564417 \nL 157.819917 94.876276 \nL 158.302378 93.564417 \nL 158.784838 83.069541 \nL 159.749759 83.064417 \nL 160.23222 84.376276 \nL 160.71468 84.376276 \nL 161.197141 85.688136 \nL 161.679601 75.19326 \nL 162.644522 75.188136 \nL 163.126982 76.499995 \nL 165.056824 76.5 \nL 165.539285 66.005124 \nL 166.021745 66 \nL 166.504206 55.505124 \nL 166.986666 55.5 \nL 167.469127 45.005124 \nL 167.951587 45 \nL 168.434048 34.505124 \nL 168.916508 34.5 \nL 169.398969 24.005124 \nL 169.881429 24 \nL 170.36389 13.505124 \nL 172.293731 13.505124 \nL 172.776192 24 \nL 174.223573 24.005124 \nL 174.706034 34.5 \nL 176.153415 34.505124 \nL 176.635876 45 \nL 178.083257 45.005124 \nL 178.565718 55.5 \nL 180.013099 55.505124 \nL 180.49556 66 \nL 181.942941 66.005124 \nL 182.425401 76.5 \nL 182.907862 76.5 \nL 183.390322 86.994876 \nL 184.837704 87 \nL 185.320164 97.494876 \nL 186.767546 97.5 \nL 187.250006 107.994876 \nL 188.697388 108 \nL 189.179848 118.494876 \nL 190.627229 118.5 \nL 191.10969 128.994876 \nL 192.557071 129 \nL 193.039532 139.494876 \nL 194.004453 139.5 \nL 196.899216 131.628843 \nL 200.758899 131.628841 \nL 201.24136 132.9407 \nL 201.72382 132.940698 \nL 202.688741 130.316979 \nL 203.171202 130.316979 \nL 203.653662 131.628838 \nL 204.136123 121.133963 \nL 204.618583 119.822103 \nL 205.101044 119.822103 \nL 205.583504 121.133963 \nL 206.548425 121.128838 \nL 207.030886 122.440698 \nL 207.513346 121.128838 \nL 207.995807 110.633963 \nL 208.960728 110.628838 \nL 209.443188 111.940698 \nL 209.925648 111.940698 \nL 210.408109 113.252557 \nL 210.890569 111.940698 \nL 211.37303 101.445822 \nL 212.337951 101.440698 \nL 212.820411 102.752557 \nL 213.302872 102.752557 \nL 213.785332 104.064417 \nL 214.267793 102.752557 \nL 214.750253 92.257682 \nL 215.715174 92.252557 \nL 216.197635 93.564417 \nL 216.680095 93.564417 \nL 217.162556 94.876276 \nL 217.645016 93.564417 \nL 218.127477 83.069541 \nL 219.092397 83.064417 \nL 219.574858 84.376276 \nL 220.057318 84.376276 \nL 220.539779 85.688136 \nL 221.022239 75.19326 \nL 221.98716 75.188136 \nL 222.469621 76.499995 \nL 224.399463 76.5 \nL 224.881923 66.005124 \nL 225.364384 66 \nL 225.846844 55.505124 \nL 226.329305 55.5 \nL 226.811765 45.005124 \nL 227.294226 45 \nL 227.776686 34.505124 \nL 228.259146 34.5 \nL 228.741607 24.005124 \nL 229.224067 24 \nL 229.706528 13.505124 \nL 231.63637 13.505124 \nL 232.11883 24 \nL 233.566212 24.005124 \nL 234.048672 34.5 \nL 235.496054 34.505124 \nL 235.978514 45 \nL 237.425896 45.005124 \nL 237.908356 55.5 \nL 239.355737 55.505124 \nL 239.838198 66 \nL 241.285579 66.005124 \nL 241.76804 76.5 \nL 241.76804 76.5 \n\" clip-path=\"url(#pc59a0b66da)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 55.345312 145.8 \nL 55.345312 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 250.645312 145.8 \nL 250.645312 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 55.345312 145.8 \nL 250.645313 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 55.345312 7.2 \nL 250.645313 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pc59a0b66da\">\n   <rect x=\"55.345312\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-4 fw=all -->\n",
    "\n",
    "Each cycle is one training step: memory climbs through the forward pass as\n",
    "activations accumulate, then falls through the backward pass as they are\n",
    "released — exactly the anatomy, now measured. The JAX side plans instead\n",
    "of counts. Ahead-of-time compilation (that section) hands\n",
    "back an object whose `memory_analysis()` reports what the compiler\n",
    "*decided* to allocate, before any memory is touched:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-5 fw=all -->\n",
    "\n",
    "The contrast is worth stating plainly: **PyTorch counts allocations as\n",
    "they happen; XLA plans memory at compile time.** Neither is better; they\n",
    "reflect the eager-versus-traced split of that section, and\n",
    "knowing which mental model your framework uses tells you where to look\n",
    "when memory surprises you. Just do not compare the two frameworks' digits\n",
    "head-to-head: a compiler's plan is not an allocator's high-water mark.\n",
    "\n",
    "## Mixed Precision\n",
    "\n",
    "The format ladder of that section promised that every halving of\n",
    "width wins twice — double the tensor-core throughput and half the bytes.\n",
    "Mixed precision cashes that promise for training. The idea: run the\n",
    "compute-heavy operations (matmuls, convolutions) in bf16, where tensor\n",
    "cores are fastest and activations are half-size, while keeping a few\n",
    "numerically delicate things (the master weights, the loss reduction) in\n",
    "fp32. On our Ada card the arithmetic win is real and robust — but only if\n",
    "you know what your baseline is running, the lesson that section\n",
    "flagged. Whether plain fp32 matmuls use the tensor cores' tf32 path is a\n",
    "framework default: PyTorch ships with tf32 *off*, so a naive fp32\n",
    "baseline is needlessly slow and inflates the apparent speedup; JAX on\n",
    "this card runs tf32-class dot compute *by default*, so its naive baseline\n",
    "is already fair — but is not the true fp32 it appears to be. We time all\n",
    "three configurations, once, so the difference is unmistakable:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-mixed-precision"
   },
   "source": [
    "net = nn.Sequential(nn.Flatten(),\n",
    "                    nn.Linear(64 * 64, 8192), nn.ReLU(),\n",
    "                    nn.Linear(8192, 8192), nn.ReLU(),\n",
    "                    nn.Linear(8192, 10)).to(gpu)\n",
    "X = torch.randn(2048, 1, 64, 64, device=gpu)\n",
    "y = torch.randint(0, 10, (2048,), device=gpu)\n",
    "opt = torch.optim.SGD(net.parameters(), lr=0.01)\n",
    "loss = nn.CrossEntropyLoss()\n",
    "\n",
    "def step(autocast, dtype=None):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    if autocast:\n",
    "        with torch.autocast('cuda', dtype=dtype):\n",
    "            l = loss(net(X), y)\n",
    "    else:\n",
    "        l = loss(net(X), y)\n",
    "    l.backward(); opt.step()\n",
    "\n",
    "torch.set_float32_matmul_precision('highest')  # Unfair baseline: tf32 off\n",
    "print(d2l.Benchmark(lambda: step(False), desc='fp32 (tf32 off)'))\n",
    "torch.set_float32_matmul_precision('high')     # Fair baseline: tf32 on\n",
    "print(d2l.Benchmark(lambda: step(False), desc='tf32'))\n",
    "print(d2l.Benchmark(lambda: step(True, torch.bfloat16), desc='bf16 autocast'))"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "fp32 (tf32 off): 22.58 ms/call\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "tf32: 17.79 ms/call\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "bf16 autocast: 9.78 ms/call\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-6 fw=all -->\n",
    "\n",
    "Against the fair tf32 baseline, bf16 runs about one and a half to two\n",
    "times as fast here — the exact ratio is framework- and shape-dependent —\n",
    "and the three timings tell the same two-step story in both tabs: moving\n",
    "true fp32 dots onto the tf32 tensor cores already buys a sizable step,\n",
    "and bf16 then adds at least as much again. Both are genuine tensor-core\n",
    "wins, not measurement artifacts.\n",
    "Note what the PyTorch tab does *not* use: a `GradScaler`. Loss scaling\n",
    "exists to keep tiny fp16 gradients from underflowing fp16's narrow\n",
    "5-bit-exponent range; bf16 shares fp32's 8-bit exponent\n",
    "(the figure), so its gradients do not underflow and no\n",
    "scaler is needed. fp16-plus-`GradScaler` is the pre-Ampere legacy path,\n",
    "worth one sentence and an exercise. The JAX tab makes the philosophical\n",
    "difference visible: precision in JAX is *explicit* — you thread dtypes\n",
    "through the computation, and `jax.default_matmul_precision` sets how fp32\n",
    "dots are *computed* (true fp32 versus tf32-class tensor-core compute)\n",
    "without changing what any tensor *stores*. On this card JAX's default\n",
    "matches `'high'`, tf32-class, bit for bit — which is why the middle\n",
    "timing is the fair baseline. PyTorch decides per-operation what to cast;\n",
    "JAX makes you say it. Both land on the same design point: bf16 halves the\n",
    "bytes and doubles the tensor-core rate.\n",
    "\n",
    "## Activation Checkpointing\n",
    "\n",
    "Mixed precision halves activation memory; activation checkpointing can cut\n",
    "it much further, by refusing to *store* most activations at all. Recall\n",
    "from the figure that activations are held from the\n",
    "moment the forward pass produces them until the backward pass consumes\n",
    "them — that storage is the peak. Checkpointing makes a different trade:\n",
    "store only a few activations (at block boundaries), and when the backward\n",
    "pass needs the ones in between, *recompute them* by re-running that block's\n",
    "forward pass. You pay extra compute — at most one extra forward pass — to\n",
    "avoid holding a large fraction of the activations; the cells below\n",
    "measure both sides of the trade. This is the same\n",
    "\"recompute rather than store\" argument the\n",
    "Mamba kernel of that section made one level down in the hardware\n",
    "[@Chen.Xu.Zhang.ea.2016]."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-activation-checkpointing"
   },
   "source": [
    "from torch.utils.checkpoint import checkpoint\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d=1024):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))\n",
    "    def forward(self, x):\n",
    "        return x + self.net(x)\n",
    "\n",
    "blocks = nn.ModuleList([Block().to(gpu) for _ in range(16)])\n",
    "X = torch.randn(16384, 1024, device=gpu, requires_grad=True)\n",
    "\n",
    "def run(use_ckpt):\n",
    "    h = X\n",
    "    for blk in blocks:\n",
    "        # use_reentrant=False is the modern, mandatory-since-2.9 form\n",
    "        h = checkpoint(blk, h, use_reentrant=False) if use_ckpt else blk(h)\n",
    "    return h.sum()\n",
    "\n",
    "for use_ckpt in (False, True):\n",
    "    for p in blocks.parameters():  # clear grads so the peaks are comparable\n",
    "        p.grad = None\n",
    "    X.grad = None\n",
    "    torch.cuda.reset_peak_memory_stats()\n",
    "    run(use_ckpt).backward()\n",
    "    peak = torch.cuda.max_memory_allocated() / 1e6\n",
    "    t = d2l.Benchmark(lambda: run(use_ckpt).backward()).time\n",
    "    tag = 'checkpointed' if use_ckpt else 'store all'\n",
    "    print(f'{tag}: peak {peak:.0f} MB, fwd+bwd {1000 * t:.1f} ms')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "store all: peak 4318 MB, fwd+bwd 59.9 ms\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "checkpointed: peak 2306 MB, fwd+bwd 71.1 ms\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-7 fw=all -->\n",
    "\n",
    "Checkpointing releases about 2 GB of held activations on this 16-block\n",
    "stack — the measured peak falls by nearly half in the allocator's count\n",
    "and by more than half in the compiler's plan — and the completion-timed\n",
    "step runs only about 15–20% longer, comfortably under the one-extra-forward bound, since the\n",
    "backward pass dominates the step. It is a trade you make when memory is\n",
    "the binding constraint: checkpointing normally costs time, and its\n",
    "purpose is memory (the rare configurations where recomputation beats a\n",
    "cache-hostile stash do not change the rule). JAX exposes the same\n",
    "mechanism as `jax.checkpoint` (equivalently `jax.remat`), with a nicer\n",
    "knob: policies like `checkpoint_dots` let you say *which* intermediates to\n",
    "keep, and `jax.remat`'s companion tools let you *watch* what was saved —\n",
    "the introspection theme JAX does uniquely well. The trade is identical;\n",
    "the interface reflects each framework's temperament, imperative versus\n",
    "declarative.\n",
    "\n",
    "## Gradient Accumulation\n",
    "\n",
    "The final technique buys the *statistical* effect of a large batch without\n",
    "its memory cost — and discharges a promise made all the way back in\n",
    "that section. A large batch improves gradient estimates and\n",
    "tensor-core utilization (its intensity is higher, that section),\n",
    "but a large batch's activations may not fit. Gradient accumulation splits\n",
    "the desired *global* batch into $k$ *micro-batches*, runs forward and\n",
    "backward on each while *summing* the gradients, and only then takes one\n",
    "optimizer step. The activations of just one micro-batch are live at a\n",
    "time, so peak memory follows the micro-batch, while the update sees the\n",
    "full global batch: $B_{\\textrm{global}} = B_{\\textrm{micro}} \\times k$\n",
    "(times the number of devices, once that section adds them). The\n",
    "parity check — that accumulating $k$ micro-batches matches one full-batch\n",
    "step — is worth seeing, because it is the whole correctness claim:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "memory-precision-gradient-accumulation"
   },
   "source": [
    "net = nn.Linear(1024, 1).to(gpu)\n",
    "X = torch.randn(128, 1024, device=gpu)\n",
    "y = torch.randn(128, 1, device=gpu)\n",
    "loss = nn.MSELoss()\n",
    "\n",
    "# Full batch: one backward over all 128 rows.\n",
    "net.zero_grad()\n",
    "loss(net(X), y).backward()\n",
    "full = net.weight.grad.clone()\n",
    "\n",
    "# Accumulated: 4 micro-batches of 32, gradients summed, scaled by 1/4.\n",
    "net.zero_grad()\n",
    "for i in range(4):\n",
    "    xb, yb = X[i * 32:(i + 1) * 32], y[i * 32:(i + 1) * 32]\n",
    "    (loss(net(xb), yb) / 4).backward()   # /4 because MSE already averages\n",
    "accumulated = net.weight.grad.clone()\n",
    "print(f'max gradient difference: {(full - accumulated).abs().max():.2e}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "max gradient difference: 1.19e-07\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=memory-precision-md-8 fw=all -->\n",
    "\n",
    "The gradients match to floating-point noise. The one subtlety is the\n",
    "averaging: if your loss already averages over the batch (as MSE and\n",
    "cross-entropy do), each micro-batch gradient must be scaled by $1/k$ so\n",
    "the sum reproduces the mean. The equivalence also rests on conditions\n",
    "this tiny example satisfies silently: an additive per-example loss,\n",
    "equal-size micro-batches (or a properly weighted mean), the same dropout\n",
    "and RNG draws, no batch-coupled state such as BatchNorm — whose\n",
    "statistics would be computed per micro-batch — and clipping or optimizer\n",
    "steps applied only at the global boundary. Accumulation is the right tool\n",
    "when you are memory-bound and want a bigger *effective* batch; it is the\n",
    "wrong tool when you want to go faster: the arithmetic is identical but\n",
    "arrives as more launches over smaller, lower-intensity GEMMs, so by this\n",
    "chapter's own accounting it is usually slightly *slower* — it only\n",
    "rearranges *when* the optimizer step happens.\n",
    "\n",
    "## The Ladder So Far\n",
    "\n",
    "Four sections in, the escalation is worth naming as a checklist. When a\n",
    "model does not fit or does not run fast enough, in order of cost:\n",
    "\n",
    "1. **It's slow, bandwidth- or overhead-bound** → compile it\n",
    "   (that section). Usually free.\n",
    "2. **It's slow, compute-bound, or it's tight on memory** → drop to bf16\n",
    "   (mixed precision). Around $1.5$–$2\\times$ speed and half the\n",
    "   activation bytes.\n",
    "3. **It still doesn't fit** → checkpoint activations. Large memory saving\n",
    "   for a modest increase in time.\n",
    "4. **It still doesn't fit, and you want a bigger effective batch** →\n",
    "   accumulate gradients. Bigger effective batch at micro-batch memory —\n",
    "   about the same wall-clock, never faster.\n",
    "5. **It's still too slow, or genuinely too big for one card** → add\n",
    "   devices (that section).\n",
    "\n",
    "The next section takes that last step — and finds that adding devices is\n",
    "where the accounting gets hardest.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* A training step's memory is parameters ($4P$) + gradients ($4P$) +\n",
    "  optimizer states (Adam: $8P$) + activations (batch × depth × width, the\n",
    "  term that explodes). About $16P$ bytes before activations; peak lands\n",
    "  early in the backward pass.\n",
    "* PyTorch measures memory with runtime counters and allocation snapshots;\n",
    "  JAX reports the compiler's ahead-of-time memory *plan*. Counting versus\n",
    "  planning — the eager-versus-traced split again.\n",
    "* Mixed-precision bf16 raises tensor-core throughput (1.5–2× over a fair\n",
    "  tf32 baseline here, framework- and shape-dependent) and halves\n",
    "  activation bytes; measure it against that *fair* tf32 baseline — the\n",
    "  frameworks' fp32 defaults differ — and skip the `GradScaler` (bf16\n",
    "  shares fp32's exponent range — only fp16 needs loss scaling).\n",
    "* Activation checkpointing recomputes instead of stores: peak memory cut\n",
    "  roughly in half or better on a deep stack (about 2 GB of activations\n",
    "  released) for about 15–20% more step time here. It normally costs time; its\n",
    "  purpose is memory — use it when memory binds.\n",
    "* Gradient accumulation reaches a large *effective* batch at micro-batch\n",
    "  memory cost — $B_{\\textrm{global}} = B_{\\textrm{micro}} \\times k$ — at\n",
    "  the same wall-clock; scale per-micro-batch losses by $1/k$ when the loss\n",
    "  already averages.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Take the GPT of that section and, using only the anatomy of\n",
    "   that section, compute the largest $(\\textrm{width},\n",
    "   \\textrm{depth})$ trainable in 24 GB under three regimes: fp32 + Adam,\n",
    "   bf16 + Adam, bf16 + Adam + checkpointing (assume activations dominate).\n",
    "   Then verify one point against `max_memory_allocated`.\n",
    "1. Derive the per-token activation bytes of one `TransformerBlock`\n",
    "   (that section) as a function of width and sequence length, and\n",
    "   check your formula against the measured peak from the snapshot plot.\n",
    "1. Implement \"checkpoint every $\\sqrt{n}$-th block\" for the deep block\n",
    "   stack above and measure peak memory and step time. Why is $\\sqrt{n}$ the\n",
    "   memory-optimal checkpointing interval?\n",
    "1. Break fp16 (not bf16) deliberately: construct a step whose gradients\n",
    "   are small enough to underflow fp16's subnormal floor (e.g., scale the\n",
    "   loss by $10^{-4}$ before calling backward), run it under\n",
    "   `torch.autocast(dtype=torch.float16)` without a `GradScaler`, and\n",
    "   verify that gradients flush to zero while the same step under bf16\n",
    "   survives. Then enable a `GradScaler`, confirm the gradients return,\n",
    "   and explain what the scaler did in terms of\n",
    "   the figure. Why does a full training run fail this\n",
    "   way only *sometimes* — what makes the failure workload-dependent?\n",
    "1. With a fixed global batch of 512, sweep the micro-batch size $\\in\n",
    "   \\{512, 256, 128, 64\\}$ (so $k \\in \\{1, 2, 4, 8\\}$), and plot peak\n",
    "   memory and wall-clock per optimizer step. Confirm that memory follows\n",
    "   the micro-batch while the update is unchanged, and explain the small\n",
    "   time overhead of large $k$."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
