examples/smoothquant_opt_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SmoothQuant on OPT-13B\n",
    "\n",
    "### Guangxuan Xiao\\*, Ji Lin\\*, Mickael Seznec, Julien Demouth, Song Han\n",
    "\n",
    "In this notebook, we use OPT-13B model to demonstrate SmoothQuant can use 8-bit for both weights and activations to achieve the same accuracy as FP16 models. Unlike previous method [[Dettmers *et al.*, 2022]](https://arxiv.org/abs/2208.07339), SmoothQuant enables fully INT8 GEMMs for linear layers and does not require high precision numbers to represent outliers. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates SmoothQuant on OPT-13B in consideration of the user's resouce constraints. We have tested SmoothQuant on up to 176 billion parameter models (OPT-175B, BLOOM-176B, GLM-130B). You can also adjust the model name to validate SmoothQuant on other models. `../act_scales/` provides the activation channel scales for OPT and BLOOM models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to run this notebook, you need to install the following packages:\n",
    "\n",
    "- smoothquant\n",
    "- PyTorch\n",
    "- Transformers\n",
    "- Accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers.models.opt.modeling_opt import OPTAttention, OPTDecoderLayer, OPTForCausalLM\n",
    "from transformers import GPT2Tokenizer\n",
    "from smoothquant.smooth import smooth_lm\n",
    "from smoothquant.fake_quant import W8A8Linear"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we simulate the 8-bit dynamic per-tensor weight and activation quantization with FP16, i.e., fake quantization. We have implemented the real 8-bit quantization with INT8 CUTLASS GEMM kernels for both PyTorch and FasterTransformer. Please stay tuned for the release."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def quantize_model(model, weight_quant='per_tensor', act_quant='per_tensor', quantize_bmm_input=True):\n",
    "    for name, m in model.model.named_modules():\n",
    "        if isinstance(m, OPTDecoderLayer):\n",
    "            m.fc1 = W8A8Linear.from_float(m.fc1, weight_quant=weight_quant, act_quant=act_quant)\n",
    "            m.fc2 = W8A8Linear.from_float(m.fc2, weight_quant=weight_quant, act_quant=act_quant)\n",
    "        elif isinstance(m, OPTAttention):\n",
    "            # Her we simulate quantizing BMM inputs by quantizing the output of q_proj, k_proj, v_proj\n",
    "            m.q_proj = W8A8Linear.from_float(\n",
    "                m.q_proj, weight_quant=weight_quant, act_quant=act_quant, quantize_output=quantize_bmm_input)\n",
    "            m.k_proj = W8A8Linear.from_float(\n",
    "                m.k_proj, weight_quant=weight_quant, act_quant=act_quant, quantize_output=quantize_bmm_input)\n",
    "            m.v_proj = W8A8Linear.from_float(\n",
    "                m.v_proj, weight_quant=weight_quant, act_quant=act_quant, quantize_output=quantize_bmm_input)\n",
    "            m.out_proj = W8A8Linear.from_float(m.out_proj, weight_quant=weight_quant, act_quant=act_quant)\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following is an evaluator to see the performance of the model. We use a toy dataset (the first 1000 examples in the validation set of the Lambada dataset) to evaluate the model. You can replace it with your own dataset. The conclusion should be the same."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "**In this demo, we have simplified the evaluation by using the first 1,000 samples from the LAMBADA dataset's validation set. We employ the \"Last Token Prediction Accuracy\" as our evaluation metric. This approximate evaluation is intended for demonstration purposes, providing simple but meaningful comparisons of relative performance between methods. For a more strict assessment, we recommend using the [lm-eval-harness](https://github.com/EleutherAI/lm-evaluation-harness) to obtain the \"Last Word Prediction Accuracy\" for the LAMBADA dataset, which is the reported metric in our paper.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Evaluator:\n",
    "    def __init__(self, dataset, tokenizer, device):\n",
    "        self.dataset = dataset\n",
    "        self.tokenizer = tokenizer\n",
    "        self.device = device\n",
    "\n",
    "        # tokenize the dataset\n",
    "        def tokenize_function(examples):\n",
    "            example = self.tokenizer(examples['text'])\n",
    "            return example\n",
    "\n",
    "        self.dataset = self.dataset.map(tokenize_function, batched=True)\n",
    "        self.dataset.set_format(type='torch', columns=['input_ids'])\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def evaluate(self, model):\n",
    "        model.eval()\n",
    "        # The task is to predict the last word of the input.\n",
    "        total, hit = 0, 0\n",
    "        for batch in self.dataset:\n",
    "            input_ids = batch['input_ids'].to(self.device).unsqueeze(0)\n",
    "            label = input_ids[:, -1]\n",
    "            outputs = model(input_ids)\n",
    "            last_token_logits = outputs.logits[:, -2, :]\n",
    "            pred = last_token_logits.argmax(dim=-1)\n",
    "            total += label.size(0)\n",
    "            hit += (pred == label).sum().item()\n",
    "        acc = hit / total\n",
    "        return acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('facebook/opt-13b')\n",
    "dataset = load_dataset('lambada', split='validation[:1000]')\n",
    "evaluator = Evaluator(dataset, tokenizer, 'cuda')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FP16 Model Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first check the performance of the original FP16 model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_fp16 = OPTForCausalLM.from_pretrained('facebook/opt-13b', torch_dtype=torch.float16, device_map='auto')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original model (fp16) accuracy: 0.786\n"
     ]
    }
   ],
   "source": [
    "acc_fp16 = evaluator.evaluate(model_fp16)\n",
    "print(f'Original model (fp16) accuracy: {acc_fp16}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then quantize the model to W8A8 and check the performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive W8A8 Quantized Model Accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_w8a8 = quantize_model(model_fp16)\n",
    "print(model_w8a8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Naive W8A8 quantized model accuracy: 0.048\n"
     ]
    }
   ],
   "source": [
    "acc_w8a8 = evaluator.evaluate(model_w8a8)\n",
    "print(f'Naive W8A8 quantized model accuracy: {acc_w8a8}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see there is a significant accuracy drop. This is consistent with LLM.int8()'s finding: when the model size increases larger than 6.7B, systematic outliers will emerge in activations, which makes fully INT8 quantization impossible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SmoothQuant W8A8 Quantized Model Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's smooth the model, quantize it, and check the performance! In `../act_scales`, we provide the activation scales for OPT and BLOOM models. You can also use this notebook to test quantizing those models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = OPTForCausalLM.from_pretrained('facebook/opt-13b', torch_dtype=torch.float16, device_map='auto')\n",
    "act_scales = torch.load('../act_scales/opt-13b.pt')\n",
    "smooth_lm(model, act_scales, 0.5)\n",
    "model_smoothquant_w8a8 = quantize_model(model)\n",
    "print(model_smoothquant_w8a8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the smoothed model has the same accuracy as the FP16 model. This is because SmoothQuant smooths the outliers in activations and moves the quantization difficulty from activations to weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SmoothQuant W8A8 quantized model accuracy: 0.793\n"
     ]
    }
   ],
   "source": [
    "acc_smoothquant_w8a8 = evaluator.evaluate(model_smoothquant_w8a8)\n",
    "print(f'SmoothQuant W8A8 quantized model accuracy: {acc_smoothquant_w8a8}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 (conda)",
   "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.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c458cb81aeeb610631c72e4cc4799f00f630d4dfa7a554b37f8134a7fe160cb8"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}