Added export checkpoint steps in QAT Notebook

Author: farshadghodsian

examples/llm_qat/notebooks/QAT_Walkthrough.ipynb

@@ -13,7 +13,7 @@
    "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",
+    "**Quantization Aware Training (QAT)** is a method that learn the effects of quantization during neural network post-training to preserve accuracy when deploying models in very-low-precision formats. QAT inserts quantizer 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."
    ]
@@ -145,12 +145,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "b6af94af-1de6-4cb1-959b-98fb3f4e1932",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AutoConfig, Mxfp4Config\n",
+    "from transformers import AutoConfig\n",
     "from trl import ModelConfig\n",
     "\n",
     "model_args = ModelConfig(\n",
@@ -165,15 +165,7 @@
     "    \"torch_dtype\": model_args.torch_dtype,\n",
     "    \"use_cache\": False,\n",
     "    \"device_map\": \"auto\",\n",
-    "}\n",
-    "\n",
-    "# Dequantize if the model is in MXFP4 format (for gpt-oss family of models)\n",
-    "config = AutoConfig.from_pretrained(model_args.model_name_or_path)\n",
-    "if (\n",
-    "    getattr(config, \"quantization_config\", {})\n",
-    "    and config.quantization_config.get(\"quant_method\", None) == \"mxfp4\"\n",
-    "):\n",
-    "    model_kwargs[\"quantization_config\"] = Mxfp4Config(dequantize=True)"
+    "}"
    ]
   },
   {
@@ -188,23 +180,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "5427ffdc-ee1f-4a81-b30b-ee06d978c4fe",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "467ffc9ed8034a819d4c9ab2bdbf4dd6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:03<00:00,  1.29it/s]\n"
+     ]
     }
    ],
    "source": [
@@ -229,7 +214,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "ac8a8e2a-fba1-4377-a528-da7808e71cfd",
    "metadata": {},
    "outputs": [],
@@ -256,7 +241,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "8eac6b5a-185c-4ea8-820e-2b3fcb593077",
    "metadata": {},
    "outputs": [],
@@ -282,7 +267,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "id": "0bf60614-99a0-48b0-85a8-1d88cd7c72ba",
    "metadata": {},
    "outputs": [],
@@ -322,64 +307,7 @@
    "execution_count": 10,
    "id": "05fdfbed-43ff-4f85-8084-a5cace7ca8ab",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "87d380ce00d646df98f57b8aecdaa0bc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Tokenizing train dataset:   0%|          | 0/900 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d638a968356f4819a1c090bf8b0e1336",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Truncating train dataset:   0%|          | 0/900 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0ee749ab465a4c33814e27bb56518e8f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Tokenizing eval dataset:   0%|          | 0/100 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a38242fe7e894579a70d1808574ae361",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Truncating eval dataset:   0%|          | 0/100 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "from trl import SFTTrainer\n",
     "\n",
@@ -413,7 +341,7 @@
     "\n",
     "import modelopt.torch.quantization as mtq\n",
     "\n",
-    "# MXFP4_MLP_WEIGHT_ONLY_CFG doesn't need calibration, but other quantization configurations may require it.\n",
+    "# Some configs don't need calibration, but other quantization configurations may require it.\n",
     "quantization_config = mtq.NVFP4_DEFAULT_CFG\n",
     "calib_size = 128\n",
     "\n",
@@ -438,7 +366,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "id": "b8611ea6-526f-4761-b456-8340abf56d0a",
    "metadata": {
     "scrolled": true
@@ -550,119 +478,141 @@
    "id": "99b10dfe-dacc-4d8d-96ea-afb6ac9e5bc7",
    "metadata": {},
    "source": [
-    "Once you have completed the above QAT workflow you should now have a model in the checkpoint folder `./llama3.1-8b-instruct-multilingual-reasoner/checkpoint-45` which contains the model files including the checkpoints and tokenizer. You can use this folder to serve the QAT NVFP4 model in TensorRT-LLM via Docker.\n",
-    "\n",
-    "**TRTLLM-Serve command**\n",
-    "\n",
-    "First we will save the trtllm-serve command to launch the inference server as an environment variabel:"
+    "Once you have completed the above QAT workflow you should now have a model in the checkpoint folder `./llama3.1-8b-instruct-multilingual-reasoner/checkpoint-45` which contains the model files including the checkpoints and tokenizer. You can use this folder to serve the QAT NVFP4 model in TensorRT-LLM via Docker."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "3994c997-e1a8-472e-ada7-96cf97156dd0",
+   "cell_type": "markdown",
+   "id": "c40a9b0a-49ba-4860-80d6-eabe06d06e5e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "env: TRTLLM_SERVE_CMD=trtllm-serve /app/tensorrt_llm/qat/checkpoint-45 --max_batch_size 1 --max_num_tokens 1024 --max_seq_len 4096 --tp_size 8 --pp_size 1 --host 0.0.0.0 --port 8000 --kv_cache_free_gpu_memory_fraction 0.95\n"
-     ]
-    }
-   ],
    "source": [
-    "%env TRTLLM_SERVE_CMD=trtllm-serve /app/tensorrt_llm/qat/checkpoint-45 --max_batch_size 1 --max_num_tokens 1024 --max_seq_len 4096 --tp_size 8 --pp_size 1 --host 0.0.0.0 --port 8000 --kv_cache_free_gpu_memory_fraction 0.95"
+    "## Running TensorRT-LLM Docker"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "0bcf77bb-b8ac-4969-a11e-c697e3b5760b",
    "metadata": {},
    "source": [
-    "**Launching the Docker container**\n",
+    "The easiest way to get started with TensorRT-LLMM is to run a TensorRT-LLM docker container. Visit the [NGC TensorRT-LLM Release page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) to find the most up-to-date NGC container image to use.\n",
     "\n",
-    "Next we will use this to run a TensorRT-LLM docker container to server the model. Visit the [NGC TensorRT-LLM Release page](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) to find the most up-to-date NGC container image to use.\n",
-    "\n",
-    "Run the following Docker command to start the TensorRT-LLM container with the above trtllm-serve command to launch the server in the background (change the image tag to match latest release):"
+    "Open a new bash shell and run the following Docker command to start the TensorRT-LLM container in interactive mode (change the image tag to match latest release):"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 60,
-   "id": "6e8dfc5d-3d41-4c14-9c9a-27b80002c320",
+   "execution_count": null,
+   "id": "2b736ebf-5970-4157-8fe1-c2242d3c3950",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "9b8741a1681ab6e054f533ea40cd65b51967783b41611fdb722df188ec77d31c\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "%%sh\n",
-    "docker run --rm --ipc=host -d \\\n",
+    "docker run --rm --ipc=host -it \\\n",
     "  --ulimit stack=67108864   --ulimit memlock=-1 \\\n",
     "  --gpus all   -p 8000:8000   -e TRTLLM_ENABLE_PDL=1 \\\n",
     "  -v ~/.cache:/root/.cache:rw --name tensorrt_llm \\\n",
-    "  -v <path-to-qat-model-checkpoint>/llama3.1-8b-instruct-multilingual-reasoner/:/app/tensorrt_llm/qat \\\n",
-    "  nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc1  $TRTLLM_SERVE_CMD"
+    "  -v $(pwd)/llama3.1-8b-instruct-multilingual-reasoner/:/app/tensorrt_llm/qat \\\n",
+    "  nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc2  /bin/bash"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0509e8ee-519d-4905-8663-c49e02712403",
+   "metadata": {},
+   "source": [
+    "## Exporting Quantized Model for deployment\n",
+    "Before deploying the model with TensorRT-LLM you will need to export the model checkpoint files. This is similar to the step you take for a quantized PTQ Model. To export the unified Hugging Face checkpoints, which can be deployed on TensorRT-LLM Pytorch, vLLM and SGLang you will need to run the [huggingface_example.sh](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/examples/llm_ptq/scripts/huggingface_example.sh) script found in the TensorRT Model Optimizer repo. "
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "c6369960-ec55-4073-922f-b9bc43ce5c08",
+   "id": "515f0e2b-0b27-44b7-bec7-2fcda65df140",
+   "metadata": {},
+   "source": [
+    "**Clone the TensorRT-LLM Model Optimizer repo inside the docker container**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "363b64fc-6d18-4b40-bd65-e90e55305b03",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "## Validating TensorRT-LLM Server"
+    "git clone https://github.com/NVIDIA/TensorRT-Model-Optimizer.git"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "d1a8dd26-1a99-41bb-8e22-9da9435834ed",
+   "id": "67dbd985-f7ae-4f33-a450-64f4ce5b2c4a",
    "metadata": {},
    "source": [
-    "You can check the Docker logs to see when the TensorRT-LLM server is ready to serve requests."
+    "**Install Model Opt prerequisites**"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d495aeec-7281-4d07-8cc0-38383b124da4",
+   "id": "44be970c-3567-4108-8009-25b6312c1eb8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cd TensorRT-Model-Optimizer/\n",
+    "pip install -e ."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6962b7ef-d8fc-4f71-a841-82045774df2d",
+   "metadata": {},
+   "source": [
+    "**Run HuggingFace checkpoint conversion script**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cff980a0-b6e6-439d-9ee6-633e9978f97d",
    "metadata": {
     "scrolled": true
    },
    "outputs": [],
    "source": [
-    "!docker logs tensorrt_llm"
+    "# set export path for converted checkpoints. The script saves the converted checkpoint in ${ROOT_SAVE_PATH}/saved_models_${MODEL_FULL_NAME}\n",
+    "export ROOT_SAVE_PATH=/app/tensorrt_llm\n",
+    "\n",
+    "# run conversion script\n",
+    "cd ..\n",
+    "bash TensorRT-Model-Optimizer/examples/llm_ptq/scripts/huggingface_example.sh --model $(pwd)/qat/checkpoint-45/ --quant nvfp4 --export_fmt hf"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "b1cfeb9a-0a60-47b0-a37f-2861ac6b0a2e",
+   "id": "156b1088-7cc6-4fcc-9e4d-0f6dd464bb20",
    "metadata": {},
    "source": [
-    "You can also check the TensorRT-LLM endpoint health to see if it is ready. It should return a **\"200\"** status if it is ready."
+    "## Serving the Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "01b5b230-47af-49bb-8227-d8c5d676e685",
+   "metadata": {},
+   "source": [
+    "Run the following trtllm-serve command to launch the inference server"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 63,
-   "id": "03d9cd22-20ab-415d-b6bc-f19470d9860c",
+   "execution_count": null,
+   "id": "3994c997-e1a8-472e-ada7-96cf97156dd0",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Status: 200\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "!curl -s -o /dev/null -w \"Status: %{http_code}\\n\" \"http://localhost:8000/health\""
+    "trtllm-serve /app/tensorrt_llm/saved_models_checkpoint-45_dense_nvfp4_tp1_pp1 \\\n",
+    "  --max_batch_size 1 --max_num_tokens 1024 \\\n",
+    "  --max_seq_len 4096 --tp_size 8 --pp_size 1 \\\n",
+    "  --host 0.0.0.0 --port 8000 \\\n",
+    "  --kv_cache_free_gpu_memory_fraction 0.95"
    ]
   },
   {
@@ -678,7 +628,7 @@
    "id": "8fba7c89-1ce0-4f7c-822e-b72786a25199",
    "metadata": {},
    "source": [
-    "Run the below example curl command to send an inference request to the server."
+    "In another terminal run the below example curl command to send an inference request to the server."
    ]
   },
   {
@@ -705,7 +655,6 @@
     }
    ],
    "source": [
-    "%%sh\n",
     "curl localhost:8000/v1/chat/completions -H \"Content-Type: application/json\" -d '{\n",
     "    \"model\": \"meta-llama/Llama-3.1-8B-Instruct-QAT\",\n",
     "    \"messages\": [\n",