{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b81b2e2",
   "metadata": {},
   "source": [
    "# Quantization\n",
    "\n",
    "Lecture 19 | CMU ANLP Fall 2025 | Instructor: Sean Welleck\n",
    "\n",
    "This notebook shows basic quantization concepts.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355f9204",
   "metadata": {},
   "source": [
    "#### Absmax quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "03f8f5d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original array:     [-3.  1.  2.  4.]\n",
      "Quantized array:    [-95  32  64 127]\n",
      "Dequantized array:  [-2.99212598  1.00787402  2.01574803  4.        ]\n",
      "Mean Squared Error: 9.300018600037166e-05\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "x = np.array([-3.0, 1.0, 2.0, 4.0])\n",
    "\n",
    "if np.max(np.abs(x)) != 0:\n",
    "    scale = 127.0 / np.max(np.abs(x))\n",
    "else:\n",
    "    scale = 1.0\n",
    "\n",
    "x_quantized = np.round(x * scale).astype(np.int8)\n",
    "x_dequantized = x_quantized.astype(np.float32) / scale\n",
    "\n",
    "print(\"Original array:    \", x)\n",
    "print(\"Quantized array:   \", x_quantized)\n",
    "print(\"Dequantized array: \", x_dequantized)\n",
    "\n",
    "mse = np.mean((x - x_dequantized) ** 2)\n",
    "print(\"Mean Squared Error:\", mse)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e658dfcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def absmax_quantize(x):\n",
    "    if np.max(np.abs(x)) != 0:\n",
    "        scale = 127.0 / np.max(np.abs(x))\n",
    "    else:\n",
    "        scale = 1.0\n",
    "\n",
    "    x_quantized = np.round(x * scale).astype(np.int8)\n",
    "    x_dequantized = x_quantized.astype(np.float32) / scale\n",
    "\n",
    "    return x_quantized, x_dequantized\n",
    "\n",
    "def analyze(x, x_quantized, x_dequantized):\n",
    "    print(\"Original array:    \", x)\n",
    "    print(\"Quantized array:   \", x_quantized)\n",
    "    print(\"Dequantized array: \", x_dequantized)\n",
    "    mse = np.mean((x - x_dequantized) ** 2)\n",
    "    print(f\"Mean Squared Error: {mse:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "274ebcb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original array:     [-3.112  1.567  2.789  4.345]\n",
      "Quantized array:    [-91  46  82 127]\n",
      "Dequantized array:  [-3.11334646  1.57377953  2.80543307  4.345     ]\n",
      "Mean Squared Error: 0.000079\n"
     ]
    }
   ],
   "source": [
    "x = np.array([-3.112, 1.567, 2.789, 4.345])\n",
    "x_quantized, x_dequantized = absmax_quantize(x)\n",
    "analyze(x, x_quantized, x_dequantized)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e90bbdd",
   "metadata": {},
   "source": [
    "### Zero-point quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9e95b8cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original array:     [-3.  1.  2.  4.]\n",
      "Quantized array:    [-128   17   54  127]\n",
      "Dequantized array:  [-2.99215686  0.98823529  2.00392157  4.00784314]\n",
      "Scale 36.42857142857143\n",
      "Zero-point -19\n",
      "Mean Squared Error: 6.92041522491351e-05\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "x = np.array([-3.0, 1.0, 2.0, 4.0])\n",
    "x_min = np.min(x)\n",
    "x_max = np.max(x)\n",
    "\n",
    "qmin, qmax = -128, 127\n",
    "\n",
    "if x_max - x_min != 0:\n",
    "    scale = (qmax - qmin) / (x_max - x_min)\n",
    "else:\n",
    "    scale = 1.0\n",
    "\n",
    "zero_point = np.round(qmin - x_min * scale)\n",
    "zero_point = np.clip(zero_point, qmin, qmax).astype(np.int8)\n",
    "\n",
    "x_quantized = np.round(x * scale + zero_point)\n",
    "x_quantized = np.clip(x_quantized, qmin, qmax).astype(np.int8)\n",
    "\n",
    "x_dequantized = (x_quantized.astype(np.float32) - zero_point) / scale\n",
    "\n",
    "print(\"Original array:    \", x)\n",
    "print(\"Quantized array:   \", x_quantized)\n",
    "print(\"Dequantized array: \", x_dequantized)\n",
    "print(\"Scale\", scale)\n",
    "print(\"Zero-point\", zero_point)\n",
    "mse = np.mean((x - x_dequantized) ** 2)\n",
    "print(\"Mean Squared Error:\", mse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b2a05999",
   "metadata": {},
   "outputs": [],
   "source": [
    "def zero_point_quantize(x):\n",
    "    x_min = np.min(x)\n",
    "    x_max = np.max(x)\n",
    "\n",
    "    qmin, qmax = -128, 127\n",
    "\n",
    "    if x_max - x_min != 0:\n",
    "        scale = (qmax - qmin) / (x_max - x_min)\n",
    "    else:\n",
    "        scale = 1.0\n",
    "        \n",
    "    zero_point = np.round(qmin - x_min * scale)\n",
    "    zero_point = np.clip(zero_point, qmin, qmax).astype(np.int8)\n",
    "    x_quantized = np.round(x * scale + zero_point)\n",
    "    x_quantized = np.clip(x_quantized, qmin, qmax).astype(np.int8)\n",
    "    x_dequantized = (x_quantized.astype(np.float32) - zero_point) / scale\n",
    "    return x_quantized, x_dequantized"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "201d39fc",
   "metadata": {},
   "source": [
    "### Outlier example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "910e4659",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "With outlier:\n",
      "\n",
      "Absmax Quantization:\n",
      "Original array:     [ -0.3   0.1   0.2   0.4  -0.3   0.1   0.2   0.4  -0.3   0.1   0.2 100. ]\n",
      "Quantized array:    [  0   0   0   1   0   0   0   1   0   0   0 127]\n",
      "Dequantized array:  [  0.           0.           0.           0.78740157   0.\n",
      "   0.           0.           0.78740157   0.           0.\n",
      "   0.         100.        ]\n",
      "Mean Squared Error: 0.060013\n",
      "\n",
      "Zero-point Quantization:\n",
      "Original array:     [ -0.3   0.1   0.2   0.4  -0.3   0.1   0.2   0.4  -0.3   0.1   0.2 100. ]\n",
      "Quantized array:    [-128 -127 -126 -126 -128 -127 -126 -126 -128 -127 -126  127]\n",
      "Dequantized array:  [-0.39333333  0.          0.39333333  0.39333333 -0.39333333  0.\n",
      "  0.39333333  0.39333333 -0.39333333  0.          0.39333333 99.90666667]\n",
      "Mean Squared Error: 0.014756\n",
      "\n",
      "Without outlier:\n",
      "\n",
      "Absmax Quantization:\n",
      "Original array:     [-0.3  0.1  0.2  0.4 -0.3  0.1  0.2  0.4 -0.3  0.1  0.2  0.4]\n",
      "Quantized array:    [-95  32  64 127 -95  32  64 127 -95  32  64 127]\n",
      "Dequantized array:  [-0.2992126  0.1007874  0.2015748  0.4       -0.2992126  0.1007874\n",
      "  0.2015748  0.4       -0.2992126  0.1007874  0.2015748  0.4      ]\n",
      "Mean Squared Error: 0.000001\n",
      "\n",
      "Zero-point Quantization:\n",
      "Original array:     [-0.3  0.1  0.2  0.4 -0.3  0.1  0.2  0.4 -0.3  0.1  0.2  0.4]\n",
      "Quantized array:    [-128   17   54  127 -128   17   54  127 -128   17   54  127]\n",
      "Dequantized array:  [-0.29921569  0.09882353  0.20039216  0.40078431 -0.29921569  0.09882353\n",
      "  0.20039216  0.40078431 -0.29921569  0.09882353  0.20039216  0.40078431]\n",
      "Mean Squared Error: 0.000001\n"
     ]
    }
   ],
   "source": [
    "x_outlier = np.array([-0.3, 0.1, 0.2, 0.4, -0.3, 0.1, 0.2, 0.4, -0.3, 0.1, 0.2, 100.0])\n",
    "x_non_outlier = np.array([-0.3, 0.1, 0.2, 0.4, -0.3, 0.1, 0.2, 0.4, -0.3, 0.1, 0.2, 0.4])\n",
    "\n",
    "# Run each kind of quantization on outliers and non-outliers\n",
    "for x in [x_outlier, x_non_outlier]:\n",
    "    print(\"\\nWith outlier:\" if x is x_outlier else \"\\nWithout outlier:\")\n",
    "\n",
    "    print(\"\\nAbsmax Quantization:\")\n",
    "    x_quantized, x_dequantized = absmax_quantize(x)\n",
    "    analyze(x, x_quantized, x_dequantized)\n",
    "\n",
    "    print(\"\\nZero-point Quantization:\")\n",
    "    x_quantized, x_dequantized = zero_point_quantize(x)\n",
    "    analyze(x, x_quantized, x_dequantized)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ebd6f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# return the scale, which will be useful for the matrix multiplication example\n",
    "def absmax_quantize(x):\n",
    "    if np.max(np.abs(x)) != 0:\n",
    "        scale = 127.0 / np.max(np.abs(x))\n",
    "    else:\n",
    "        scale = 1.0\n",
    "\n",
    "    x_quantized = np.round(x * scale).astype(np.int8)\n",
    "    return x_quantized, scale\n",
    "\n",
    "def dequantize(x_quantized, scale):\n",
    "    x_dequantized = x_quantized.astype(np.float32) / scale\n",
    "    return x_dequantized"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a00dea7b",
   "metadata": {},
   "source": [
    "### Matrix multiplication example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a110d30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W (Original):\n",
      " [[ 1.05     0.746    0.05606]\n",
      " [ 0.5435  -0.7725   0.8413 ]\n",
      " [ 1.889    0.2646   2.355  ]\n",
      " [-0.1483  -1.635    0.542  ]]\n",
      "X (Original):\n",
      " [[-2.084e+00 -2.498e+01 -1.690e+00 -7.336e-02]\n",
      " [-1.869e-01 -3.581e+01 -1.252e+00 -5.234e-01]\n",
      " [-4.277e-01  2.052e+01 -6.880e-01  3.694e-01]\n",
      " [ 1.839e+00  1.438e+02  1.112e+00 -1.826e+00]\n",
      " [ 1.476e+00 -2.453e+01  1.535e+00 -6.396e-01]\n",
      " [-1.732e+00 -6.819e+01  1.178e+00 -2.742e-01]\n",
      " [ 1.381e+00 -6.009e+01  3.245e-01  9.492e-01]\n",
      " [ 1.128e+00 -7.227e-01  1.777e+00 -1.615e+00]]\n",
      "W_quantized:\n",
      " [[ 57  40   3]\n",
      " [ 29 -42  45]\n",
      " [102  14 127]\n",
      " [ -8 -88  29]]\n",
      "X_quantized:\n",
      " [[ -2 -22  -1   0]\n",
      " [  0 -32  -1   0]\n",
      " [  0  18  -1   0]\n",
      " [  2 127   1  -2]\n",
      " [  1 -22   1  -1]\n",
      " [ -2 -60   1   0]\n",
      " [  1 -53   0   1]\n",
      " [  1  -1   2  -1]]\n",
      "Y (Quantized MatMul Result before scaling):\n",
      " [[ -854   830 -1123]\n",
      " [-1030  1330 -1567]\n",
      " [  420  -770   683]\n",
      " [ 3915 -5064  5790]\n",
      " [ -471  1066  -889]\n",
      " [-1752  2454 -2579]\n",
      " [-1488  2178 -2353]\n",
      " [  240   198   183]]\n",
      "Dequantized MatMul Result:\n",
      " [[ -17.931759    17.427822   -23.580053 ]\n",
      " [ -21.627296    27.92651    -32.902885 ]\n",
      " [   8.818897   -16.167978    14.3412075]\n",
      " [  82.20473   -106.33071    121.57481  ]\n",
      " [  -9.889764    22.383202   -18.666666 ]\n",
      " [ -36.787403    51.527557   -54.15223  ]\n",
      " [ -31.244095    45.732285   -49.406826 ]\n",
      " [   5.03937      4.1574802    3.8425198]]\n",
      "Float16 MatMul Result:\n",
      " [[ -18.95    17.42   -25.16 ]\n",
      " [ -21.95    28.05   -33.38 ]\n",
      " [   9.34   -16.95    15.81 ]\n",
      " [  82.44  -106.4    122.7  ]\n",
      " [  -8.79    21.5    -17.28 ]\n",
      " [ -36.62    52.12   -54.84 ]\n",
      " [ -30.73    46.     -49.2  ]\n",
      " [   4.387    4.51     2.766]]\n",
      "Mean Squared Error between float16 and quantized matmul: 0.6285514\n"
     ]
    }
   ],
   "source": [
    "# example for matrix multiplication \n",
    "\n",
    "outliers = True\n",
    "\n",
    "h = 4\n",
    "o = 3\n",
    "T = 8\n",
    "\n",
    "W = np.random.randn(h, o).astype(np.float16)\n",
    "X = np.random.randn(T, h).astype(np.float16)\n",
    "\n",
    "if outliers:\n",
    "    # Introduce outlier features in X\n",
    "    X[:, 1] *= 100.0\n",
    "\n",
    "W_quantized, W_scale = absmax_quantize(W)\n",
    "X_quantized, X_scale = absmax_quantize(X)\n",
    "\n",
    "matmul_scale = W_scale * X_scale\n",
    "\n",
    "# In reality, we would use a int8xint8 -> int32 kernel\n",
    "Y = X_quantized.astype(np.int32).dot(W_quantized.astype(np.int32))\n",
    "Y_dequantized = Y.astype(np.float32) / matmul_scale\n",
    "\n",
    "print(\"W (Original):\\n\", W)\n",
    "print(\"X (Original):\\n\", X)\n",
    "print(\"W_quantized:\\n\", W_quantized)\n",
    "print(\"X_quantized:\\n\", X_quantized)\n",
    "print(\"Y (Quantized MatMul Result before scaling):\\n\", Y)\n",
    "\n",
    "print(\"Dequantized MatMul Result:\\n\", Y_dequantized)\n",
    "# Compare with float16 matmul\n",
    "Y_float16 = np.matmul(X.astype(np.float16), W.astype(np.float16))\n",
    "print(\"Float16 MatMul Result:\\n\", Y_float16)\n",
    "mse = np.mean((Y_float16 - Y_dequantized) ** 2)\n",
    "print(\"Mean Squared Error between float16 and quantized matmul:\", mse)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f2cc501",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anlp",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}