{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What rounding a weight costs\n",
    "\n",
    "The affine map, the error it leaves, and why one outlier ruins a whole row.\n",
    "\n",
    "From chapter 24, [Efficiency: Smaller, Cheaper, Faster](https://nlp.jcrlabz.com/book/efficient/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/quantization.py`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "Nothing. This example uses only the Python standard library.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nothing to install: this example uses only the Python standard library.\n",
    "import sys; print(sys.version.split()[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The script\n",
    "\n",
    "Everything the example defines, in one cell. Run it and the definitions\n",
    "are live in the notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Quantisation, worked as Chapter 23 works it.\n",
    "\n",
    "Four things, all reproducing the chapter's tables:\n",
    "\n",
    "    the memory table   what each precision costs for a 7B model\n",
    "    one tensor         the affine map, rounded and dequantised, with its error\n",
    "    the outlier        why one large weight ruins the whole tensor\n",
    "    per-channel        the fix, and what it costs to store\n",
    "\n",
    "    python3 quantization.py             # all four\n",
    "    python3 quantization.py --memory    # bytes per parameter\n",
    "    python3 quantization.py --round     # int8 and int4 on eight weights\n",
    "    python3 quantization.py --outlier   # one big weight, and the damage\n",
    "    python3 quantization.py --channel   # per-channel scales\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "\n",
    "# Eight weights from one row of a hidden layer. Small, spread either side of\n",
    "# zero, and deliberately not symmetric: this is what most of a trained model\n",
    "# looks like.\n",
    "WEIGHTS = [0.12, -0.45, 0.33, -0.08, 0.51, -0.22, 0.07, -0.37]\n",
    "\n",
    "# The same row after one outlier arrives. Large activations and weights in a\n",
    "# few coordinates are a measured property of trained transformers, not an\n",
    "# invention for this example.\n",
    "OUTLIER = WEIGHTS + [4.20]\n",
    "\n",
    "PRECISIONS = [(\"fp32\", 4.0), (\"fp16\", 2.0), (\"int8\", 1.0), (\"int4\", 0.5)]\n",
    "PARAMS_7B = 7_000_000_000\n",
    "\n",
    "\n",
    "def show_memory():\n",
    "    print(\"\\nSTEP 1  what a precision costs\\n\")\n",
    "    print(\"  precision   bytes/param      a 7B model     versus fp16\")\n",
    "    print(\"  \" + \"-\" * 56)\n",
    "    fp16 = PARAMS_7B * 2.0 / 1e9\n",
    "    for name, b in PRECISIONS:\n",
    "        gb = PARAMS_7B * b / 1e9\n",
    "        print(f\"  {name:9}   {b:11.1f}   {gb:9.1f} GB   {gb / fp16:9.2f}x\")\n",
    "    print(\"  \" + \"-\" * 56)\n",
    "    print(\"\\n  Read the int4 row. A model that needed a data-centre card at fp16\")\n",
    "    print(\"  fits on a consumer GPU at 3.5 GB, and nothing was retrained.\")\n",
    "    print(\"  Weights are stored in 16 bits and rarely need 16 bits to be used.\")\n",
    "\n",
    "\n",
    "def quantise(ws, bits):\n",
    "    \"\"\"Symmetric affine quantisation. Returns scale, codes, dequantised values.\"\"\"\n",
    "    qmax = 2 ** (bits - 1) - 1              # 127 for int8, 7 for int4\n",
    "    scale = max(abs(w) for w in ws) / qmax\n",
    "    codes = [max(-qmax - 1, min(qmax, round(w / scale))) for w in ws]\n",
    "    deq = [c * scale for c in codes]\n",
    "    return scale, codes, deq\n",
    "\n",
    "\n",
    "def errors(ws, deq):\n",
    "    errs = [abs(w - d) for w, d in zip(ws, deq)]\n",
    "    rms = (sum(e * e for e in errs) / len(errs)) ** 0.5\n",
    "    return max(errs), rms\n",
    "\n",
    "\n",
    "def show_round(ws=WEIGHTS, title=\"STEP 2  one tensor, rounded and put back\"):\n",
    "    print(f\"\\n{title}\\n\")\n",
    "    print(\"  A weight is stored as an integer code and a shared scale:\\n\")\n",
    "    print(\"      q = round(w / S),      S = max|w| / q_max,      w_hat = q * S\\n\")\n",
    "    for bits in (8, 4):\n",
    "        scale, codes, deq = quantise(ws, bits)\n",
    "        qmax = 2 ** (bits - 1) - 1\n",
    "        mx, rms = errors(ws, deq)\n",
    "        print(f\"  int{bits}:  q_max = {qmax},  S = max|w| / {qmax} = {scale:.6f}\\n\")\n",
    "        print(\"        w        w/S      q     w_hat      error\")\n",
    "        print(\"     \" + \"-\" * 44)\n",
    "        for w, c, d in zip(ws, codes, deq):\n",
    "            print(f\"   {w:+7.2f}  {w / scale:+9.2f}  {c:+5d}  {d:+8.4f}   {abs(w - d):8.5f}\")\n",
    "        print(\"     \" + \"-\" * 44)\n",
    "        print(f\"     max error {mx:.5f}   rms error {rms:.5f}\\n\")\n",
    "    s8, _, d8 = quantise(ws, 8)\n",
    "    s4, _, d4 = quantise(ws, 4)\n",
    "    _, r8 = errors(ws, d8)\n",
    "    _, r4 = errors(ws, d4)\n",
    "    print(f\"  Halving the bit width multiplies the rms error by {r4 / r8:.1f}, because\")\n",
    "    print(f\"  the step between neighbouring codes grows from {s8:.6f} to {s4:.6f}.\")\n",
    "\n",
    "\n",
    "def show_outlier():\n",
    "    print(\"\\nSTEP 3  one large weight ruins the row\\n\")\n",
    "    s_clean, _, _ = quantise(WEIGHTS, 8)\n",
    "    s_out, codes, deq = quantise(OUTLIER, 8)\n",
    "    print(\"  Add a single outlier of +4.20 to the same eight weights. The scale\")\n",
    "    print(\"  is set by the largest value, so every other weight now shares a\")\n",
    "    print(f\"  step {s_out / s_clean:.1f} times coarser.\\n\")\n",
    "    print(f\"  int8 scale without the outlier  S = {s_clean:.6f}\")\n",
    "    print(f\"  int8 scale with it              S = {s_out:.6f}   ({s_out / s_clean:.1f}x coarser)\\n\")\n",
    "    print(\"        w      q     w_hat      error    error %\")\n",
    "    print(\"     \" + \"-\" * 46)\n",
    "    for w, c, d in zip(OUTLIER, codes, deq):\n",
    "        pct = 100 * abs(w - d) / abs(w) if w else 0.0\n",
    "        flag = \"   <- outlier\" if abs(w) > 1 else \"\"\n",
    "        print(f\"   {w:+7.2f}  {c:+5d}  {d:+8.4f}   {abs(w - d):8.5f}  {pct:7.2f}%{flag}\")\n",
    "    print(\"     \" + \"-\" * 46)\n",
    "    small = [(w, d) for w, d in zip(OUTLIER, deq) if abs(w) <= 1]\n",
    "    worst = max(100 * abs(w - d) / abs(w) for w, d in small)\n",
    "    print(f\"\\n  The outlier itself is fine, at {100 * abs(4.20 - deq[-1]) / 4.20:.2f} per cent. The small weights\")\n",
    "    print(f\"  now carry up to {worst:.2f} per cent error, and there are far more of them.\")\n",
    "    print(\"  This is why naive per-tensor quantisation of a transformer fails.\")\n",
    "\n",
    "\n",
    "def show_channel():\n",
    "    print(\"\\nSTEP 4  give each group its own scale\\n\")\n",
    "    print(\"  Split the row into groups and quantise each with its own scale, so\")\n",
    "    print(\"  one extreme value cannot set the step for everything else.\\n\")\n",
    "    groups = [(\"the eight normal weights\", WEIGHTS), (\"the outlier alone\", [4.20])]\n",
    "    _, _, deq_all = quantise(OUTLIER, 8)\n",
    "    _, rms_all = errors(OUTLIER, deq_all)\n",
    "    rms_parts, n = 0.0, 0\n",
    "    for name, g in groups:\n",
    "        s, _, d = quantise(g, 8)\n",
    "        mx, rms = errors(g, d)\n",
    "        rms_parts += sum((a - b) ** 2 for a, b in zip(g, d))\n",
    "        n += len(g)\n",
    "        print(f\"  {name:26} S = {s:.6f}   max error {mx:.5f}\")\n",
    "    rms_grouped = (rms_parts / n) ** 0.5\n",
    "    print(f\"\\n  rms error, one scale for the row   {rms_all:.5f}\")\n",
    "    print(f\"  rms error, two scales              {rms_grouped:.5f}\")\n",
    "    print(f\"  improvement                        {rms_all / rms_grouped:.0f}x\\n\")\n",
    "    print(\"  The cost is the scales themselves. One fp16 scale per group of 64\")\n",
    "    print(\"  weights adds 2 bytes per 64, which is 0.25 bits per weight, so int4\")\n",
    "    print(\"  with group scales really costs 4.25 bits. That is the trade every\")\n",
    "    print(\"  quantisation scheme is negotiating.\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"--memory\", action=\"store_true\")\n",
    "    ap.add_argument(\"--round\", action=\"store_true\")\n",
    "    ap.add_argument(\"--outlier\", action=\"store_true\")\n",
    "    ap.add_argument(\"--channel\", action=\"store_true\")\n",
    "    a = ap.parse_args()\n",
    "    run_all = not (a.memory or a.round or a.outlier or a.channel)\n",
    "    if a.memory or run_all:\n",
    "        show_memory()\n",
    "    if a.round or run_all:\n",
    "        show_round()\n",
    "    if a.outlier or run_all:\n",
    "        show_outlier()\n",
    "    if a.channel or run_all:\n",
    "        show_channel()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running it\n",
    "\n",
    "The script is a command line tool, so each run below sets the arguments\n",
    "and calls `main()`. The heading shows the equivalent shell command.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "def run(*args):\n",
    "    sys.argv = [\"quantization.py\", *args]\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 quantization.py`\n",
    "\n",
    "Memory, rounding, the outlier and the repair.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 quantization.py --round`\n",
    "\n",
    "The same eight weights at int8 and int4. Halving the bit width multiplies the rms error by 20.6, because the step grows from 0.004016 to 0.072857.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--round\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 quantization.py --outlier`\n",
    "\n",
    "One weight of +4.20 makes the scale 8.2 times coarser. The outlier is then stored exactly while the small weights carry up to 17.32 per cent error.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--outlier\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 quantization.py --channel`\n",
    "\n",
    "Two scales instead of one cut the rms error thirteenfold, at a cost of a quarter of a bit per weight, so int4 with group scales really costs 4.25 bits.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--channel\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Efficiency: Smaller, Cheaper, Faster](https://nlp.jcrlabz.com/book/efficient/).\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "quantization.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}