Implement 4-bit quant for train script and inference

sia/llm_engine/qwq_llm_engine.ipynb

@@ -12,7 +12,7 @@
       "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
       "Unsloth: Failed to patch Gemma3ForConditionalGeneration.\n",
       "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
-      "INFO 04-23 16:36:10 [__init__.py:239] Automatically detected platform cuda.\n"
+      "INFO 04-23 16:42:57 [__init__.py:239] Automatically detected platform cuda.\n"
      ]
     }
    ],
@@ -93,7 +93,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a99acfd42fa645abae77144125a734b7",
+       "model_id": "4f302eb9995d47fab2aa6339aa00a8d8",
        "version_major": 2,
        "version_minor": 0
       },
@@ -109,7 +109,7 @@
     "# Load model\n",
     "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
     "    model_path,\n",
-    "    gpu_memory_utilization = 0.8, # Reduce if out of memory\n",
+    "    gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
     ")"
    ]
   },
@@ -117,6 +117,53 @@
    "cell_type": "code",
    "execution_count": 6,
    "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Qwen2ForCausalLM(\n",
+       "  (model): Qwen2Model(\n",
+       "    (embed_tokens): Embedding(152064, 5120, padding_idx=151654)\n",
+       "    (layers): ModuleList(\n",
+       "      (0-63): 64 x Qwen2DecoderLayer(\n",
+       "        (self_attn): Qwen2Attention(\n",
+       "          (q_proj): Linear4bit(in_features=5120, out_features=5120, bias=True)\n",
+       "          (k_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
+       "          (v_proj): Linear4bit(in_features=5120, out_features=1024, bias=True)\n",
+       "          (o_proj): Linear4bit(in_features=5120, out_features=5120, bias=False)\n",
+       "          (rotary_emb): LlamaRotaryEmbedding()\n",
+       "        )\n",
+       "        (mlp): Qwen2MLP(\n",
+       "          (gate_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
+       "          (up_proj): Linear4bit(in_features=5120, out_features=27648, bias=False)\n",
+       "          (down_proj): Linear4bit(in_features=27648, out_features=5120, bias=False)\n",
+       "          (act_fn): SiLU()\n",
+       "        )\n",
+       "        (input_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
+       "        (post_attention_layernorm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
+       "      )\n",
+       "    )\n",
+       "    (norm): Qwen2RMSNorm((5120,), eps=1e-05)\n",
+       "    (rotary_emb): LlamaRotaryEmbedding()\n",
+       "  )\n",
+       "  (lm_head): Linear(in_features=5120, out_features=152064, bias=False)\n",
+       ")"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# enable unsloth optimizations\n",
+    "FastLanguageModel.for_inference(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
    "outputs": [
     {
      "name": "stderr",
@@ -138,7 +185,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -147,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -160,7 +207,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -173,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -185,7 +232,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -200,7 +247,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -209,7 +256,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -223,15 +270,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "<write_stdout>\n",
-      "Hello! I'm just a large language model, so I don't have feelings, but thank you for asking. How can I assist you today?</write_stdout><|im_end|>"
+      "<write_stdout>Hi, I'm an AI assistant. I don't have feelings, but I'm here to help you. How can I assist you today?</write_stdout><|im_end|>"
      ]
     }
    ],
@@ -242,7 +288,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [

sia/llm_engine/qwq_llm_engine.py

@@ -6,8 +6,6 @@ from threading import Thread
 from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
 from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
-import os
-import torch
 
 from . import LlmEngine
 from .. import util
@@ -38,10 +36,7 @@ class QwQLlmEngine(LlmEngine):
         # Load model
         self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
             model_path,
-            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,
+            gpu_memory_utilization = 0.5, # Reduce if out of memory
         )
 
         # enable unsloth optimizations