GPTJ-group-quantized.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e421bdb2",
   "metadata": {},
   "source": [
    "Copyright (c) 2023 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fff53455",
   "metadata": {},
   "source": [
    "# Faster Text Generation with GPT-J using 4-bit Weight Quantization on IPUs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a67dd960",
   "metadata": {},
   "source": [
    "The speed of text generation with large language models is often limited by the time it takes to read a model state from memory. One way to alleviate this issue is to: \n",
    " * compress the model state for storage in low-bandwidth, external memory and for communication with high-bandwidth on-chip memory\n",
    " * decompress the model state on-chip into a number format you can compute with (for example float16).\n",
    " \n",
    "Recently, many neural network practitioners have found that compressing model parameters to just 4 bits has minimal effect on the quality of model outputs.\n",
    "\n",
    "Group quantisation is a simple approach for compressing model parameters to 4 bits with no finetuning and is described in\n",
    "[\"FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU\"](https://arxiv.org/abs/2303.06865). \n",
    "\n",
    "Here we will show you how to apply this technique to GPT-J on IPUs. \n",
    "\n",
    "In the notebook \"Text Generation with GPT-J 6B on IPUs\" `GPTJ-generative-inference.ipynb` you learned how to generate text with GPT-J, an accessible 6B parameter language model. You saw: \n",
    "\n",
    "- how GPT-J performs on NLP tasks using both a base and fine-tuned checkpoint. \n",
    "- the effects on output quality from adjustments to prompt structure.\n",
    "- throughput improvements from batching text queries.\n",
    "\n",
    "In this notebook you will:\n",
    "\n",
    "- compress GPT-J weights to 4 bits, using 4x less memory.\n",
    "- speed up GPT-J inference by ~1.5x with minimal degradation of MNLI task performance.\n",
    "- see the trade-off between speed and accuracy."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b2b88b28",
   "metadata": {},
   "source": [
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------|--------------|\n",
    "|   NLP   |  Question answering | GPT-J | Glue-MNLI| Inference | 16 | 30min (+1h for final cell)|"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c498edbf",
   "metadata": {},
   "source": [
    "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e934b1e7",
   "metadata": {},
   "source": [
    "## Dependencies and configuration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d546a139",
   "metadata": {},
   "source": [
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f2571d11",
   "metadata": {},
   "source": [
    "Install the dependencies the notebook needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3947002b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install  -r requirements.txt\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2ce871a4",
   "metadata": {},
   "source": [
    "To make it easier for you to run this demo, we read in some configuration related to the environment you are running the notebook in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4626a0e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "number_of_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "if number_of_ipus < 4:\n",
    "    raise ValueError(\"This notebook is designed to run with at least 4 IPUs\")\n",
    "\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"./exe_cache/\")\n",
    "os.environ[\"POPART_CACHE_DIR\"] = executable_cache_dir\n",
    "checkpoint_directory = os.getenv(\"CHECKPOINT_DIR\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50942060",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gc-monitor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "21c4ef02",
   "metadata": {},
   "source": [
    "## How to generate text with GPT-J on the IPU with 4-bit weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "365f2688",
   "metadata": {},
   "source": [
    "We start by showing you how to generate text using GPT-J with 4-bit weights on the Graphcore IPU. As we did in our previous notebook `GPTJ-generative-inference.ipynb`, we load a configuration to create a pipeline object that we can use for generating text interactively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "366b19e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Setup ---\n",
    "import run_inference\n",
    "\n",
    "config, *_ = run_inference.gptj_config_setup(\n",
    "    \"config/inference.yml\", \"release\", \"gpt-j-gq-4bit\"\n",
    ")\n",
    "print(config.dumps_yaml())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dd11ec0c",
   "metadata": {},
   "source": [
    "We start by using the same config as we did previously, with a single change. We modify `config.execution.group_quantise_weights` to some value `n` greater than 0. This config option divides the columns of weight matrices into groups of `n`. Values within each group are then binned into one of 16 values (4 bits), and we store scale and bias factors in float16 to allow us to convert from 4 bits, back to float16 when we compute. This means the column dimension of the weights matrix must be divisible by `n`. Since native 4-bit formats don't exist, we store four consecutive 4-bit integers in a 16-bit integer value. As a result, `n` must be divisible by 4.\n",
    "\n",
    "Since GPT-J weights all have a column dimension of `4096`, a convenient choice for `n` is 64. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "193979c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Size of groups for 4-bit quantization\n",
    "config.execution.group_quantise_weights = 64"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5e966300",
   "metadata": {},
   "source": [
    "Next, we download and compress EleutherAI's pretrained weights from Hugging Face. Compressing weights takes a bit of time, but you only need to do this once when you are happy with your quantization hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7aa855",
   "metadata": {},
   "outputs": [],
   "source": [
    "import api\n",
    "\n",
    "int4_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"EleutherAI/gpt-j-6b\",\n",
    "    sequence_length=512,\n",
    "    micro_batch_size=1,\n",
    "    output_length=20,\n",
    "    print_live=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "003a029a",
   "metadata": {},
   "source": [
    "Let's ask our newly quantized GPT-J model the same question we asked in our first notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c68db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = int4_model(\"What is the capital of France?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "968a9e49",
   "metadata": {},
   "source": [
    "The answer is correct, if a bit repetitive! But this behaviour is expected for models of this size that have not been fine-tuned for particular types of prompts or instructions.\n",
    "\n",
    "So how fast was it? We added a timer to our pipeline to measure token generation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8902752",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    f\"Average token generation time for our compressed model is {int4_model.token_generation_time:.3f} secs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e99bff03",
   "metadata": {},
   "outputs": [],
   "source": [
    "int4_model.detach()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7d9717f7",
   "metadata": {},
   "source": [
    "Let's also compare with an uncompressed float16 model to look at the speedup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "653dd249",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.execution.group_quantise_weights = 0\n",
    "fp16_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"EleutherAI/gpt-j-6b\",\n",
    "    sequence_length=512,\n",
    "    micro_batch_size=1,\n",
    "    output_length=20,\n",
    "    print_live=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3149cba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = fp16_model(\"What is the capital of France?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d9a09c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    f\"Average token generation time for base model is {fp16_model.token_generation_time:.3f} secs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be9de5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "speedup = fp16_model.token_generation_time / int4_model.token_generation_time\n",
    "print(f\"int4 model generates tokens {speedup:.3f}x faster than fp16 model!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843471a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fp16_model.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eedb165b",
   "metadata": {},
   "outputs": [],
   "source": [
    "del int4_model, fp16_model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "813743ea",
   "metadata": {},
   "source": [
    "The compressed model is definitely faster! We can see that the answers are similar, but not exactly the same. Both models answered the question correctly, then started producing related text to fill the output quota of 20 tokens.\n",
    "\n",
    "Next, we will try to quantify the change in model quality from quantising weights."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c9f92206",
   "metadata": {},
   "source": [
    "## Compressing model weights impacts task performance (but not by much!)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "32809662",
   "metadata": {},
   "source": [
    "Although interactively asking simple questions to language models is important for getting a subjective feel for model quality, we would prefer to have more objective assessments where possible. In this section we will return to the MNLI task, in which we use a language model to classify the logical relationship between a premise and a hypothesis (entailment, contradiction, or neutral). Once again we'll download the MNLI fine-tuned weights from Hugging Face, but this time we'll also compress them using the same scheme.\n",
    "\n",
    "First, let's set our config for the MNLI validation task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe4de90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The number of tokens generated before stopping\n",
    "# Note the model will stop before this if it generates an <|endoftext|> token\n",
    "config.inference.output_length = 5\n",
    "# The number of prompts which will be processed at once\n",
    "config.execution.micro_batch_size = 12\n",
    "# The maximum tokenized sequence length (input + generated) handled by the model\n",
    "config.model.sequence_length = 256\n",
    "\n",
    "# Size of groups for 4-bit quantization\n",
    "config.execution.group_quantise_weights = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b662261c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "dataset = datasets.load_dataset(\"glue\", \"mnli\", split=\"validation_mismatched\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "691add5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_accuracy(predictions):\n",
    "    mnli_classes = [\"entailment\", \"neutral\", \"contradiction\", \"unknown\"]\n",
    "    correct = [\n",
    "        pred == mnli_classes[actual]\n",
    "        for pred, actual in zip(predictions, dataset[:][\"label\"])\n",
    "    ]\n",
    "    return sum(correct) * 100 / len(predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf92e14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_int4_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"Graphcore/gptj-mnli\",\n",
    "    sequence_length=256,\n",
    "    print_live=True,\n",
    ")\n",
    "mnli_int4_pipeline = api.GPTJEntailmentPipeline.from_gptj_pipeline(mnli_int4_model)\n",
    "int4_out = mnli_int4_pipeline(\n",
    "    premise=dataset[:][\"premise\"],\n",
    "    hypothesis=dataset[:][\"hypothesis\"],\n",
    "    print_live=False,\n",
    "    output_length=5,\n",
    ")\n",
    "\n",
    "int4_acc = compute_accuracy(int4_out)\n",
    "print(f\"Compressed model accuracy is {int4_acc:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e9998f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_int4_model.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7086331c",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.execution.group_quantise_weights = 0\n",
    "mnli_fp16_model = api.GPTJPipeline(\n",
    "    config,\n",
    "    \"Graphcore/gptj-mnli\",\n",
    "    sequence_length=256,\n",
    "    print_live=True,\n",
    ")\n",
    "mnli_fp16_pipeline = api.GPTJEntailmentPipeline.from_gptj_pipeline(mnli_fp16_model)\n",
    "fp16_out = mnli_fp16_pipeline(\n",
    "    premise=dataset[:][\"premise\"],\n",
    "    hypothesis=dataset[:][\"hypothesis\"],\n",
    "    print_live=False,\n",
    "    output_length=5,\n",
    ")\n",
    "\n",
    "fp16_acc = compute_accuracy(fp16_out)\n",
    "print(f\"Base model accuracy is {fp16_acc:.2f}%\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ce357677",
   "metadata": {},
   "source": [
    "You can see that quantising to 4 bits results in just a 1.27% degradation in accuracy! \n",
    "\n",
    "Shall we see what happens to accuracy and speed when we vary the group size? If you decide to run the `for` loop below you will cycle through group sizes of 16, 32, 128, and 256, and rerun our MNLI validation pipeline for each group size. This will take a while as we need to compress the checkpoint again for every new group size! If you just want to see the results, scroll down to where we have plotted this for you.\n",
    "\n",
    "<img src=\"./imgs/gq-speed-accuracy-tradeoff.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "867a821e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnli_fp16_model.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f2f817",
   "metadata": {},
   "outputs": [],
   "source": [
    "times = {\n",
    "    0: mnli_fp16_pipeline.token_generation_time,\n",
    "    64: mnli_int4_pipeline.token_generation_time,\n",
    "}\n",
    "accuracies = {0: fp16_acc, 64: int4_acc}\n",
    "\n",
    "for gs in [16, 32, 128, 256]:\n",
    "    config.execution.group_quantise_weights = gs\n",
    "    mnli_int4_model = api.GPTJPipeline(\n",
    "        config,\n",
    "        \"Graphcore/gptj-mnli\",\n",
    "        sequence_length=256,\n",
    "        print_live=True,\n",
    "    )\n",
    "    mnli_int4_pipeline = api.GPTJEntailmentPipeline.from_gptj_pipeline(mnli_int4_model)\n",
    "    int4_out = mnli_int4_pipeline(\n",
    "        premise=dataset[:][\"premise\"],\n",
    "        hypothesis=dataset[:][\"hypothesis\"],\n",
    "        print_live=False,\n",
    "        output_length=5,\n",
    "    )\n",
    "\n",
    "    int4_acc = compute_accuracy(int4_out)\n",
    "    times[gs] = mnli_int4_pipeline.token_generation_time\n",
    "    accuracies[gs] = int4_acc\n",
    "    mnli_int4_pipeline.detach()\n",
    "    del mnli_int4_pipeline, mnli_int4_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a5ba00",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "keys = list(times.keys())\n",
    "keys.sort()\n",
    "for k in keys:\n",
    "    plt.plot(times[k] * 1000, accuracies[k], marker=\"o\")\n",
    "\n",
    "\n",
    "def bits_per_param(n):\n",
    "    # n = group size\n",
    "    # 4 bits per param, plus float 16 scale and bias for each group\n",
    "    return (n * 4 + 2 * 16) / n\n",
    "\n",
    "\n",
    "plt.legend(\n",
    "    [\"16 bits per param (uncompressed)\"]\n",
    "    + [f\"{bits_per_param(k)} bits per param (group_size={k})\" for k in keys[1:]]\n",
    ")\n",
    "plt.xlabel(\"Batched token generation time (ms)\")\n",
    "plt.ylabel(\"MNLI accuracy (%)\")\n",
    "plt.title(\n",
    "    \"Finetuned GPT-J (6B) MNLI speed/accuracy \\n trade-off for varying int4 quantization group size\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fe21c19",
   "metadata": {},
   "source": [
    "You can see that the uncompressed checkpoint produces the most accurate results, but is also the slowest. As you increase `group_size`, token generation time gets smaller and accuracy degrades quite smoothly. For group sizes > 128, you can see that accuracy degrades very sharply, without much improvement in speed. The best tradeoff looks to be for `group_size=64`, since you need a significantly slower model to obtain a marginal improvement in accuracy, whereas an only slightly faster model produces quite a big dropoff in accuracy. Depending on your requirements, you might decide differently, for example you may want a little bit more quality even if it means your model runs slower. \n",
    "\n",
    "If you want to try this for your own models and use-cases you can perform an evaluation like this to help you decide!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4d9d516f",
   "metadata": {},
   "source": [
    "## Conclusions and next steps"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bff1ebf2",
   "metadata": {},
   "source": [
    "In this notebook, we showed how:\n",
    "1. to run 4-bit inference on GPT-J by changing the `group_quantise_weights` config option.\n",
    "2. compressing weights to 4-bits with group quantisation gives a 2x speed up over float16 on text generation on IPUs for a batch size of 1.\n",
    "3. quantising weights results in just a small drop off in accuracy on the MNLI entailment task.\n",
    "4. groups of size 64 provide a good tradeoff between speed and accuracy.\n",
    "\n",
    "Efficiently decompressing 4-bit integers to float16 values is surprisingly difficult for any hardware. We have some guidelines for [how to write efficient custom C++ code for the IPU](https://docs.graphcore.ai/projects/poplar-user-guide/en/latest/poplar_programs.html), and a description of how we wrote the code for our [custom op](https://github.com/graphcore/popxl-addons/blob/master/popxl_addons/ops/group_quantize_decompress/group_quantize_decompressx.cpp#L51-L109) and the [tile vertex (per thread kernel)](https://github.com/graphcore/popxl-addons/blob/master/popxl_addons/ops/group_quantize_decompress/group_quantize_decompress_codelet.cpp) for quickly decompressing weights. \n",
    "\n",
    "You can also run our other notebooks exploring other efficiency wins for deep learning such as:\n",
    "- [Converting FLAN-T5 XL to float16 for faster inference](https://ipu.dev/tvxZ3Q)\n",
    "- [A how-to on training GPT models in float8 on IPUs with unit scaling](https://ipu.dev/qXfm2a)\n",
    "- [Accelerating transformers with packing for fine-tuning and inference](https://ipu.dev/q6HAUX)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e7624634",
   "metadata": {},
   "source": [
    "Try out the other [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to see how how IPUs perform on other tasks."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}