Use Unsloth for QwQ inference

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
+}