RHAIENG-332: (feat): add basic llm compressor test (opendatahub-io#1972)

Author: atheo89

jupyter/pytorch+llmcompressor/test/test_notebook.ipynb

@@ -0,0 +1,80 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc39f0e9",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import unittest\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from llmcompressor.modifiers.quantization import QuantizationModifier\n",
+    "from llmcompressor.transformers import oneshot\n",
+    "\n",
+    "class TestQuantizationProcess(unittest.TestCase):\n",
+    "\n",
+    "    def test_quantization_and_save(self):\n",
+    "        # Load model\n",
+    "        model_stub = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\"\n",
+    "        model_name = model_stub.split(\"/\")[-1]\n",
+    "\n",
+    "        model = AutoModelForCausalLM.from_pretrained(\n",
+    "            model_stub,\n",
+    "            torch_dtype=\"auto\",\n",
+    "        )\n",
+    "\n",
+    "        tokenizer = AutoTokenizer.from_pretrained(model_stub)\n",
+    "\n",
+    "        # Configure the quantization algorithm and scheme\n",
+    "        recipe = QuantizationModifier(\n",
+    "            targets=\"Linear\",\n",
+    "            scheme=\"FP8_DYNAMIC\",\n",
+    "            ignore=[\"lm_head\"],\n",
+    "        )\n",
+    "\n",
+    "        # Apply quantization\n",
+    "        oneshot(\n",
+    "            model=model,\n",
+    "            recipe=recipe,\n",
+    "        )\n",
+    "\n",
+    "        # Save to disk in compressed-tensors format\n",
+    "        save_path = model_name + \"-FP8-dynamic\"\n",
+    "        model.save_pretrained(save_path)\n",
+    "        tokenizer.save_pretrained(save_path)\n",
+    "\n",
+    "        # Assertions to verify save\n",
+    "        self.assertTrue(os.path.exists(save_path), f\"Save path does not exist: {save_path}\")\n",
+    "        self.assertTrue(os.path.exists(os.path.join(save_path, \"config.json\")), \"Model config not found\")\n",
+    "        self.assertTrue(os.path.exists(os.path.join(save_path, \"tokenizer_config.json\")), \"Tokenizer config not found\")\n",
+    "\n",
+    "unittest.main(argv=[''], verbosity=2, exit=False)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7bc63bb",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}