examples/smoothquant_opt_real_int8_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SmoothQuant Real-INT8 Inference for PyTorch\n",
    "\n",
    "### Guangxuan Xiao\\*, Ji Lin\\*, Mickael Seznec, Julien Demouth, Song Han\n",
    "\n",
    "In this notebook, we use OPT-30B model to demonstrate the latency and memory advantages of SmoothQuant. We implement SmoothQuant real-INT8 inference for PyTorch with [CUTLASS](https://github.com/NVIDIA/cutlass) INT8 GEMM kernels, which are wrapped as PyTorch modules in [torch-int](https://github.com/Guangxuan-Xiao/torch-int).\n",
    "\n",
    "This notebook demonstrates SmoothQuant on OPT-30B because it is the largest model we can run both FP16 and INT8 inference on a single A100 GPU. For larger models requiring multiple GPUs, we recommend using the [FasterTransformer](https://github.com/NVIDIA/FasterTransformer) implementation of SmoothQuant.\n",
    "\n",
    "In order to run this notebook, you need to install the following packages:\n",
    "\n",
    "- [smoothquant](https://github.com/mit-han-lab/smoothquant)\n",
    "- [torch-int](https://github.com/Guangxuan-Xiao/torch-int)\n",
    "- [PyTorch](https://pytorch.org/)\n",
    "- [Transformers](https://github.com/huggingface/transformers)\n",
    "- [Accelerate](https://github.com/huggingface/accelerate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers.models.opt.modeling_opt import OPTForCausalLM\n",
    "from transformers import GPT2Tokenizer\n",
    "from smoothquant.opt import Int8OPTForCausalLM\n",
    "import os\n",
    "import gc\n",
    "from torch.nn.functional import pad\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "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 dataset. The conclusion should be the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Evaluator:\n",
    "    def __init__(self, dataset, tokenizer):\n",
    "        self.dataset = dataset\n",
    "        self.tokenizer = tokenizer\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",
    "        start = torch.cuda.Event(enable_timing=True)\n",
    "        end = torch.cuda.Event(enable_timing=True)\n",
    "        latency = 0\n",
    "        for batch in self.dataset:\n",
    "            input_ids = batch['input_ids'].cuda().unsqueeze(0)\n",
    "            label = input_ids[:, -1]\n",
    "            pad_len = 512 - input_ids.shape[1]\n",
    "            input_ids = pad(input_ids, (0, pad_len), value=1)\n",
    "            torch.cuda.synchronize()\n",
    "            start.record()\n",
    "            outputs = model(input_ids)\n",
    "            end.record()\n",
    "            torch.cuda.synchronize()\n",
    "            latency += start.elapsed_time(end)\n",
    "            last_token_logits = outputs.logits[:, -2-pad_len, :]\n",
    "            pred = last_token_logits.argmax(dim=-1)\n",
    "            total += label.size(0)\n",
    "            hit += (pred == label).sum().item()\n",
    "\n",
    "        acc = hit / total\n",
    "        lantecy = latency / len(self.dataset)\n",
    "        return acc, lantecy\n",
    "\n",
    "\n",
    "def print_model_size(model):\n",
    "    # https://discuss.pytorch.org/t/finding-model-size/130275\n",
    "    param_size = 0\n",
    "    for param in model.parameters():\n",
    "        param_size += param.nelement() * param.element_size()\n",
    "    buffer_size = 0\n",
    "    for buffer in model.buffers():\n",
    "        buffer_size += buffer.nelement() * buffer.element_size()\n",
    "\n",
    "    size_all_mb = (param_size + buffer_size) / 1024**2\n",
    "    print('Model size: {:.3f}MB'.format(size_all_mb))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('facebook/opt-30b')\n",
    "dataset = load_dataset('lambada', split='validation[:1000]')\n",
    "evaluator = Evaluator(dataset, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FP16 Model Accuracy and Latency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 57171.898MB\n",
      "FP16 accuracy: 0.807, per-sample lantecy: 263.633ms\n"
     ]
    }
   ],
   "source": [
    "model_fp16 = OPTForCausalLM.from_pretrained(\n",
    "    'facebook/opt-30b', torch_dtype=torch.float16, device_map='auto')\n",
    "print_model_size(model_fp16)\n",
    "acc_fp16, lantecy_fp16 = evaluator.evaluate(model_fp16)\n",
    "print(f'FP16 accuracy: {acc_fp16}, per-sample lantecy: {lantecy_fp16:.3f}ms')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "del model_fp16\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SmoothQuant W8A8 Quantized Model Accuracy and Latency"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We provide the already smoothed and quantized OPT model at `https://huggingface.co/mit-han-lab/opt-[MODEL-SIZE]-smoothquant`, where `[MODEL-SIZE]` can be `125m`, `1.3B`, `2.7B`, `6.7B`, `13B`, `30b`, and `66b`. You can load the INT8 model with the following code:\n",
    "\n",
    "```python\n",
    "from smoothquant.opt import Int8OPTForCausalLM\n",
    "model = Int8OPTForCausalLM.from_pretrained(\"mit-han-lab/opt-30b-smoothquant\")\n",
    "```\n",
    "\n",
    "We implement the following quantization flow for OPT models, which you can see details in [smoothquant/opt.py](../smoothquant/opt.py).\n",
    "\n",
    "![quantization flow](../figures/quantization_flow.png)\n",
    "\n",
    "You can also check [generate_act_scales.py](../examples/generate_act_scales.py) and [export_int8_model.py](../examples/export_int8_model.py) to see how we smooth, quantize and export INT8 models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 28945.603MB\n",
      "SmoothQuant INT8 accuracy: 0.798, per-sample lantecy: 212.361ms\n"
     ]
    }
   ],
   "source": [
    "model_smoothquant = Int8OPTForCausalLM.from_pretrained(\n",
    "    'mit-han-lab/opt-30b-smoothquant', torch_dtype=torch.float16, device_map='auto')\n",
    "print_model_size(model_smoothquant)\n",
    "acc_smoothquant, lantecy_smoothquant = evaluator.evaluate(model_smoothquant)\n",
    "print(\n",
    "    f'SmoothQuant INT8 accuracy: {acc_smoothquant}, per-sample lantecy: {lantecy_smoothquant:.3f}ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conlusion\n",
    "\n",
    "We can see that the SmoothQuant model has a similar accuracy as the FP16 model, but it is faster and uses less memory. This is because SmoothQuant reduces the quantization difficulty of activations and enables the use of INT8 GEMM kernels."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.15 (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.8.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b18562e22caa2a2bb5e6615862f7e7ce92f781ef7fc2a883871422ecfcd6595c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}