Demos/15. 探究非对称量化和对称量化 - INT8.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19e284d2-74ad-4b0c-86d4-8279fcd69e8e",
   "metadata": {},
   "source": [
    "# 探究非对称量化和对称量化 - INT8\n",
    "\n",
    "> 指导文章：[17. 浅谈模型量化: 非对称 vs 对称](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/17.%20浅谈模型量化：非对称%20vs%20对称.md)\n",
    "\n",
    "在线链接：[Kaggle](https://www.kaggle.com/code/aidemos/15-int8) | [Colab](https://colab.research.google.com/drive/1aFlUL8jQZEAO2ZsMXJxbUqlIXGb0m0hC?usp=sharing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a8ce26a-af40-4356-9752-e8440b6973ae",
   "metadata": {},
   "source": [
    "## 安装库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2059734-fbe3-412d-8b73-2d94a96af741",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install bitsandbytes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5e0797e-dd3e-4df3-82df-5eeaeeb084f9",
   "metadata": {},
   "source": [
    "## 应用 INT8 量化并打印模型参数\n",
    "\n",
    "观察 `transformer.h.0` 的模块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f450da-f688-423c-abaa-b76b1b2188b9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, BitsAndBytesConfig\n",
    "import torch.nn as nn\n",
    "\n",
    "model_name = 'gpt2-large'\n",
    "\n",
    "# 将模型配置为以 8-bit 量化的方式加载\n",
    "bnb_config = BitsAndBytesConfig(load_in_8bit=True)\n",
    "\n",
    "# 加载模型，并根据设备自动分配\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map='auto'\n",
    ")\n",
    "\n",
    "# 打印模型的架构\n",
    "print(model)\n",
    "\n",
    "# 打印模型的参数名称和数据类型\n",
    "print(f\"\\n打印模型 {model_name} 的参数信息:\\n\")\n",
    "for name, param in model.named_parameters():\n",
    "    print(f\"参数名称: {name}, 数据类型: {param.dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "936ada0a-9a06-44a1-a8eb-17c17d053861",
   "metadata": {},
   "source": [
    "## 代码示例\n",
    "\n",
    "尝试运行三种不同的 `fp32_values`。\n",
    "```python\n",
    "# 0\n",
    "fp32_values = torch.tensor([3.0, -5.5, 0.0, 6.0, -6.0, 2.5], dtype=torch.float32)\n",
    "# 1\n",
    "fp32_values = torch.tensor([3.0, -5.5, 0.0, 4.0, -6.0, 2.5], dtype=torch.float32)\n",
    "# 2\n",
    "fp32_values = torch.tensor([3.0, -5.5, 0.0, 8.0, -6.0, 2.5], dtype=torch.float32)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ffeba75-da11-494d-98f5-268dde4df385",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.quantization\n",
    "\n",
    "# 假设 FP32 的张量\n",
    "fp32_values = torch.tensor([3.0, -5.5, 0.0, 6.0, -6.0, 2.5], dtype=torch.float32)\n",
    "print(f\"FP32 的示例张量: {fp32_values}\\n\")\n",
    "\n",
    "# 定义 PyTorch 的量化和反量化函数\n",
    "def pytorch_quantize(fp32_tensor):\n",
    "    # 使用 min-max 范围计算缩放因子，指定 dtype 为 torch.qint8\n",
    "    q_params = torch.quantization.MinMaxObserver(dtype=torch.qint8)\n",
    "    q_params(fp32_tensor)\n",
    "    scale, zero_point = q_params.calculate_qparams()\n",
    "\n",
    "    # 量化\n",
    "    int8_tensor = torch.quantize_per_tensor(fp32_tensor, scale.item(), zero_point.item(), dtype=torch.qint8)\n",
    "    return int8_tensor, scale.item(), zero_point\n",
    "\n",
    "def pytorch_dequantize(int8_tensor):\n",
    "    # 反量化\n",
    "    fp32_tensor = int8_tensor.dequantize()\n",
    "    return fp32_tensor\n",
    "\n",
    "# 量化并获取 PyTorch 结果\n",
    "int8_tensor, scale, zero_point = pytorch_quantize(fp32_values)\n",
    "print(\"PyTorch 量化后的 INT8 数值：\", int8_tensor.int_repr())\n",
    "print(\"PyTorch 使用的 scale:\", scale, \"zero_point:\", zero_point)\n",
    "\n",
    "# 反量化\n",
    "recovered_fp32_pytorch = pytorch_dequantize(int8_tensor)\n",
    "print(\"PyTorch 反量化恢复后的 FP32 数值：\", recovered_fp32_pytorch)\n",
    "\n",
    "print(\"\\n=====================\\n\")\n",
    "\n",
    "# 对比与自定义的量化方式\n",
    "def custom_quantize_compare(fp32_values):\n",
    "    # 获取张量数值的最小值和最大值\n",
    "    x_min, x_max = fp32_values.min().item(), fp32_values.max().item()\n",
    "    \n",
    "    # 定义量化后整数数值的范围\n",
    "    qmin, qmax = -128, 127  # 对应 torch.qint8\n",
    "    \n",
    "    # 计算 scale\n",
    "    scale_custom = (x_max - x_min) / (qmax - qmin)\n",
    "    \n",
    "    # 非对称量化\n",
    "    initial_zero_point = qmin - x_min / scale_custom\n",
    "    zero_point_custom = int(round(initial_zero_point))\n",
    "    \n",
    "    # 将 zero_point 限制在 [qmin, qmax] 范围内\n",
    "    zero_point_custom = max(qmin, min(qmax, zero_point_custom))\n",
    "    \n",
    "    print(\"自定义计算的 scale:\", scale_custom, \"zero_point:\", zero_point_custom)\n",
    "    \n",
    "    def quantize(fp32_tensor, scale, zero_point):\n",
    "        # 计算量化值\n",
    "        int8_tensor = torch.round(fp32_tensor / scale) + zero_point\n",
    "        # 限制在 [qmin, qmax] 范围内\n",
    "        int8_tensor = torch.clamp(int8_tensor, qmin, qmax)\n",
    "        return int8_tensor.to(torch.int8)\n",
    "    \n",
    "    def dequantize(int8_tensor, scale, zero_point):\n",
    "        # 反量化\n",
    "        fp32_tensor = (int8_tensor.float() - zero_point) * scale\n",
    "        return fp32_tensor\n",
    "    \n",
    "    # 量化\n",
    "    int8_values_custom = quantize(fp32_values, scale_custom, zero_point_custom)\n",
    "    print(\"自定义对称量化后的 INT8 数值：\", int8_values_custom)\n",
    "    \n",
    "    # 反量化\n",
    "    recovered_fp32_custom = dequantize(int8_values_custom, scale_custom, zero_point_custom)\n",
    "    print(\"自定义对称反量化恢复后的 FP32 数值：\", recovered_fp32_custom)\n",
    "\n",
    "# 运行自定义量化并比较\n",
    "custom_quantize_compare(fp32_values)\n",
    "\n",
    "print(\"\\n=====================\\n\")\n",
    "\n",
    "# 使用 fp32_values 作为线性层参数\n",
    "\n",
    "# 定义一个简单的线性模型\n",
    "class SimpleLinearModel(nn.Module):\n",
    "    def __init__(self, weights, bias=None):\n",
    "        super(SimpleLinearModel, self).__init__()\n",
    "        # 假设输入特征数为6，输出特征数为1，用于匹配之前定义的张量，你也可以试试 (1, 6)，记得对应修改 weights.view(6, 1)\n",
    "        self.linear = nn.Linear(6, 1, bias=False)\n",
    "        # 初始化权重\n",
    "        self.linear.weight = nn.Parameter(weights.view(1, 6))  # 权重形状为 [out_features, in_features]\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "# 创建 FP32 模型\n",
    "fp32_weights = fp32_values  # [6]\n",
    "model_fp32 = SimpleLinearModel(fp32_weights)\n",
    "\n",
    "# 打印 FP32 模型的权重\n",
    "print(\"FP32 模型的权重:\\n\", model_fp32.linear.weight)\n",
    "\n",
    "# 使用默认量化配置\n",
    "model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')\n",
    "\n",
    "# 准备量化\n",
    "torch.quantization.prepare(model_fp32, inplace=True)\n",
    "\n",
    "# 量化权重\n",
    "torch.quantization.convert(model_fp32, inplace=True)\n",
    "\n",
    "# 打印量化后的权重\n",
    "print(\"量化后的 INT8 模型的权重:\\n\", model_fp32.linear.weight())\n",
    "\n",
    "# 获取量化参数\n",
    "weight_observer = model_fp32.linear.weight().q_per_channel_scales()\n",
    "weight_zero_points = model_fp32.linear.weight().q_per_channel_zero_points()\n",
    "print(\"量化权重的 scale:\", weight_observer)\n",
    "print(\"量化权重的 zero_point:\", weight_zero_points)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84d27d79-424a-404c-b076-58545107fe00",
   "metadata": {},
   "source": [
    "### 非对称量化\n",
    "\n",
    "模型量化方式修改为非对称量化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64256bcf-3e75-4b46-957f-ffdf8a5fdac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_fp32 = SimpleLinearModel(fp32_weights)\n",
    "\n",
    "# 打印 FP32 模型的权重\n",
    "print(\"FP32 模型的权重:\\n\", model_fp32.linear.weight)\n",
    "\n",
    "# ------------------------ 修改部分开始 ------------------------\n",
    "\n",
    "import torch.quantization as quant\n",
    "\n",
    "# 自定义的 qconfig，使用非对称量化\n",
    "custom_qconfig = quant.QConfig(\n",
    "    activation=quant.MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),\n",
    "    weight=quant.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_affine)\n",
    ")\n",
    "\n",
    "# 应用自定义的 qconfig 到模型\n",
    "model_fp32.qconfig = custom_qconfig\n",
    "\n",
    "# 插入量化准备步骤\n",
    "quant.prepare(model_fp32, inplace=True)\n",
    "\n",
    "# 量化权重\n",
    "quant.convert(model_fp32, inplace=True)\n",
    "\n",
    "# 打印量化后的权重\n",
    "quantized_weight = model_fp32.linear.weight()\n",
    "print(\"量化后的 INT8 模型的权重（int_repr）：\\n\", quantized_weight.int_repr())\n",
    "print(\"量化权重的 scale:\", quantized_weight.q_scale())\n",
    "print(\"量化权重的 zero_point:\", quantized_weight.q_zero_point())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f266a2fb-a305-468e-9d27-795e898ab0a8",
   "metadata": {},
   "source": [
    "### 对称量化\n",
    "\n",
    "所有方法修改为对称量化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26be6efc-3c53-4c0c-89ad-f9452cc28882",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.quantization\n",
    "\n",
    "# 假设 FP32 的张量\n",
    "fp32_values = torch.tensor([3.0, -5.5, 0.0, 6.0, -6.0, 2.5], dtype=torch.float32)\n",
    "print(f\"FP32 的示例张量: {fp32_values}\\n\")\n",
    "\n",
    "# 定义 PyTorch 的量化和反量化函数（对称量化）\n",
    "def pytorch_quantize_symmetric(fp32_tensor):\n",
    "    # 使用 MinMaxObserver 计算缩放因子，指定 dtype 为 torch.qint8，使用对称量化\n",
    "    q_params = torch.quantization.MinMaxObserver(\n",
    "        dtype=torch.qint8, qscheme=torch.per_tensor_symmetric\n",
    "    )\n",
    "    q_params(fp32_tensor)\n",
    "    scale, zero_point = q_params.calculate_qparams()\n",
    "\n",
    "    # 量化\n",
    "    int8_tensor = torch.quantize_per_tensor(fp32_tensor, scale.item(), zero_point.item(), dtype=torch.qint8)\n",
    "    return int8_tensor, scale.item(), zero_point\n",
    "\n",
    "def pytorch_dequantize(int8_tensor):\n",
    "    # 反量化\n",
    "    fp32_tensor = int8_tensor.dequantize()\n",
    "    return fp32_tensor\n",
    "\n",
    "# 量化并获取 PyTorch 结果（对称量化）\n",
    "int8_tensor, scale, zero_point = pytorch_quantize_symmetric(fp32_values)\n",
    "print(\"PyTorch 对称量化后的 INT8 数值：\", int8_tensor.int_repr())\n",
    "print(\"PyTorch 使用的 scale:\", scale, \"zero_point:\", zero_point)\n",
    "\n",
    "# 反量化\n",
    "recovered_fp32_pytorch = pytorch_dequantize(int8_tensor)\n",
    "print(\"PyTorch 反量化恢复后的 FP32 数值：\", recovered_fp32_pytorch)\n",
    "\n",
    "print(\"\\n=====================\\n\")\n",
    "\n",
    "# 对比与自定义的量化方式（对称量化）\n",
    "def custom_quantize_compare_symmetric(fp32_values):\n",
    "    # 获取张量的绝对最大值\n",
    "    x_max = fp32_values.abs().max().item()\n",
    "    \n",
    "    # 定义量化后整数数值的范围\n",
    "    qmin, qmax = -128, 127  # 对应 torch.qint8\n",
    "    \n",
    "    # 使用 PyTorch 的方式计算 scale\n",
    "    scale_custom = x_max / 127.5  # PyTorch源码：scale = max_val_pos / (float(quant_max - quant_min) / 2)\n",
    "    \n",
    "    # 对称量化时，zero_point 固定为 0\n",
    "    zero_point_custom = 0\n",
    "    \n",
    "    print(\"自定义计算的 scale:\", scale_custom, \"zero_point:\", zero_point_custom)\n",
    "    \n",
    "    # 使用 torch 的 round 函数以匹配 PyTorch 的量化行为\n",
    "    def quantize(fp32_tensor, scale, zero_point):\n",
    "        # 计算量化值\n",
    "        int8_tensor = torch.round(fp32_tensor / scale) + zero_point\n",
    "        # 限制在 [qmin, qmax] 范围内\n",
    "        int8_tensor = torch.clamp(int8_tensor, qmin, qmax)\n",
    "        return int8_tensor.to(torch.int8)\n",
    "    \n",
    "    def dequantize(int8_tensor, scale, zero_point):\n",
    "        # 反量化\n",
    "        fp32_tensor = (int8_tensor.float() - zero_point) * scale\n",
    "        return fp32_tensor\n",
    "    \n",
    "    # 量化\n",
    "    int8_values_custom = quantize(fp32_values, scale_custom, zero_point_custom)\n",
    "    print(\"自定义对称量化后的 INT8 数值：\", int8_values_custom)\n",
    "    \n",
    "    # 反量化\n",
    "    recovered_fp32_custom = dequantize(int8_values_custom, scale_custom, zero_point_custom)\n",
    "    print(\"自定义对称反量化恢复后的 FP32 数值：\", recovered_fp32_custom)\n",
    "    \n",
    "    return recovered_fp32_custom  # 返回以便后续使用\n",
    "\n",
    "# 运行自定义对称量化并比较\n",
    "print(\"运行自定义对称量化并比较：\")\n",
    "recovered_fp32_custom = custom_quantize_compare_symmetric(fp32_values)\n",
    "\n",
    "print(\"\\n=====================\\n\")\n",
    "\n",
    "# 使用 fp32_values 作为线性层参数\n",
    "\n",
    "# 定义一个简单的线性模型\n",
    "class SimpleLinearModel(nn.Module):\n",
    "    def __init__(self, weights, bias=None):\n",
    "        super(SimpleLinearModel, self).__init__()\n",
    "        # 假设输入特征数为6，输出特征数为1\n",
    "        self.linear = nn.Linear(6, 1, bias=False)\n",
    "        # 初始化权重\n",
    "        with torch.no_grad():\n",
    "            self.linear.weight = nn.Parameter(weights.view(1, 6))  # 权重形状为 [out_features, in_features]\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "# 创建 FP32 模型\n",
    "fp32_weights = fp32_values  # [6]\n",
    "model_fp32 = SimpleLinearModel(fp32_weights)\n",
    "\n",
    "# 打印 FP32 模型的权重\n",
    "print(\"FP32 模型的权重:\\n\", model_fp32.linear.weight)\n",
    "\n",
    "# 和之前非对称保持一致方便对比\n",
    "\n",
    "import torch.quantization as quant\n",
    "\n",
    "# 自定义的 qconfig，使用对称量化\n",
    "custom_qconfig = quant.QConfig(\n",
    "    activation=quant.MinMaxObserver.with_args(\n",
    "        dtype=torch.quint8, qscheme=torch.per_tensor_symmetric\n",
    "    ),\n",
    "    weight=quant.MinMaxObserver.with_args(\n",
    "        dtype=torch.qint8, qscheme=torch.per_tensor_symmetric\n",
    "    )\n",
    ")\n",
    "\n",
    "# 应用自定义的 qconfig 到模型，结果和model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')一致\n",
    "model_fp32.qconfig = custom_qconfig\n",
    "\n",
    "# 准备量化\n",
    "quant.prepare(model_fp32, inplace=True)\n",
    "\n",
    "# 量化权重\n",
    "quant.convert(model_fp32, inplace=True)\n",
    "\n",
    "# 打印量化后的权重\n",
    "quantized_weight = model_fp32.linear.weight()\n",
    "print(\"量化后的 INT8 模型的权重（int_repr）：\\n\", quantized_weight.int_repr())\n",
    "print(\"量化权重的 scale:\", quantized_weight.q_scale())\n",
    "print(\"量化权重的 zero_point:\", quantized_weight.q_zero_point())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e89242a-deda-48eb-819b-ea62b2ab40dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}