Test train qwen 3

tools/train/train/qwen.ipynb

@@ -0,0 +1,316 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4c7c96e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install transformers>=4.51.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3edd43ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.\n",
+    "%pip install torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d41f8851",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install unsloth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1b5a6da6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5256285",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unsloth import FastLanguageModel, is_bfloat16_supported"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "88abe86a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, TrainingArguments\n",
+    "from trl import SFTTrainer, DataCollatorForCompletionOnlyLM"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b404a9db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01519eee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"unsloth/Qwen3-0.6B-unsloth-bnb-4bit\"\n",
+    "#model_name = \"Qwen/Qwen3-0.6B\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "04e1aad4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# load the tokenizer and the model\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31f2451c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name,\n",
+    "    gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2dddcb05",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare the model input\n",
+    "prompt = \"Give me a short introduction to large language model.\"\n",
+    "messages = [\n",
+    "    {\"role\": \"user\", \"content\": prompt}\n",
+    "]\n",
+    "text = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    tokenize=False,\n",
+    "    add_generation_prompt=True,\n",
+    "    enable_thinking=True # Switches between thinking and non-thinking modes. Default is True.\n",
+    ")\n",
+    "model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84a7b4cc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## conduct text completion\n",
+    "#generated_ids = model.generate(\n",
+    "#    **model_inputs,\n",
+    "#    max_new_tokens=32768\n",
+    "#)\n",
+    "#output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "05be1a53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## parsing thinking content\n",
+    "#try:\n",
+    "#    # rindex finding 151668 (</think>)\n",
+    "#    index = len(output_ids) - output_ids[::-1].index(151668)\n",
+    "#except ValueError:\n",
+    "#    index = 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "64775c31",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip(\"\\n\")\n",
+    "#content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip(\"\\n\")\n",
+    "#\n",
+    "#print(\"thinking content:\", thinking_content)\n",
+    "#print(\"content:\", content)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "48477dfd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = dataset.Dataset(\"/root/sia/training/config.yaml\")\n",
+    "dataset.validate()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90b56737",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    target_modules = [\n",
+    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
+    "        \"gate_proj\", \"up_proj\", \"down_proj\",\n",
+    "    ],\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",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "21422429",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "training_args = TrainingArguments(\n",
+    "    output_dir=\"/root/models/qwen_train\",\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,
+   "id": "83be036e",
+   "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",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "886ba936",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13071f43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save_pretrained_merged(\n",
+    "    \"/root/models/qwen_merged_4bit\", \n",
+    "    tokenizer=tokenizer,\n",
+    "    save_method=\"merged_4bit_forced\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a80162e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install vllm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a934ce7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%vllm serve /root/models/qwen_merged_4bit --enable-reasoning --reasoning-parser deepseek_r1"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "notebook",
+   "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": 5
+}