Created a QAT Walkthrough notebook

Signed-off-by: Farshad Ghodsian <[email protected]>

Author: farshadghodsian

examples/llm_qat/notebooks/QAT_Walkthrough.ipynb

@@ -0,0 +1,298 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b54dd4ac-5b4f-42fb-af1a-8f3001abd08a",
+   "metadata": {},
+   "source": [
+    "# NVIDIA ModelOpt Quantization Aware Training (QAT) Walkthrough"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a695be45-1472-42bc-824e-5c992a487fa7",
+   "metadata": {},
+   "source": [
+    "**Quantization Aware Training (QAT)** is a method that simulates the effects of quantization during neural network post-training to preserve accuracy when deploying models in very-low-precision formats. Unlike post-training quantization, QAT inserts \"fake quantization\" nodes into the computational graph, mimicking the rounding and clamping operations that occur during actual quantization. This allows the model to adapt its weights and activations to mitigate accuracy loss.\n",
+    "\n",
+    "This notebook demonstrates how to apply Quantization Aware Training (QAT) to an LLM, Meta's Llama-3.1-8b in this case with NVIDIA's TensorRT Model Optimizer (ModelOpt) QAT toolkit. We walk through downloading and loading the model, calibrating it using an example dataset, specifically CNN/DailyMail dataset sample, applying NVFP4 quantization, generating outputs, and exporting the quantized model."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c3f7f931-ac38-494e-aea8-ca2cd6d05794",
+   "metadata": {},
+   "source": [
+    "## Installing Prerequisites and Dependancies"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d7d4f25f-e569-42cf-8022-bb7cc6f9ea6e",
+   "metadata": {},
+   "source": [
+    "If you haven't already, install the required dependencies for this notebook. Key dependancies include:\n",
+    "\n",
+    "- nvidia-modelopt\n",
+    "- torch\n",
+    "- transformers\n",
+    "- jupyterlab\n",
+    "\n",
+    "This repo contains a `examples/llm_qat/notebooks/requirements.txt` file that can be used to install all required dependancies."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab464a07-8a19-43a9-a715-81ccef350253",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "!pip install -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "99c6ca5d-0d08-4b6c-814f-b8a92a8469f2",
+   "metadata": {},
+   "source": [
+    "## Setting HuggingFace Token and Model for Download"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "09f3c6a7-1c9c-4254-9524-7e253528d9d7",
+   "metadata": {},
+   "source": [
+    "Set the HF_TOKEN environment variable making sure to update it to include you token (eg. `%env HF_TOKEN=hf_abdxyz...`)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb2071be-df85-4961-92b0-567830a37d71",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "%env HF_TOKEN=<YOUR_HUGGINGFACE_TOKEN>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4eab1f6a-5855-4f2a-a982-27b7e756deca",
+   "metadata": {},
+   "source": [
+    "As mentioned above, we will use Meta's **Llamma-3.1-8B-Instruct** in this example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6d25c2b1-a68b-4748-ac29-e8a893ce1762",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"meta-llama/Llama-3.1-8B-Instruct\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41b7c23f-748b-4c9e-8883-d6ca24af46ed",
+   "metadata": {},
+   "source": [
+    "## Import Required Libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "0ec71181-770a-4ee6-8760-c62cfab8340f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "53f8c53a-94ab-4968-aac8-866d850ef874",
+   "metadata": {},
+   "source": [
+    "## Download and Load Model and Tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "5f946576-83ac-45b5-a290-9a2167193e3d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "37c5f366ef204794bad4711ae6056d6c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model = AutoModelForCausalLM.from_pretrained(model_name).cuda()\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "tokenizer.pad_token = tokenizer.eos_token\n",
+    "tokenizer.padding_side = \"left\" # Setting this as tokenizer with the right padding_side may impact calibration accuracy."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fef14995-c3cd-445d-a93a-55b60229600a",
+   "metadata": {},
+   "source": [
+    "## Load Calibration Datset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "f3b618e9-fdee-46b2-8d7e-f11f1f7ada8d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# Calibration dataloader\n",
+    "calib_loader = get_dataset_dataloader(\n",
+    "    dataset_name=\"cnn_dailymail\",\n",
+    "    tokenizer=tokenizer,\n",
+    "    batch_size=8,\n",
+    "    num_samples=512,\n",
+    "    device=\"cuda\"\n",
+    ")\n",
+    "\n",
+    "forward_loop = create_forward_loop(dataloader=calib_loader)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f971d580-495f-4fb9-b390-0e723f6d1c18",
+   "metadata": {},
+   "source": [
+    "## Set Quatization Config and Quantize the Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7c35bf66-fe51-48ff-809b-884de91fbaf2",
+   "metadata": {},
+   "source": [
+    "Applying QAT with Model Optimizer is fairly straightforward. QAT supports the same quantization formats as the PTQ workflow, including key formats such as FP8, NVFP4, MXFP4, INT8, and INT4. In this example, we are using the default NVFP4 config which quantizes weight and activation to NVFP4 format. \n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "51c0c1bb-2804-45ae-873f-e33388458e04",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Registered <class 'transformers.models.llama.modeling_llama.LlamaAttention'> to _QuantAttention for KV Cache quantization\n",
+      "Inserted 771 quantizers\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████████████████████| 64/64 [01:14<00:00,  1.16s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import modelopt.torch.quantization as mtq\n",
+    "\n",
+    "config = mtq.NVFP4_DEFAULT_CFG\n",
+    "\n",
+    "# # Define forward loop for calibration\n",
+    "# def forward_loop(model):\n",
+    "#     for data in calib_set:\n",
+    "#         model(data)\n",
+    "\n",
+    "# quantize the model and prepare for QAT\n",
+    "model = mtq.quantize(model, config, forward_loop) \n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c1a15f93-ee06-42a5-ab3b-ca3428a62fe7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import modelopt.torch.opt as mto\n",
+    "torch.save(mto.modelopt_state(model), \"modelopt_quantizer_states.pt\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "95411f4c-b1d3-4e82-9afb-2608bd21a9a4",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "SyntaxError",
+     "evalue": "incomplete input (1262456024.py, line 1)",
+     "output_type": "error",
+     "traceback": [
+      "  \u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[14]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[31m    \u001b[39m\u001b[31mtrainer = Trainer(model=model, processing_class=tokenizer, args=training_args, **data_module\u001b[39m\n                                                                                                ^\n\u001b[31mSyntaxError\u001b[39m\u001b[31m:\u001b[39m incomplete input\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer = Trainer(model=model, processing_class=tokenizer, args=training_args, **data_module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5ff221a-d807-450b-a099-6481cb3b00d0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/llm_qat/notebooks/requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+torchaudio
+nvidia-modelopt[all]
+transformers
+ipywidgets
+jupyterlab