{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 14: Quantization"
      ],
      "id": "cf1870f1-246a-4719-a4cf-a817ede9d043"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "b6c5f053-dc24-45b2-a7ea-b00f18a74bc8"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "71478b2c-3422-4f32-9586-d8ef29fff5eb"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "You have trained a model (Modules 06–08). Every weight in it is a\n",
        "**32-bit float** — 4 bytes each. A 7-billion-parameter model is\n",
        "therefore 28 GB just to *store*, and every matmul reads all of it from\n",
        "memory. **Quantization** is the single most effective way to shrink\n",
        "that: round each weight onto a small grid of integers and keep one\n",
        "floating-point **scale** to undo the rounding at run time.\n",
        "\n",
        "Store 8-bit integers instead of 32-bit floats and the model is **4×\n",
        "smaller**; store 4-bit integers and it is **8× smaller** — often the\n",
        "difference between a model fitting on your GPU or not. The cost is a\n",
        "small, *measurable* rounding error, and this module builds the whole\n",
        "machinery from scratch so you can see exactly what that error is and how\n",
        "to keep it tiny.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "- Why a trained weight has far more precision than inference needs\n",
        "- **Symmetric (absmax)** and **affine (zero-point)** quantization, from\n",
        "  the formulas\n",
        "- How to quantize a tensor to int8 / int4 and measure the error in\n",
        "  **SNR**\n",
        "- Why **per-channel** scales beat a single per-tensor scale (the outlier\n",
        "  problem)\n",
        "- How to build a **`QuantizedLinear`** that replaces `nn.Linear` at a\n",
        "  fraction of the memory\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "- [Module 01: Tensors](../m01_tensors/lesson.qmd) — shapes,\n",
        "  broadcasting, dtypes\n",
        "- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the linear\n",
        "  layers whose weights we quantize\n",
        "- [Module 07: Training](../m07_training/lesson.qmd) — where those\n",
        "  weights come from\n",
        "\n",
        "## Intuition: A Weight Is a Number With Too Many Digits\n",
        "\n",
        "A trained weight might be `0.0417328...`. During inference, does the\n",
        "model really need all those digits? Almost never. The weights in a given\n",
        "tensor span some range — say roughly `[-0.2, 0.2]` — and if you chop\n",
        "that range into a few hundred evenly spaced buckets, rounding each\n",
        "weight to the nearest bucket barely changes the model’s output.\n",
        "\n",
        "That is the whole idea. Pick a grid of integer levels, store which\n",
        "bucket each weight fell into (a small integer), and remember **one\n",
        "scale** — the width of a bucket — to turn buckets back into real\n",
        "numbers:\n",
        "\n",
        "$$w \\approx \\text{scale} \\times q, \\qquad q \\in \\{-127, \\dots, 127\\}\\ \\text{(int8)}$$\n",
        "\n",
        "Eight bits gives $2^8 = 256$ buckets; four bits gives just $16$. Fewer\n",
        "buckets = a coarser grid = more rounding error. The art of quantization\n",
        "is choosing the grid — how wide, how centred, and how many independent\n",
        "grids to keep — so that the error stays negligible while the integers\n",
        "stay tiny.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> Quantization trades **precision for size**. A float carries ~7 decimal\n",
        "> digits of precision that a trained, redundant network mostly does not\n",
        "> use. Rounding weights onto a 256-level (int8) grid typically changes\n",
        "> outputs by well under a percent, while cutting the model’s memory —\n",
        "> and its memory-bandwidth cost, which usually dominates inference — by\n",
        "> 4×.\n",
        "\n",
        "## The Math: Symmetric and Affine Quantization\n",
        "\n",
        "**Symmetric (absmax).** The simplest scheme centres the grid on zero.\n",
        "For $b$ bits we use the signed range $[-q_{\\max}, q_{\\max}]$ with\n",
        "$q_{\\max} = 2^{b-1}-1$ (so $127$ for int8). One scale covers the whole\n",
        "tensor:\n",
        "\n",
        "$$s = \\frac{\\max_i |w_i|}{q_{\\max}}, \\qquad q_i = \\text{clamp}\\!\\left(\\text{round}\\!\\left(\\frac{w_i}{s}\\right),\\, -q_{\\max},\\, q_{\\max}\\right), \\qquad \\hat{w}_i = s\\, q_i$$\n",
        "\n",
        "Because the grid is symmetric, $0.0$ maps exactly to the integer $0$ —\n",
        "no offset needed. This is the default for **weights**, which are roughly\n",
        "zero-centred.\n",
        "\n",
        "**Affine (zero-point).** When the values are *one-sided* — think the\n",
        "output of a ReLU, all $\\geq 0$ — a symmetric grid wastes half its levels\n",
        "on negatives that never occur. Affine quantization maps the true range\n",
        "$[w_{\\min}, w_{\\max}]$ onto the *unsigned* range $[0, 2^b - 1]$ using an\n",
        "integer **zero-point** $z$ (the code that represents $0.0$):\n",
        "\n",
        "$$s = \\frac{w_{\\max} - w_{\\min}}{2^b - 1}, \\qquad z = \\text{round}\\!\\left(-\\frac{w_{\\min}}{s}\\right), \\qquad q_i = \\text{clamp}\\!\\left(\\text{round}\\!\\left(\\tfrac{w_i}{s}\\right) + z,\\, 0,\\, 2^b - 1\\right)$$\n",
        "\n",
        "and we recover $\\hat{w}_i = s\\,(q_i - z)$. No levels are wasted, so\n",
        "affine is the better choice for activations.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> Symmetric needs only a scale; affine needs a scale **and** a\n",
        "> zero-point but uses the full code range. Weights → symmetric\n",
        "> (zero-centred, and a zero-point would slow the matmul). Activations →\n",
        "> affine (often one-sided). Both are just a linear map between reals and\n",
        "> integers; the only lossy step is `round`.\n",
        "\n",
        "## Code: Quantize a Tensor from Scratch\n",
        "\n",
        "Everything lives in `quantization.py`. Load it and quantize a weight\n",
        "tensor to int8, then read back the reconstruction error."
      ],
      "id": "c7983ab5-e82f-4171-bd91-afc5529b31c7"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "original[0]:    [-0.11258398741483688, -0.11523602157831192, -0.025057857856154442, -0.04338788241147995]\n",
            "integers[0]:    [-68, -69, -15, -26]   (dtype torch.int8)\n",
            "scale:          0.001666\n",
            "dequantized[0]: [-0.11330000311136246, -0.11490000039339066, -0.02500000037252903, -0.043299999088048935]"
          ]
        }
      ],
      "source": [
        "import importlib.util\n",
        "import sys\n",
        "from pathlib import Path\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "spec = importlib.util.spec_from_file_location(\"quantization\", Path(\"quantization.py\").resolve())\n",
        "q = importlib.util.module_from_spec(spec)\n",
        "sys.modules[\"quantization\"] = q\n",
        "spec.loader.exec_module(q)\n",
        "\n",
        "torch.manual_seed(0)\n",
        "w = torch.randn(4, 4) * 0.1                       # a small weight matrix\n",
        "qi, scale = q.absmax_quantize(w, num_bits=8)\n",
        "\n",
        "print(\"original[0]:   \", w[0].tolist())\n",
        "print(\"integers[0]:   \", qi[0].tolist(), f\"  (dtype {qi.dtype})\")\n",
        "print(f\"scale:          {scale.item():.6f}\")\n",
        "print(\"dequantized[0]:\", q.dequantize(qi, scale)[0].round(decimals=4).tolist())"
      ],
      "id": "e2a334a0"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The four floats became four small integers plus one shared scale. Now\n",
        "measure how much precision that cost, in int8 versus int4:"
      ],
      "id": "6e402ca8-4bb4-44ec-a2a3-df3aba2d0ce5"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "int8: max error 0.00092   SNR  39.5 dB\n",
            "int4: max error 0.01664   SNR  14.3 dB"
          ]
        }
      ],
      "source": [
        "w = torch.randn(512, 512) * 0.05\n",
        "for bits in (8, 4):\n",
        "    qi, s = q.absmax_quantize(w, num_bits=bits)\n",
        "    err = q.quantization_error(w, q.dequantize(qi, s))\n",
        "    print(f\"int{bits}: max error {err['max_abs']:.5f}   SNR {err['snr_db']:5.1f} dB\")"
      ],
      "id": "2a7bb115"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "int8 reconstructs the tensor at ~40 dB (rounding error ~10 000× smaller\n",
        "than the signal); int4, with only 16 levels, drops to the mid-teens.\n",
        "**Signal-to-noise ratio** is the honest way to compare: every extra bit\n",
        "adds ~6 dB, so int8 is worth about 24 dB more than int4.\n",
        "\n",
        "## Per-Tensor vs Per-Channel: The Outlier Problem\n",
        "\n",
        "One scale for a whole matrix has a weakness. Neural-network weight\n",
        "matrices routinely contain a few **outlier channels** whose magnitudes\n",
        "dwarf the rest. The absmax scale is set by that single largest value, so\n",
        "every *other* channel is forced onto a needlessly coarse grid.\n",
        "\n",
        "The fix is **per-channel** quantization: give each output channel (row)\n",
        "its own scale. One extra float per row buys back the precision the\n",
        "outlier stole. Watch it on a matrix where channel 0 is 20× louder than\n",
        "the rest:"
      ],
      "id": "9f2241c2-8674-4468-a7a2-d6e7572d5e26"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "torch.manual_seed(0)\n",
        "w = torch.randn(24, 64) * 0.05\n",
        "w[0] *= 20.0                                       # one loud output channel\n",
        "\n",
        "qt, st = q.absmax_quantize(w, num_bits=8)          # one scale for everything\n",
        "qc, sc = q.quantize_per_channel(w, num_bits=8, axis=0)   # one scale per row\n",
        "\n",
        "# Per-channel mean-squared error under each scheme (for the visual below).\n",
        "mse_tensor = ((w - q.dequantize(qt, st)) ** 2).mean(dim=1)\n",
        "mse_chan = ((w - q.dequantize(qc, sc)) ** 2).mean(dim=1)\n",
        "\n",
        "ojs_define(quantChan = {\n",
        "    \"perTensor\": mse_tensor.tolist(),\n",
        "    \"perChannel\": mse_chan.tolist(),\n",
        "    \"outlier\": 0,\n",
        "})"
      ],
      "id": "4290a64f"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "per-tensor  : SNR 29.1 dB\n",
            "per-channel : SNR 42.8 dB   ← the quiet channels keep their precision"
          ]
        }
      ],
      "source": [
        "et = q.quantization_error(w, q.dequantize(qt, st))\n",
        "ec = q.quantization_error(w, q.dequantize(qc, sc))\n",
        "print(f\"per-tensor  : SNR {et['snr_db']:.1f} dB\")\n",
        "print(f\"per-channel : SNR {ec['snr_db']:.1f} dB   ← the quiet channels keep their precision\")"
      ],
      "id": "609efff3"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Per-tensor lets one loud channel drag every quiet channel’s error up by\n",
        "the same factor it is louder — here ~20×. Per-channel absorbs the\n",
        "outlier into its own row and leaves the rest untouched, recovering well\n",
        "over 10 dB of SNR. This is exactly why production weight quantization is\n",
        "**always** per-channel (and why methods like LLM.int8() go further and\n",
        "pull genuine outlier *features* out into fp16).\n",
        "\n",
        "## Code: A Quantized Linear Layer\n",
        "\n",
        "Now assemble the piece that actually saves memory in a transformer: a\n",
        "drop-in for `nn.Linear` whose weight is stored per-channel int8.\n",
        "`QuantizedLinear.from_linear` quantizes a trained layer; its `forward`\n",
        "dequantizes the weight and runs an ordinary linear — the faithful,\n",
        "readable stand-in for what an integer kernel does (integer matmul, then\n",
        "a per-channel rescale)."
      ],
      "id": "ab445f6e-7380-44b5-b4c4-a853ccd9aebb"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "output relative error: 0.3824%\n",
            "fp32 weight bytes: 1,048,576\n",
            "int8 stored bytes: 264,192  (3.97× smaller)"
          ]
        }
      ],
      "source": [
        "torch.manual_seed(0)\n",
        "linear = nn.Linear(512, 512)\n",
        "qlinear = q.QuantizedLinear.from_linear(linear, num_bits=8)\n",
        "\n",
        "x = torch.randn(8, 512)\n",
        "ref = linear(x)\n",
        "out = qlinear(x)\n",
        "\n",
        "rel = (out - ref).norm() / ref.norm()\n",
        "print(f\"output relative error: {rel.item():.4%}\")     # well under 1%\n",
        "\n",
        "mem = qlinear.memory_bytes()\n",
        "print(f\"fp32 weight bytes: {mem['fp32_weight_bytes']:,}\")\n",
        "print(f\"int8 stored bytes: {mem['weight_bytes'] + mem['scale_bytes']:,.0f}\"\n",
        "      f\"  ({mem['compression']:.2f}× smaller)\")"
      ],
      "id": "de2dd1f1"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The quantized layer’s output is within a fraction of a percent of the\n",
        "original, at ~4× less memory. Scale that up to a whole model:"
      ],
      "id": "771c23f5-9b34-4ab7-ae28-570cdff2d489"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  7B:  int32=  28.0GB  int16=  14.0GB  int8=   7.0GB  int4=   3.5GB  \n",
            " 13B:  int32=  52.0GB  int16=  26.0GB  int8=  13.0GB  int4=   6.5GB  \n",
            " 70B:  int32= 280.0GB  int16= 140.0GB  int8=  70.0GB  int4=  35.0GB  "
          ]
        }
      ],
      "source": [
        "for params, name in [(7e9, \"7B\"), (13e9, \"13B\"), (70e9, \"70B\")]:\n",
        "    line = f\"{name:>4}:  \"\n",
        "    for bits in (32, 16, 8, 4):\n",
        "        line += f\"int{bits}={q.model_footprint(int(params), bits)/1e9:6.1f}GB  \"\n",
        "    print(line)"
      ],
      "id": "3651afb1"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A 70B model is 280 GB in fp32 — many GPUs — but 35 GB in int4, within\n",
        "reach of a single high-memory card. That table is why quantization is\n",
        "not optional at the frontier.\n",
        "\n",
        "## Interactive Exploration\n",
        "\n",
        "### Step Through the Pipeline\n",
        "\n",
        "Every quantization is the same six steps: measure the range, form a\n",
        "scale, divide, round, store the integers, and dequantize on the way back\n",
        "out. Step through them and watch where the lossy `round` sits."
      ],
      "id": "4dad1a6d-2d21-4e99-a2a0-48f3d771087f"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "8ef09e99-ba67-407a-a8f8-d24b2bf84556"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "c01a2896-48b2-4b82-b09e-ef88f8f4dc0f"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "3510f0d2-6e20-4426-860b-a06ead0b2cc9"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Drive the Grid Yourself\n",
        "\n",
        "Slide the bit-width and switch schemes. Fewer bits coarsen the grid; the\n",
        "arrows show each weight snapping to its nearest level, and the readouts\n",
        "track the error. Below it, the per-channel chart (built from the tested\n",
        "code above) shows the outlier channel’s damage — and how per-channel\n",
        "scales contain it."
      ],
      "id": "4b2c2c5e-8faa-493b-ab60-5e0676155be8"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "6f81c14d-38db-4db0-a468-bb1a7c85798e"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "da73007a-7ba6-4e6f-b350-88b20d5437e9"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "d7e2941f-94af-4099-bcbe-22ca1f84836d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "96e68758-256a-4834-a53d-267bbf4cf930"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "71b4e02b-ae15-4b86-be03-a7607829a12a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Watch int4 break.** In the grid explorer, drag bits from 8 down\n",
        ">     to 2. At int2 there are only 4 levels — the snapped dots barely\n",
        ">     resemble the weights, and the SNR collapses. This is why naive\n",
        ">     int4 needs help (grouping, GPTQ).\n",
        "> 2.  **Re-centre with affine.** At low bit-widths, switch from\n",
        ">     *symmetric* to *affine*. The grid slides to hug the data’s actual\n",
        ">     range instead of straddling zero, and the error drops — the payoff\n",
        ">     of spending a zero-point.\n",
        "> 3.  **Read the per-channel chart.** The red bars (per-tensor) tower\n",
        ">     over the green (per-channel) for every quiet channel, because the\n",
        ">     outlier set their shared scale. Only the outlier channel itself is\n",
        ">     a near-tie.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "1.  **Forgetting to clamp.** After `round(w/s)` a value can land at\n",
        "    $\\pm(q_{\\max}+1)$ from floating-point error. Without `clamp` it\n",
        "    overflows the integer type and corrupts the weight. Always clamp to\n",
        "    the representable range.\n",
        "\n",
        "2.  **Symmetric quantization on skewed data.** A symmetric grid centred\n",
        "    on zero wastes half its levels on a range the data never visits\n",
        "    (e.g. post-ReLU activations). Use affine — with a zero-point —\n",
        "    whenever the values are one-sided.\n",
        "\n",
        "3.  **Per-tensor scales on weights.** One outlier channel forces a\n",
        "    coarse grid on the whole matrix. Per-channel (per-row) scales are\n",
        "    nearly free — one float per channel — and are the standard for\n",
        "    weight quantization.\n",
        "\n",
        "4.  **Expecting int4 to just work.** Four bits is 16 levels; plain\n",
        "    per-channel int4 usually loses too much. Real 4-bit uses smaller\n",
        "    **groups** (a scale per 64 or 128 weights) and error-correcting\n",
        "    methods like GPTQ/AWQ. This module builds the foundation those\n",
        "    extend.\n",
        "\n",
        "5.  **Confusing storage bits with compute.** Quantizing to int8 shrinks\n",
        "    *storage* and memory bandwidth immediately. Getting the *speedup*\n",
        "    also needs an integer matmul kernel; the `QuantizedLinear` here\n",
        "    dequantizes-then-matmuls to stay readable, which is correct but not\n",
        "    itself faster.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: int4 vs int8 error\n",
        "\n",
        "Quantize the same tensor to int8 and int4 and confirm int4’s SNR is\n",
        "roughly 24 dB lower (about 6 dB per bit)."
      ],
      "id": "487b6499-eb43-4646-ace7-359bbb412833"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "int8 39.6 dB,  int4 14.5 dB,  gap 25.2 dB"
          ]
        }
      ],
      "source": [
        "x = torch.randn(10_000) * 0.1\n",
        "q8, s8 = q.absmax_quantize(x, num_bits=8)\n",
        "q4, s4 = q.absmax_quantize(x, num_bits=4)\n",
        "snr8 = q.quantization_error(x, q.dequantize(q8, s8))[\"snr_db\"]\n",
        "snr4 = q.quantization_error(x, q.dequantize(q4, s4))[\"snr_db\"]\n",
        "print(f\"int8 {snr8:.1f} dB,  int4 {snr4:.1f} dB,  gap {snr8 - snr4:.1f} dB\")"
      ],
      "id": "fc62f928"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Build the per-channel win\n",
        "\n",
        "Take a matrix, make one channel an outlier, and show per-channel MSE is\n",
        "far below per-tensor MSE. (Hint: `quantize_per_channel(..., axis=0)` vs\n",
        "`absmax_quantize`.)"
      ],
      "id": "81f82fa7-c295-4e38-be53-205dc85a0b88"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Your implementation here"
      ],
      "id": "0374b09c"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: Footprint of a model that fits\n",
        "\n",
        "What is the largest model (in billions of parameters) that fits in 24 GB\n",
        "of GPU memory at int4, leaving 4 GB for activations? Use\n",
        "`model_footprint`."
      ],
      "id": "ad256619-3977-4acc-8c13-6288382fa0f6"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "~40B parameters fit in int4 within a 20 GB weight budget"
          ]
        }
      ],
      "source": [
        "# budget = 20 GB for weights; solve params from model_footprint(params, 4)\n",
        "budget_bytes = 20e9\n",
        "params = budget_bytes / q.bytes_per_param(4)\n",
        "print(f\"~{params/1e9:.0f}B parameters fit in int4 within a 20 GB weight budget\")"
      ],
      "id": "9d63974e"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **Quantization stores integers, not floats.** Round each weight onto\n",
        "    a grid of $2^b$ levels and keep a scale to undo it:\n",
        "    $w \\approx s\\,q$. int8 is 4× smaller than fp32, int4 is 8×.\n",
        "\n",
        "2.  **Two schemes.** *Symmetric* (absmax) centres the grid on zero —\n",
        "    best for weights. *Affine* adds a zero-point to use the full code\n",
        "    range — best for one-sided activations.\n",
        "\n",
        "3.  **SNR is the honest metric.** Every bit is worth ~6 dB; int8\n",
        "    reconstructs weights near 40 dB, int4 near the mid-teens.\n",
        "\n",
        "4.  **Per-channel beats per-tensor.** A single outlier channel wrecks a\n",
        "    per-tensor scale; one scale per row costs almost nothing and\n",
        "    recovers the lost precision.\n",
        "\n",
        "5.  **`QuantizedLinear` is the payoff.** Stored per-channel int8, it\n",
        "    matches the fp32 layer to a fraction of a percent at ~4× less memory\n",
        "    — the building block of every quantized transformer.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "Quantization is the first lever of **fast inference**. From here the\n",
        "frontier adds **speculative decoding** (a small draft model proposes,\n",
        "the big model verifies), **paged attention** (vLLM’s KV-cache memory\n",
        "manager), and lower-bit methods (**GPTQ**, **AWQ**) that push past int4\n",
        "with grouping and error correction — each building on the integer\n",
        "round-trip you just implemented.\n",
        "\n",
        "## Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- Jacob et al., *Quantization and Training of Neural Networks for\n",
        "  Efficient Integer-Arithmetic-Only Inference* (2017) — the\n",
        "  affine/zero-point scheme. <https://arxiv.org/abs/1712.05877>\n",
        "- Dettmers et al., *LLM.int8(): 8-bit Matrix Multiplication for\n",
        "  Transformers at Scale* (2022) — outlier features and why plain int8\n",
        "  breaks at scale. <https://arxiv.org/abs/2208.07339>\n",
        "- Frantar et al., *GPTQ: Accurate Post-Training Quantization for\n",
        "  Generative Pre-trained Transformers* (2022) — accurate one-shot 4-bit\n",
        "  weight quantization. <https://arxiv.org/abs/2210.17323>\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- Lin et al., *AWQ: Activation-aware Weight Quantization* (2023) —\n",
        "  protecting the salient weights. <https://arxiv.org/abs/2306.00978>"
      ],
      "id": "3b6e1a81-db63-4413-aa21-27ec106cb967"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3"
    },
    "language_info": {
      "name": "python",
      "codemirror_mode": {
        "name": "ipython",
        "version": "3"
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.14"
    }
  }
}