wip vllm

sia/llm_engine/qwq_vllm.ipynb

@@ -0,0 +1,255 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# vLLM Streaming Implementation\n",
+    "\n",
+    "This notebook demonstrates how to implement streaming capability with vLLM, comparable to the unsloth implementation.\n",
+    "\n",
+    "First, let's make sure we have vLLM installed:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:31 [__init__.py:239] Automatically detected platform cuda.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "from vllm import SamplingParams\n",
+    "from transformers import AutoTokenizer\n",
+    "import sys"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temperature = 0.6\n",
+    "model_path = \"/root/models/current\"\n",
+    "xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_path,\n",
+    "    legacy=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "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": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    tokenize=False,\n",
+    "    add_generation_prompt=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define sampling parameters\n",
+    "sampling_params = SamplingParams(\n",
+    "    temperature=temperature,\n",
+    "    top_p=0.95,\n",
+    "    max_tokens=512,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:40 [config.py:585] This model supports multiple tasks: {'generate', 'score', 'reward', 'embed', 'classify'}. Defaulting to 'generate'.\n",
+      "WARNING 04-25 19:36:42 [config.py:664] bitsandbytes quantization is not fully optimized yet. The speed can be slower than non-quantized models.\n",
+      "WARNING 04-25 19:36:42 [arg_utils.py:1854] --quantization bitsandbytes is not supported by the V1 Engine. Falling back to V0. \n",
+      "INFO 04-25 19:36:42 [llm_engine.py:241] Initializing a V0 LLM engine (v0.8.2) with config: model='/root/models/current', speculative_config=None, tokenizer='/root/models/current', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=bitsandbytes, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=bitsandbytes, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=/root/models/current, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=None, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"max_capture_size\":256}, use_cached_outputs=False, \n",
+      "INFO 04-25 19:36:42 [cuda.py:291] Using Flash Attention backend.\n",
+      "INFO 04-25 19:36:43 [parallel_state.py:954] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
+      "INFO 04-25 19:36:43 [model_runner.py:1110] Starting to load model /root/models/current...\n",
+      "INFO 04-25 19:36:43 [loader.py:1155] Loading weights with BitsAndBytes quantization. May take a while ...\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8b9f3cb293484cac932e6cedd841c813",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "54f8aa5eefdb43d8bc07274044a8bc1c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:36:51 [model_runner.py:1146] Model loading took 18.0523 GB and 8.113452 seconds\n",
+      "INFO 04-25 19:36:55 [worker.py:267] Memory profiling takes 3.23 seconds\n",
+      "INFO 04-25 19:36:55 [worker.py:267] the current vLLM instance can use total_gpu_memory (47.53GiB) x gpu_memory_utilization (0.90) = 42.78GiB\n",
+      "INFO 04-25 19:36:55 [worker.py:267] model weights take 18.05GiB; non_torch_memory takes 0.06GiB; PyTorch activation peak memory takes 1.59GiB; the rest of the memory reserved for KV Cache is 23.08GiB.\n",
+      "INFO 04-25 19:36:55 [executor_base.py:111] # cuda blocks: 5907, # CPU blocks: 1024\n",
+      "INFO 04-25 19:36:55 [executor_base.py:116] Maximum concurrency for 4096 tokens per request: 23.07x\n",
+      "INFO 04-25 19:36:58 [model_runner.py:1442] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Capturing CUDA graph shapes: 100%|██████████| 35/35 [01:01<00:00,  1.75s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO 04-25 19:38:00 [model_runner.py:1570] Graph capturing finished in 61 secs, took 1.98 GiB\n",
+      "INFO 04-25 19:38:00 [llm_engine.py:447] init engine (profile, create kv cache, warmup model) took 68.57 seconds\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from vllm import LLM, SamplingParams\n",
+    "import time\n",
+    "\n",
+    "# Initialize LLM\n",
+    "llm = LLM(\n",
+    "    model=model_path,\n",
+    "    tensor_parallel_size=1,\n",
+    "    max_model_len=4096,\n",
+    "    quantization=\"bitsandbytes\",\n",
+    "    load_format=\"bitsandbytes\",\n",
+    "    trust_remote_code=True,\n",
+    "    # Enable streaming\n",
+    "    enable_lora=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting generation with token-by-token output:\n",
+      "<think>\n",
+      "Okay, the user greeted me with \"Hi, how are you?\" I need to respond appropriately. Let me see... The instructions say to always use the <write_stdout> XML tag. So first, I should acknowledge their greeting and state that I'm an AI, then ask how I can assist them. Keep it friendly and helpful. Let me make sure I don't add any extra information beyond that. Just a simple response. Alright, that should work.\n",
+      "</think>\n",
+      "\n",
+      "<write_stdout>\n",
+      "Hello! I'm just a computer program, but I'm here to help you. How can I assist you today?\n",
+      "</write_stdout>"
+     ]
+    }
+   ],
+   "source": [
+    "previous_text = \"\"\n",
+    "print(\"Starting generation with token-by-token output:\")\n",
+    "\n",
+    "# Try with direct iteration over the generator\n",
+    "for output in llm.generate(prompt, sampling_params, use_tqdm=False):\n",
+    "    if hasattr(output, 'outputs') and output.outputs and len(output.outputs) > 0:\n",
+    "        generated_text = output.outputs[0].text\n",
+    "        if len(generated_text) > len(previous_text):\n",
+    "            new_text = generated_text[len(previous_text):]\n",
+    "            sys.stdout.write(new_text)\n",
+    "            sys.stdout.flush()\n",
+    "            previous_text = generated_text"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sia",
+   "language": "python",
+   "name": "sia"
+  },
+  "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
+}