Add Colab training notebook

colab_training.ipynb

@@ -0,0 +1,344 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Qwen2.5-7B QLoRA Training on Colab\n",
+    "\n",
+    "Google Colab Pro (A100) での学習用ノートブック\n",
+    "\n",
+    "**推奨**: Colab Pro ($10/月) 以上、A100 GPU"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. 環境セットアップ"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# GPU確認\n",
+    "!nvidia-smi"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Google Driveマウント（チェックポイント保存用）\n",
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')\n",
+    "\n",
+    "# 作業ディレクトリ作成\n",
+    "!mkdir -p /content/drive/MyDrive/qwen-training/checkpoints\n",
+    "!mkdir -p /content/drive/MyDrive/qwen-training/output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 依存関係インストール\n",
+    "!pip install -q torch==2.2.0 torchvision==0.17.0\n",
+    "!pip install -q transformers==4.46.0 datasets peft==0.13.0 trl==0.11.0\n",
+    "!pip install -q bitsandbytes accelerate huggingface_hub safetensors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# HuggingFaceログイン\n",
+    "from huggingface_hub import login\n",
+    "login()  # トークンを入力"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. 設定"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 設定\n",
+    "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
+    "OUTPUT_MODEL_ID = \"hajimemat/qwen2.5-7b-glaive-fc-lora-colab\"  # 変更可\n",
+    "DATASET_NAME = \"glaiveai/glaive-function-calling-v2\"\n",
+    "\n",
+    "# Google Driveに保存\n",
+    "CHECKPOINT_DIR = \"/content/drive/MyDrive/qwen-training/checkpoints\"\n",
+    "FINAL_OUTPUT_DIR = \"/content/drive/MyDrive/qwen-training/output\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. データセット準備"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "def convert_glaive_to_chatml(example):\n",
+    "    parts = []\n",
+    "    if example.get(\"system\"):\n",
+    "        parts.append(f\"<|im_start|>system\\n{example['system']}<|im_end|>\")\n",
+    "    \n",
+    "    chat = example.get(\"chat\", \"\")\n",
+    "    if chat:\n",
+    "        current_role = None\n",
+    "        current_content = []\n",
+    "        for line in chat.split(\"\\n\"):\n",
+    "            line = line.strip()\n",
+    "            if line.startswith(\"USER:\"):\n",
+    "                if current_role and current_content:\n",
+    "                    content = \"\\n\".join(current_content).strip()\n",
+    "                    if content:\n",
+    "                        parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
+    "                current_role = \"user\"\n",
+    "                current_content = [line[5:].strip()]\n",
+    "            elif line.startswith(\"ASSISTANT:\"):\n",
+    "                if current_role and current_content:\n",
+    "                    content = \"\\n\".join(current_content).strip()\n",
+    "                    if content:\n",
+    "                        parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
+    "                current_role = \"assistant\"\n",
+    "                current_content = [line[10:].strip()]\n",
+    "            elif current_role:\n",
+    "                current_content.append(line)\n",
+    "        if current_role and current_content:\n",
+    "            content = \"\\n\".join(current_content).strip()\n",
+    "            if content:\n",
+    "                parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
+    "    return {\"text\": \"\\n\".join(parts)}\n",
+    "\n",
+    "print(f\"Loading dataset: {DATASET_NAME}\")\n",
+    "dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
+    "print(f\"Original: {len(dataset)} examples\")\n",
+    "\n",
+    "dataset = dataset.map(convert_glaive_to_chatml, remove_columns=dataset.column_names, num_proc=4)\n",
+    "dataset = dataset.filter(lambda x: len(x[\"text\"]) > 50)\n",
+    "dataset = dataset.shuffle(seed=42)\n",
+    "split = dataset.train_test_split(test_size=0.02, seed=42)\n",
+    "\n",
+    "print(f\"Train: {len(split['train'])}, Test: {len(split['test'])}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. モデル準備"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments\n",
+    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
+    "\n",
+    "# QLoRA量子化設定\n",
+    "bnb_config = BitsAndBytesConfig(\n",
+    "    load_in_4bit=True,\n",
+    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
+    "    bnb_4bit_use_double_quant=True,\n",
+    ")\n",
+    "\n",
+    "# LoRA設定\n",
+    "lora_config = LoraConfig(\n",
+    "    r=64,\n",
+    "    lora_alpha=16,\n",
+    "    lora_dropout=0.05,\n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    bias=\"none\",\n",
+    "    task_type=\"CAUSAL_LM\",\n",
+    ")\n",
+    "\n",
+    "# トークナイザー\n",
+    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)\n",
+    "tokenizer.padding_side = \"right\"\n",
+    "if tokenizer.pad_token is None:\n",
+    "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "\n",
+    "# モデル\n",
+    "print(f\"Loading model: {BASE_MODEL}\")\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    BASE_MODEL,\n",
+    "    quantization_config=bnb_config,\n",
+    "    device_map=\"auto\",\n",
+    "    attn_implementation=\"sdpa\",\n",
+    "    trust_remote_code=True,\n",
+    ")\n",
+    "\n",
+    "model = prepare_model_for_kbit_training(model)\n",
+    "model = get_peft_model(model, lora_config)\n",
+    "model.print_trainable_parameters()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. 学習実行"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import SFTTrainer\n",
+    "\n",
+    "training_args = TrainingArguments(\n",
+    "    output_dir=CHECKPOINT_DIR,\n",
+    "    num_train_epochs=1,\n",
+    "    per_device_train_batch_size=4,\n",
+    "    per_device_eval_batch_size=4,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    learning_rate=2e-4,\n",
+    "    weight_decay=0.01,\n",
+    "    warmup_ratio=0.03,\n",
+    "    lr_scheduler_type=\"cosine\",\n",
+    "    optim=\"paged_adamw_8bit\",\n",
+    "    bf16=True,\n",
+    "    logging_steps=10,\n",
+    "    save_steps=200,\n",
+    "    save_total_limit=3,\n",
+    "    eval_strategy=\"steps\",\n",
+    "    eval_steps=200,\n",
+    "    report_to=\"none\",\n",
+    "    gradient_checkpointing=True,\n",
+    "    save_safetensors=True,\n",
+    ")\n",
+    "\n",
+    "trainer = SFTTrainer(\n",
+    "    model=model,\n",
+    "    train_dataset=split[\"train\"],\n",
+    "    eval_dataset=split[\"test\"],\n",
+    "    args=training_args,\n",
+    "    peft_config=lora_config,\n",
+    "    tokenizer=tokenizer,\n",
+    "    max_seq_length=1024,\n",
+    "    packing=False,\n",
+    "    dataset_text_field=\"text\",\n",
+    ")\n",
+    "\n",
+    "# チェックポイントから再開\n",
+    "import os\n",
+    "resume_from = None\n",
+    "if os.path.exists(CHECKPOINT_DIR):\n",
+    "    checkpoints = [d for d in os.listdir(CHECKPOINT_DIR) if d.startswith(\"checkpoint-\")]\n",
+    "    if checkpoints:\n",
+    "        latest = max(checkpoints, key=lambda x: int(x.split(\"-\")[1]))\n",
+    "        resume_from = os.path.join(CHECKPOINT_DIR, latest)\n",
+    "        print(f\"Resuming from: {resume_from}\")\n",
+    "\n",
+    "# 学習開始\n",
+    "trainer.train(resume_from_checkpoint=resume_from)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. 保存とアップロード"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ローカル保存\n",
+    "print(f\"Saving to {FINAL_OUTPUT_DIR}\")\n",
+    "trainer.save_model(FINAL_OUTPUT_DIR)\n",
+    "tokenizer.save_pretrained(FINAL_OUTPUT_DIR)\n",
+    "\n",
+    "# HuggingFaceにアップロード\n",
+    "print(f\"Uploading to {OUTPUT_MODEL_ID}\")\n",
+    "try:\n",
+    "    trainer.model.push_to_hub(OUTPUT_MODEL_ID, private=True)\n",
+    "    tokenizer.push_to_hub(OUTPUT_MODEL_ID, private=True)\n",
+    "    print(f\"Done! https://huggingface.co/{OUTPUT_MODEL_ID}\")\n",
+    "except Exception as e:\n",
+    "    print(f\"Upload failed: {e}\")\n",
+    "    print(\"Model saved locally in Google Drive\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. クイックテスト（オプション）"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 簡単な推論テスト\n",
+    "from peft import PeftModel\n",
+    "\n",
+    "test_prompt = \"\"\"<|im_start|>system\n",
+    "You are a helpful assistant with access to functions.\n",
+    "<|im_end|>\n",
+    "<|im_start|>user\n",
+    "What's the weather in Tokyo?\n",
+    "<|im_end|>\n",
+    "<|im_start|>assistant\n",
+    "\"\"\"\n",
+    "\n",
+    "inputs = tokenizer(test_prompt, return_tensors=\"pt\").to(model.device)\n",
+    "outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)\n",
+    "print(tokenizer.decode(outputs[0], skip_special_tokens=False))"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "A100",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}