Use Unsloth for QwQ inference

scripts/install.sh

@@ -17,4 +17,5 @@ if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
     echo "Installing SIA core..."
     python3 -m venv /root/venvs/sia
     /root/venvs/sia/bin/pip install -e /root/sia
+    /root/venvs/sia/bin/ipython kernel install --name=sia
 fi

setup.py

@@ -10,12 +10,13 @@ setup(
         ],
     },
     install_requires=[
-        'torch>=2.0.0',
         'accelerate>=0.26.0',
         'aiohttp>=3.8.0',
-        'bitsandbytes>=0.41.0',
+        'bitsandbytes>=0.45',
         'dotenv-python>=0.0.1',
         'huggingface_hub>=0.16.0',
+        'ipykernel>=6.0.0', 
+        'ipywidgets>=8.0.0',
         'lxml>=4.9.0',
         'mistral-common>=1.0.0',
         'mistralai>=0.0.7',
@@ -23,7 +24,11 @@ setup(
         'psutil>=5.9.0',
         'python-dotenv>=1.0.0',
         'tiktoken>=0.4.0',
+        'torch>=2.0.0',
         'transformers>=4.30.0',
+        'trl>=0.7.8',
+        'unsloth>=2025.3',
+        'vllm==0.8.2',
         'xml_schema_validator @ file:///root/sia/lib/xml_schema_validator',
     ],
     python_requires='>=3.10',

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
@@ -27,39 +30,26 @@ class QwQLlmEngine(LlmEngine):
         """
         self._temperature = temperature
 
-        # Configure 4-bit quantization for massive memory savings
+        # Load tokenizer
-        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
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
-            padding_side="left",
-            trust_remote_code=True,
         )
         
-        # Load model with 4-bit quantization
+        # Load model
-        self._model = AutoModelForCausalLM.from_pretrained(
+        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
             model_path,
-            device_map="auto",
+            load_in_4bit = True, # False for LoRA 16bit
-            quantization_config=quantization_config,
+            fast_inference = True, # Enable vLLM fast inference
-            torch_dtype=torch.bfloat16,
+            gpu_memory_utilization = 0.8, # Reduce if out of memory
-            low_cpu_mem_usage=True,
+            tokenizer = self._tokenizer,
-            trust_remote_code=True,
         )
         
-        # 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,
         }

tools/train/setup.py

@@ -9,7 +9,7 @@ setup(
     ],
 
     install_requires=[
-        'accelerate>=0.25.0',
+        'accelerate>=0.26.0',
         'bitsandbytes>=0.45.0',
         'black>=22.0.0',
         'datasets>=2.14.6',
@@ -18,7 +18,6 @@ setup(
         'ipykernel>=6.0.0', 
         'ipywidgets>=8.0.0',
         'peft>=0.8.0',
-        'peft>=0.8.0',
         'pytest-cov>=4.0.0',
         'pytest>=7.0.0',
         'pyyaml>=6.0',

tools/train/train/qwq.ipynb

@@ -0,0 +1,208 @@
+{
+ "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, is_bfloat16_supported"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dataclasses import dataclass\n",
+    "from pathlib import Path\n",
+    "from transformers import AutoTokenizer, TrainingArguments\n",
+    "from trl import SFTTrainer, DataCollatorForCompletionOnlyLM\n",
+    "from typing import Optional, List\n",
+    "import argparse\n",
+    "import json\n",
+    "import os"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from train import qwq"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "args = qwq.Args([\"--output-dir\", \"/root/models/notebook\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = qwq.Dataset(args.config_path)\n",
+    "dataset.validate()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_seq_length = 2048 # Can increase for longer reasoning traces\n",
+    "lora_rank = 64 # Larger rank = smarter, but slower"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:\n",
+    "    tokenizer_config = json.load(f)\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    args.base_model,\n",
+    "    **tokenizer_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name = args.base_model,\n",
+    "    max_seq_length = max_seq_length,\n",
+    "    load_in_4bit = True, # False for LoRA 16bit\n",
+    "    fast_inference = True, # Enable vLLM fast inference\n",
+    "    max_lora_rank = lora_rank,\n",
+    "    gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
+    "    tokenizer = tokenizer,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128\n",
+    "    target_modules = [\n",
+    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
+    "        \"gate_proj\", \"up_proj\", \"down_proj\",\n",
+    "    ], # Remove QKVO if out of memory\n",
+    "    lora_alpha = lora_rank,\n",
+    "    lora_dropout = 0, # Supports any, but = 0 is optimized\n",
+    "    bias = \"none\",    # Supports any, but = \"none\" is optimized\n",
+    "    use_gradient_checkpointing = \"unsloth\", # True or \"unsloth\" for very long context\n",
+    "    random_state = 3407,\n",
+    "    use_rslora = False,  # We support rank stabilized LoRA\n",
+    "    loftq_config = None, # And LoftQ\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "training_args = TrainingArguments(\n",
+    "    output_dir=str(args.output_dir),\n",
+    "    num_train_epochs=3,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=16,\n",
+    "    gradient_checkpointing=True,\n",
+    "    learning_rate=2e-5,\n",
+    "    lr_scheduler_type=\"cosine\",\n",
+    "    warmup_ratio=0.05,\n",
+    "    weight_decay=0.01,\n",
+    "    fp16=not is_bfloat16_supported(),\n",
+    "    bf16=is_bfloat16_supported(),\n",
+    "    logging_steps=10,\n",
+    "    save_steps=200,\n",
+    "    save_total_limit=3,\n",
+    "    report_to=\"none\",\n",
+    "    optim=\"adamw_8bit\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = SFTTrainer(\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    args=training_args,\n",
+    "    train_dataset=dataset.to_transformers_dataset(tokenizer),\n",
+    "    dataset_text_field=\"messages\",\n",
+    "    max_seq_length=max_seq_length,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save_pretrained_merged(\n",
+    "    str(args.output_dir), \n",
+    "    tokenizer=tokenizer,\n",
+    "    #save_method=\"merged_4bit_forced\"\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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
+}