{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM Compressor Workbench -- Getting Started\n",
    "\n",
    "This notebook will demonstrate how common [LLM Compressor](https://github.com/vllm-project/llm-compressor) flows can be run on the [opendatahub/llmcompressor-workbench](https://quay.io/repository/opendatahub/llmcompressor-workbench) image.\n",
    "\n",
    "We will show how a user can compress and evaluate a Large Language Model, first without data and then with a calibration dataset.\n",
    "\n",
    "The notebook will detect if a GPU is available. If one is not available, it will demonstrate an abbreviated run, so users without GPU access can still get a feel for `llm-compressor`.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> If you are not using the Workbench image, just be sure to have lm_eval>=0.4.8 and llmcompressor>=0.5.1 installed\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1\\) Data-Free Model Compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "use_gpu = torch.cuda.is_available()\n",
    "print(f\"GPU acceleration available: {use_gpu}\")\n",
    "if not use_gpu:\n",
    "    print(\"Running in CPU-only mode - operations will be slower and use abbreviated datasets\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llmcompressor.modifiers.quantization import QuantizationModifier\n",
    "\n",
    "# model to compress\n",
    "model_id = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n",
    "\n",
    "# This recipe will quantize all Linear layers except those in the `lm_head`,\n",
    "#  which is often sensitive to quantization. The W4A16 scheme compresses\n",
    "#  weights to 4-bit integers while retaining 16-bit activations.\n",
    "recipe = QuantizationModifier(\n",
    "    targets=\"Linear\", scheme=\"W4A16\", ignore=[\"lm_head\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up model using huggingface API\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, device_map=\"auto\", torch_dtype=\"auto\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run compression using `oneshot`\n",
    "from llmcompressor import oneshot\n",
    "\n",
    "model = oneshot(model=model, recipe=recipe, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model and tokenizer\n",
    "model_dir = \"./\" + model_id.split(\"/\")[-1] + \"-W4A16\"\n",
    "model.save_pretrained(model_dir)\n",
    "tokenizer.save_pretrained(model_dir);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2\\) Evaluate compressed model using open-source `lm_eval` framework\n",
    "\n",
    "We will evaluate the performance of the model on the [`wikitext`](https://paperswithcode.com/dataset/wikitext-2) language modeling dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lm_eval\n",
    "from lm_eval.utils import make_table\n",
    "\n",
    "results = lm_eval.simple_evaluate(\n",
    "    model=\"vllm\" if use_gpu else \"hf\",\n",
    "    model_args={\n",
    "        \"pretrained\": model_dir,\n",
    "        \"add_bos_token\": True,\n",
    "        \"device\": \"auto\"\n",
    "    },\n",
    "    tasks=[\"wikitext\"],\n",
    "    batch_size=\"auto\" if use_gpu else 4,\n",
    "    limit=None if use_gpu else 4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Tasks  |Version|Filter|n-shot|    Metric     |   | Value |   |Stderr|\n",
      "|--------|------:|------|-----:|---------------|---|------:|---|------|\n",
      "|wikitext|      2|none  |     0|bits_per_byte  |↓  | 0.7586|±  |   N/A|\n",
      "|        |       |none  |     0|byte_perplexity|↓  | 1.6918|±  |   N/A|\n",
      "|        |       |none  |     0|word_perplexity|↓  |16.6397|±  |   N/A|\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(make_table(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3\\) Calibrated Compression with a Dataset\n",
    "\n",
    "Some more advanced compression algorithms require a small dataset of calibration samples that are meant to be a representative random subset of the language the model will see at inference.\n",
    "\n",
    "We will show how the previous section can be augmented with a calibration dataset and GPTQ, one of the first published LLM compression algorithms.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> This will take several minutes if no GPU is available\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will use a new recipe running GPTQ (https://arxiv.org/abs/2210.17323)\n",
    "# to reduce error caused by quantization. GPTQ requires a calibration dataset.\n",
    "from llmcompressor.modifiers.quantization import GPTQModifier\n",
    "\n",
    "recipe = GPTQModifier(targets=\"Linear\", scheme=\"W4A16\", ignore=[\"lm_head\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Create the calibration dataset, using Huggingface datasets API\n",
    "dataset_id = \"HuggingFaceH4/ultrachat_200k\"\n",
    "\n",
    "# Select number of samples. 512 samples is a good place to start.\n",
    "# Increasing the number of samples can improve accuracy.\n",
    "num_calibration_samples = 512 if use_gpu else 4\n",
    "max_sequence_length = 2048 if use_gpu else 16\n",
    "\n",
    "# Load dataset\n",
    "ds = load_dataset(dataset_id, split=\"train_sft\")\n",
    "# Shuffle and grab only the number of samples we need\n",
    "ds = ds.shuffle(seed=42).select(range(num_calibration_samples))\n",
    "\n",
    "# Preprocess and tokenize into format the model uses\n",
    "def preprocess(example):\n",
    "    text = tokenizer.apply_chat_template(\n",
    "            example[\"messages\"],\n",
    "            tokenize=False,\n",
    "        )\n",
    "    return tokenizer(\n",
    "        text,\n",
    "        padding=False,\n",
    "        max_length=max_sequence_length,\n",
    "        truncation=True,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "\n",
    "ds = ds.map(preprocess, remove_columns=ds.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# oneshot modifies model in-place, so reload\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, device_map=\"auto\", torch_dtype=\"auto\"\n",
    ")\n",
    "# run oneshot again, with dataset\n",
    "model = oneshot(\n",
    "    model=model,\n",
    "    dataset=ds,\n",
    "    recipe=recipe,\n",
    "    max_seq_length=max_sequence_length,\n",
    "    num_calibration_samples=num_calibration_samples,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model and tokenizer\n",
    "model_dir = \"./\" + model_id.split(\"/\")[-1] + \"-GPTQ-W4A16\"\n",
    "model.save_pretrained(model_dir)\n",
    "tokenizer.save_pretrained(model_dir);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4\\) Rerun `lm_eval`\n",
    "\n",
    "Note that perplexity score has improved (lower is better) for this `TinyLlama` model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = lm_eval.simple_evaluate(\n",
    "    model=\"vllm\" if use_gpu else \"hf\",\n",
    "    model_args={\n",
    "        \"pretrained\": model_dir,\n",
    "        \"add_bos_token\": True,\n",
    "        \"device\": \"auto\"\n",
    "    },\n",
    "    tasks=[\"wikitext\"],\n",
    "    batch_size=\"auto\" if use_gpu else 4,\n",
    "    limit=None if use_gpu else 4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Tasks  |Version|Filter|n-shot|    Metric     |   | Value |   |Stderr|\n",
      "|--------|------:|------|-----:|---------------|---|------:|---|------|\n",
      "|wikitext|      2|none  |     0|bits_per_byte  |↓  | 0.7497|±  |   N/A|\n",
      "|        |       |none  |     0|byte_perplexity|↓  | 1.6814|±  |   N/A|\n",
      "|        |       |none  |     0|word_perplexity|↓  |16.0972|±  |   N/A|\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(make_table(results))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}