Use Unsloth for QwQ inference

sia/llm_engine/qwq_llm_engine.ipynb

@@ -0,0 +1,218 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
+    "from unsloth import FastLanguageModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "from threading import Thread\n",
+    "from transformers import AutoTokenizer, TextIteratorStreamer, pipeline\n",
+    "from xml_schema_validator import XmlLogitsProcessor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temperature = 0.6\n",
+    "model_path = \"/root/models/notebook\"\n",
+    "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_path,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load model\n",
+    "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_path,\n",
+    "    load_in_4bit = True, # False for LoRA 16bit\n",
+    "    fast_inference = True, # Enable vLLM fast inference\n",
+    "    gpu_memory_utilization = 0.8, # Reduce if out of memory\n",
+    "    tokenizer = tokenizer,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create inference pipeline with memory-efficient settings\n",
+    "pipeline = pipeline(\n",
+    "    \"text-generation\",\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    return_full_text=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    tokenize=False,\n",
+    "    add_generation_prompt=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "streamer = TextIteratorStreamer(\n",
+    "    tokenizer,\n",
+    "    skip_prompt=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_kwargs = {\n",
+    "    \"text_inputs\": text,\n",
+    "    \"do_sample\": True,\n",
+    "    \"temperature\": temperature,\n",
+    "    \"streamer\": streamer,\n",
+    "    \"use_cache\": True,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_kwargs[\"logits_processor\"] = [logits_processor.copy()]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generation_thread = Thread(\n",
+    "    target=pipeline,\n",
+    "    kwargs=generation_kwargs\n",
+    ")\n",
+    "\n",
+    "generation_thread.start()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for text in streamer:\n",
+    "    print(text, end=\"\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
+      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
+      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
+      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "generation_thread.join()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sia",
+   "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
+}

sia/llm_engine/qwq_llm_engine.py

@@ -1,6 +1,9 @@
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel
+
 from pathlib import Path
 from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline, BitsAndBytesConfig
+from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
 from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
 import os
@@ -26,40 +29,27 @@ class QwQLlmEngine(LlmEngine):
             xml_schema_text: Optional XML schema to validate against
         """
         self._temperature = temperature
-                
-        # Configure 4-bit quantization for massive memory savings
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4"
-        )
 
-        # Load tokenizer first - this uses minimal memory
+        # Load tokenizer
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
-            padding_side="left",
-            trust_remote_code=True,
         )
         
-        # Load model with 4-bit quantization
-        self._model = AutoModelForCausalLM.from_pretrained(
+        # Load model
+        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
             model_path,
-            device_map="auto",
-            quantization_config=quantization_config,
-            torch_dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
+            load_in_4bit = True, # False for LoRA 16bit
+            fast_inference = True, # Enable vLLM fast inference
+            gpu_memory_utilization = 0.8, # Reduce if out of memory
+            tokenizer = self._tokenizer,
         )
         
-        # Create inference pipeline with memory-efficient settings
+        # Create inference pipeline
         self._pipeline = pipeline(
             "text-generation",
             model=self._model,
             tokenizer=self._tokenizer,
             return_full_text=False,
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
         )
 
         if xml_schema_text:
@@ -102,7 +92,6 @@ class QwQLlmEngine(LlmEngine):
             "text_inputs": text,
             "do_sample": True,
             "temperature": self._temperature,
-            "max_new_tokens": self.token_limit(),
             "streamer": streamer,
             "use_cache": True,
         }