training/optimize-llama-2-gptq.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantize open LLMs using optimum and GPTQ\n",
    "\n",
    "The Hugging Face Optimum team collaborated with AutoGPTQ library to provide a simple API that apply GPTQ quantization on language models. With GPTQ quantization open LLMs to 8, 4, 3 or even 2 bits to run them on smaller Hardware without a big drop of performance. \n",
    "\n",
    "In the blog, you will learn how to:\n",
    "\n",
    "1. Setup our development environment\n",
    "2. Prepare quantization dataset\n",
    "3. Load and Quantize Model\n",
    "4. Test performance and inference speed\n",
    "5. Bonus: Run Inference with Text Generation Inference\n",
    "  \n",
    "But we before we get started lets take quick look on what GPTQ does. \n",
    "\n",
    "_Note: This tutorial was created and run on a g5.2xlarge AWS EC2 Instance, including an NVIDIA A10G GPU._\n",
    "\n",
    "\n",
    "## What is GPTQ?\n",
    "\n",
    "[GPTQ](https://arxiv.org/abs/2210.17323) is a post-training quantziation method to compress LLMs, like GPT. GPTQ compresses GPT models by reducing the number of bits needed to store each weight in the model, from 32 bits down to just 3-4 bits. This means the model takes up much less memory, so it can run on less Hardware, e.g. Single GPU for 13B Llama2 models. GPTQ analyzes each layer of the model separately and approximating the weights in a way that preserves the overall accuracy.\n",
    "\n",
    "The main benefits are:\n",
    "* Quantizes the weights of the model layer-by-layer to 4 bits instead of 16 bits, this reduces the needed memory by 4x.\n",
    "* Quantization is done gradually to minimize the accuracy loss from quantization.\n",
    "* Achieves same latency as fp16 model, but 4x less memory usage, sometimes faster due to custom kernels, e.g. [Exllama](https://github.com/turboderp/exllama)\n",
    "* Quantized weights can be saved to disk for a head of time quantization.\n",
    "\n",
    "_Note: GPTQ quantization only works for text model for now. Futhermore, the quantization process can take a lot of time. You check on the [Hugging Face Hub](https://huggingface.co/models?search=gptq) if there is not already a GPTQ quantized version of the model you want to use._\n",
    "\n",
    "--- \n",
    "\n",
    "## 1. Setup our development environment\n",
    "\n",
    "Let's start coding, but first, install our dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"transformers==4.32.1\" \"optimum==1.12.0\" \"auto-gptq==0.4.2\" \"accelerate==0.22.0\" \"safetensors>=0.3.1\" --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare quantization dataset\n",
    "\n",
    "GPTQ is a post-training quantization method, so we need to prepare a dataset to quantize our model. We can either use a dataset from the [Hugging Face Hub](https://huggingface.co/datasets) or use our own dataset. In this blog, we are going to use the [WikiText](https://huggingface.co/datasets/wikitext) dataset from the Hugging Face Hub. The dataset is used to quantize the weights to minimize the performance loss. It is recommended to use a quantization dataset with atleast `128` samples.\n",
    "\n",
    "_Note: [TheBloke](https://huggingface.co/TheBloke) a very active community member is contributing hundreds of gptq weights to the Hugging Face Hub. He mostly uses wikitext as quantization dataset for general domain models._\n",
    "\n",
    "If you want to use, e.g. your fine-tuning dataset for quantization you can provide it as a list instead of the \"id\", check out this [example](https://colab.research.google.com/drive/1_TIrmuKOFhuRRiTWN94iLKUFu6ZX4ceb).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset id from Hugging Face \n",
    "dataset_id = \"wikitext2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load and Quantize Model\n",
    "\n",
    "Optimum integrates GPTQ quantization in the `optimum.qptq` namespace with a `GPTQQuantizer`. The quantizer takes our dataset (id or list), bits, and model_seqlen as input. For more customization check [here](https://github.com/huggingface/optimum/blob/234a427450a7dcc978b227fa627ebcdab1764318/optimum/gptq/quantizer.py#L76).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.gptq import GPTQQuantizer\n",
    "\n",
    "# GPTQ quantizer\n",
    "quantizer = GPTQQuantizer(bits=4, dataset=dataset_id, model_seqlen=4096)\n",
    "quantizer.quant_method = \"gptq\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we have created our Quantizer we can load our model using Transformers. In our example we will quantize a [Llama 2 7B](https://huggingface.co/philschmid/llama-2-7b-instruction-generator), which we trained in my other blog post [\"Extended Guide: Instruction-tune Llama 2\"](https://www.philschmid.de/instruction-tune-llama-2). We are going to load our model in `fp16` since GPTQ adopts a mixed int4/fp16 quantization scheme where weights are quantized as int4 while activations remain in float16. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "# Hugging Face model id\n",
    "model_id = \"philschmid/llama-2-7b-instruction-generator\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False) # bug with fast tokenizer\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, torch_dtype=torch.float16) # we load the model in fp16 on purpose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we loaded our model we are ready to quantize it. \n",
    "_Note: Quantization can take process can take a lot of time depending on one's hardware. For this example the quantization on a single A10G GPU for a 7B model took ~minutes._ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import json\n",
    "\n",
    "# quantize the model \n",
    "quantized_model = quantizer.quantize_model(model, tokenizer)\n",
    "\n",
    "# save the quantize model to disk\n",
    "save_folder = \"quantized_llama\"\n",
    "quantized_model.save_pretrained(save_folder, safe_serialization=True)\n",
    "\n",
    "# load fresh, fast tokenizer and save it to disk\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id).save_pretrained(save_folder)\n",
    "\n",
    "# save quantize_config.json for TGI \n",
    "with open(os.path.join(save_folder, \"quantize_config.json\"), \"w\", encoding=\"utf-8\") as f:\n",
    "  quantizer.disable_exllama = False\n",
    "  json.dump(quantizer.to_dict(), f, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "since the model was partially offloaded it set `disable_exllama` to `True` to avoid an error. For inference and production load we want to leverage the exllama kernels. Therefore we need to change the `config.json`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(save_folder, \"config.json\"), \"r\", encoding=\"utf-8\") as f:\n",
    "  config = json.load(f)\n",
    "  config[\"quantization_config\"][\"disable_exllama\"] = False\n",
    "  with open(os.path.join(save_folder, \"config.json\"), \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(config, f, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test performance and inference speed\n",
    "\n",
    "Since the latest release of transformers we can load any GPTQ quantized model directly using the `AutoModelForCausalLM` class this. You can either load already quantized models from Hugging Face, e.g. [TheBloke/Llama-2-13B-chat-GPTQ](https://huggingface.co/TheBloke/Llama-2-13B-chat-GPTQ) or models you quantized yourself. Since we want to test here the results of our quantization we are going to load our quantized model from disk and compare it to our non quantize model. \n",
    "\n",
    "First lets our our non quantized model and test it on a simple prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time \n",
    "\n",
    "# The prompt is based on the fine-tuning from the model: https://www.philschmid.de/instruction-tune-llama-2#4-test-model-and-run-inference\n",
    "prompt = \"\"\"### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM.\n",
    "\n",
    "### Input:\n",
    "Dear [boss name],\n",
    "\n",
    "I'm writing to request next week, August 1st through August 4th,\n",
    "off as paid time off.\n",
    "\n",
    "I have some personal matters to attend to that week that require\n",
    "me to be out of the office. I wanted to give you as much advance\n",
    "notice as possible so you can plan accordingly while I am away.\n",
    "\n",
    "Thank you, [Your name]\n",
    "\n",
    "### Response:\n",
    "\"\"\"\n",
    "\n",
    "# helper function to generate text and measure latency\n",
    "def generate_helper(pipeline,prompt=prompt):\n",
    "    # warm up\n",
    "    for i in range(5):\n",
    "      _ = pipeline(\"Warm up\")\n",
    "\n",
    "    # measure latency in a simple way \n",
    "    start = time.time()\n",
    "    out = pipeline(prompt, max_new_tokens=100, do_sample=True, top_p=0.9,temperature=0.9)\n",
    "    end = time.time()\n",
    "    \n",
    "    generated_text = out[0][\"generated_text\"][len(prompt):]\n",
    "    \n",
    "    latency_per_token_in_ms = ((end-start)/len(pipeline.tokenizer(generated_text)[\"input_ids\"]))*1000\n",
    "    \n",
    "    # return the generated text and the latency\n",
    "    return {\"text\": out[0][\"generated_text\"][len(prompt):], \"latency\": f\"{round(latency_per_token_in_ms,2)}ms/token\"}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can load the vanilla transformers model and run inference using the `pipeline` class. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline\n",
    "\n",
    "# Hugging Face model id\n",
    "model_id = \"philschmid/llama-2-7b-instruction-generator\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, device_map=\"auto\", torch_dtype=torch.float16) # we load the model in fp16 on purpose\n",
    "\n",
    "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "lets create our vanilla base line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "\n",
    "vanilla_res = generate_helper(pipe)\n",
    "\n",
    "print(f\"Latency: {vanilla_res['latency']}\")\n",
    "print(f\"GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB\")\n",
    "print(f\"Generated Instruction: {vanilla_res['text']}\")\n",
    "\n",
    "# Latency: 37.49ms/token\n",
    "# GPU memory: 12.62 GB\n",
    "# Generated Instruction: Write a request for PTO letter to my boss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# clean up \n",
    "del pipe\n",
    "del model \n",
    "del tokenizer\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we have now our baseline we can test and validate our GPTQ quantize weights. Therefore we will use the new `gptq` integration into the `AutoModelForCausalLM` class where we can directly load the `gptq` weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline\n",
    "\n",
    "# path to gptq weights\n",
    "model_id = \"quantized_llama\"\n",
    "\n",
    "q_tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "q_model = AutoModelForCausalLM.from_pretrained(model_id, device_map=\"auto\", torch_dtype=torch.float16)\n",
    "\n",
    "qtq_pipe = pipeline(\"text-generation\", model=q_model, tokenizer=q_tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can test our quantized model on the same prompt as our baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpq_res = generate_helper(qtq_pipe)\n",
    "\n",
    "print(f\"Latency: {gpq_res['latency']}\")\n",
    "print(f\"GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB\")\n",
    "print(f\"Generated Instruction: {gpq_res['text']}\")\n",
    "\n",
    "# Latency: 36.0ms/token\n",
    "# GPU memory: 3.83 GB\n",
    "# Generated Instruction: Write a letter requesting time off"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For comparison the vanilla model needed ~12.6GB Memory and the GPTQ model needed ~3.8GB Memory, with equal performance. GPTQ allowed us to save ~4x memory (don't forget pytorch has default kernels). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Bonus: Run Inference with Text Generation Inference\n",
    "\n",
    "Text Generation Inference supports GPTQ model for more efficient deployments. We simply need to provide `gptq` as `QUANTIZE` environment variable when starting our container. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model=\"/home/ubuntu/test-gptq\"\n",
    "num_shard=1\n",
    "quantize=\"gptq\"\n",
    "max_input_length=1562\n",
    "max_total_tokens=4096 # 4096\n",
    "\n",
    "!docker run --gpus all -ti -p 8080:80 \\\n",
    "  -e MODEL_ID=$model \\\n",
    "  -e QUANTIZE=$quantize \\\n",
    "  -e NUM_SHARD=$num_shard \\\n",
    "  -e MAX_INPUT_LENGTH=$max_input_length \\\n",
    "  -e MAX_TOTAL_TOKENS=$max_total_tokens \\\n",
    "  -v $model:$model \\\n",
    "  ghcr.io/huggingface/text-generation-inference:1.0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can invoke our container using curl. \n",
    "_Note: The first request will be slow. _"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "curl 127.0.0.1:8080/generate \\\n",
    "    -X POST \\\n",
    "    -d '{\"inputs\":\"### Instruction:\\nUse the Input below to create an instruction, which could have been used to generate the input using an LLM.\\n\\n### Input:\\nDear [boss name],\\n\\nI am writing to request next week, August 1st through August 4th,\\noff as paid time off.\\n\\nI have some personal matters to attend to that week that require\\nme to be out of the office. I wanted to give you as much advance\\nnotice as possible so you can plan accordingly while I am away.\\n\\nThank you, [Your name]\\n\\n### Response:\",\"parameters\":{\"temperature\":0.2, \"top_p\": 0.95, \"max_new_tokens\": 256}}' \\\n",
    "    -H 'Content-Type: application/json'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With Text Generation inference we are achieving ~`22.942983ms` latency per token, which is 2x faster than transformers. If you plan to deploy your model in production, I would recommend to use Text Generation Inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}