From 9b2c2ece70c6a064686349e10876837178d9ca30 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Wed, 23 Apr 2025 19:38:56 +0000
Subject: [PATCH] Flash attention in notebook

---
 sia/llm_engine/qwq_llm_engine.ipynb | 170 +++++++---------------------
 1 file changed, 38 insertions(+), 132 deletions(-)

diff --git a/sia/llm_engine/qwq_llm_engine.ipynb b/sia/llm_engine/qwq_llm_engine.ipynb
index 037343c..faccc31 100644
--- a/sia/llm_engine/qwq_llm_engine.ipynb
+++ b/sia/llm_engine/qwq_llm_engine.ipynb
@@ -1,21 +1,19 @@
 {
  "cells": [
   {
-   "cell_type": "code",
-   "execution_count": 1,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "🦥 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:42:57 [__init__.py:239] Automatically detected platform cuda.\n"
-     ]
-    }
-   ],
+   "source": [
+    "source /root/venvs/sia/bin/activate\n",
+    "apt-get update && apt-get install -y cuda-toolkit\n",
+    "pip install flash-attn --no-build-isolation"
+   ]
+  },
+  {
+   "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"
@@ -23,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,126 +33,48 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "temperature = 0.6\n",
-    "model_path = \"/root/models/notebook_merged_4bit\"\n",
+    "model_path = \"/root/models/current\"\n",
     "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Load tokenizer\n",
     "tokenizer = AutoTokenizer.from_pretrained(\n",
     "    model_path,\n",
+    "    legacy=False,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "==((====))==  Unsloth 2025.3.19: Fast Qwen2 patching. Transformers: 4.51.3. vLLM: 0.8.2.\n",
-      "   \\\\   /|    NVIDIA RTX 6000 Ada Generation. Num GPUs = 1. Max memory: 47.5 GB. Platform: Linux.\n",
-      "O^O/ \\_/ \\    Torch: 2.6.0+cu124. CUDA: 8.9. CUDA Toolkit: 12.4. Triton: 3.2.0\n",
-      "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.29.post2. FA2 = False]\n",
-      " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
-      "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Sliding Window Attention is enabled but not implemented for `eager`; unexpected results may be encountered.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "validate_env\n",
-      "device_map sequential\n",
-      "validate_env\n",
-      "device_map OrderedDict([('', 0)])\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4f302eb9995d47fab2aa6339aa00a8d8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Load model\n",
     "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
     "    model_path,\n",
     "    gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
+    "    load_in_4bit=True,\n",
+    "    attn_implementation=\"flash_attention_2\",\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "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"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# enable unsloth optimizations\n",
     "FastLanguageModel.for_inference(model)"
@@ -162,30 +82,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Device set to use cuda:0\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
+    "import torch\n",
     "# 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",
+    "    torch_dtype=torch.float16,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -194,7 +108,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -207,7 +121,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -220,7 +134,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -232,7 +146,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -247,7 +161,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -256,7 +170,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -270,17 +184,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "<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|>"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "for text in streamer:\n",
     "    print(text, end=\"\")"
@@ -288,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -301,7 +207,7 @@
   "kernelspec": {
    "display_name": "sia",
    "language": "python",
-   "name": "python3"
+   "name": "sia"
   },
   "language_info": {
    "codemirror_mode": {