fix: update notebook for Colab T4 GPU training

- Configure for CUDA/T4 GPU instead of MPS
- Add base model download step
- Use Google Drive for persistence
- Fix training and merge function calls
- Update workflow for Colab environment

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

colab_train_stack29.ipynb

@@ -4,17 +4,25 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🚀 Stack 2.9 - Local Training Notebook\n",
+    "# 🚀 Stack 2.9 - Colab Training Notebook\n",
     "\n",
-    "**Training on local Mac with MPS (Apple Silicon GPU)**\n",
+    "**Zero-cost training on Google Colab free tier with T4 GPU**\n",
     "\n",
-    "This notebook trains a LoRA adapter for Stack 2.9 on **Qwen2.5-Coder-7B** using your local machine.\n",
+    "This notebook trains a LoRA adapter for Stack 2.9 on **Qwen2.5-Coder-7B** using a free T4 GPU.\n",
+    "\n",
+    "⏱️ **Expected runtime:** 3-5 hours\n",
+    "💾 **VRAM needed:** ~12GB (fits in T4's 15GB)\n",
+    "📦 **Output:** `./training_output/` (on Google Drive)\n",
     "\n",
     "---\n",
     "\n",
-    "**Prerequisites:**\n",
-    "1. Base model downloaded to `./base_model_qwen7b/`\n",
-    "2. Run from the stack-2.9 directory\n",
+    "**CRITICAL:** All data saved to **Google Drive** to persist through disconnects.\n",
+    "\n",
+    "**Instructions:**\n",
+    "1. Runtime → Change runtime type → **GPU (T4)**\n",
+    "2. Run all cells in order\n",
+    "3. **Allow** Drive access when prompted\n",
+    "\n",
     "---\n"
    ]
   },
@@ -24,32 +32,46 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Check system and GPU\n",
-    "import platform\n",
-    "import os\n",
+    "# Check GPU\n",
+    "!nvidia-smi"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1️⃣ Mount Google Drive (REQUIRED for persistence)\n",
     "\n",
-    "print(\"=\"*60)\n",
-    "print(\"SYSTEM INFO\")\n",
-    "print(\"=\"*60)\n",
-    "print(f\"Platform: {platform.platform()}\")\n",
-    "print(f\"Python: {platform.python_version()}\")\n",
-    "print(f\"Working directory: {os.getcwd()}\")\n",
+    "Click the link, allow access, copy the auth code, paste it, and press Enter.\n",
     "\n",
-    "try:\n",
-    "    import torch\n",
-    "    print(f\"PyTorch: {torch.__version__}\")\n",
-    "    print(f\"MPS available: {torch.backends.mps.is_available()}\")\n",
-    "    if torch.backends.mps.is_available():\n",
-    "        print(f\"MPS device: {torch.backends.mps.is_built()}\")\n",
-    "except ImportError:\n",
-    "    print(\"PyTorch not installed - run: pip install torch\")"
+    "**Without Drive mounting, training will be lost if Colab disconnects!**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')\n",
+    "\n",
+    "# Set up paths on Drive - ALL OUTPUT GOES HERE\n",
+    "import os\n",
+    "BASE_PATH = \"/content/drive/MyDrive/stack-2.9\"\n",
+    "os.makedirs(BASE_PATH, exist_ok=True)\n",
+    "os.chdir(BASE_PATH)\n",
+    "print(f\"\\n✅ Working directory: {os.getcwd()}\")\n",
+    "print(f\"All outputs will be saved to: {BASE_PATH}\")\n",
+    "print(\"\\nCurrent folder contents:\")\n",
+    "!ls -la"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 1️⃣ Set Working Directory"
+    "## 2️⃣ Clone Stack 2.9 Repository"
    ]
   },
   {
@@ -58,35 +80,79 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Navigate to the stack-2.9 directory\n",
-    "import os\n",
+    "# Clone into Drive if not already there\n",
+    "if not os.path.exists('stack-2.9'):\n",
+    "    !git clone https://github.com/my-ai-stack/stack-2.9.git\n",
     "\n",
-    "STACK_PATH = \"/Users/walidsobhi/.openclaw/workspace/stack-2.9\"\n",
-    "os.chdir(STACK_PATH)\n",
+    "os.chdir('stack-2.9')\n",
+    "print(f\"Now in: {os.getcwd()}\")\n",
+    "!ls -la"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3️⃣ Install Dependencies\n",
     "\n",
-    "BASE_MODEL_PATH = \"./base_model_qwen7b\"\n",
-    "DATA_PATH = \"./data/final/train.jsonl\"\n",
-    "OUTPUT_DIR = \"./training_output\"\n",
+    "Takes 5-10 minutes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install --upgrade pip\n",
+    "!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
+    "!pip install transformers==4.40.0 peft==0.10.0 accelerate bitsandbytes==0.43.0 datasets pyyaml tqdm\n",
+    "!pip install scipy\n",
+    "print(\"\\n✅ Dependencies installed\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4️⃣ Download Base Model (Qwen2.5-Coder-7B)\n",
     "\n",
-    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "This takes ~15-20 minutes. The model is ~14GB."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MODEL_NAME = \"Qwen/Qwen2.5-Coder-7B\"\n",
+    "MODEL_DIR = \"./base_model_qwen7b\"\n",
     "\n",
-    "print(f\"✅ Working directory: {os.getcwd()}\")\n",
-    "print(f\"Base model path: {BASE_MODEL_PATH}\")\n",
-    "print(f\"Data path: {DATA_PATH}\")\n",
-    "print(f\"Output dir: {OUTPUT_DIR}\")\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "import os\n",
     "\n",
-    "# Check paths exist\n",
-    "print(\"\\n📁 Checking required files:\")\n",
-    "print(f\"  - Data exists: {os.path.exists(DATA_PATH)}\")\n",
-    "if not os.path.exists(DATA_PATH):\n",
-    "    print(f\"  ⚠️ Data not found at {DATA_PATH}\")"
+    "if not os.path.exists(MODEL_DIR):\n",
+    "    print(f\"Downloading {MODEL_NAME} to {MODEL_DIR}...\")\n",
+    "    print(\"This will take 15-20 minutes...\")\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+    "    tokenizer.save_pretrained(MODEL_DIR)\n",
+    "    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+    "    model.save_pretrained(MODEL_DIR)\n",
+    "    print(f\"✅ Model saved to {MODEL_DIR}\")\n",
+    "else:\n",
+    "    print(f\"✅ Model already exists at {MODEL_DIR}\")\n",
+    "\n",
+    "!ls -lh {MODEL_DIR} | head -10"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 2️⃣ Install Dependencies"
+    "## 5️⃣ Prepare Training Data\n",
+    "\n",
+    "Using the bundled mini dataset (5K examples) for quick prototyping."
    ]
   },
   {
@@ -95,18 +161,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install required packages\n",
-    "!pip install torch torchvision torchaudio\n",
-    "!pip install transformers peft accelerate datasets pyyaml tqdm\n",
-    "!pip install scipy\n",
-    "print(\"\\n✅ Dependencies installed\")"
+    "# Check if data exists in the repo, if not create mini dataset\n",
+    "import os\n",
+    "\n",
+    "DATA_PATH = \"./data/final/train.jsonl\"\n",
+    "\n",
+    "if os.path.exists(DATA_PATH):\n",
+    "    print(f\"✅ Training data found at {DATA_PATH}\")\n",
+    "    !wc -l {DATA_PATH}\n",
+    "else:\n",
+    "    print(\"Creating mini dataset (5K examples)...\")\n",
+    "    !python scripts/create_mini_dataset.py --size 5000 --output data_mini/train_mini.jsonl\n",
+    "    DATA_PATH = \"./data_mini/train_mini.jsonl\"\n",
+    "    !ls -lh {DATA_PATH}"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 3️⃣ Prepare Training Configuration"
+    "## 6️⃣ Prepare Training Configuration"
    ]
   },
   {
@@ -115,7 +189,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Use local training config\n",
+    "# Use Colab config and update paths\n",
     "import yaml\n",
     "\n",
     "config_path = \"./stack/training/train_config_local.yaml\"\n",
@@ -123,16 +197,18 @@
     "with open(config_path, 'r') as f:\n",
     "    config = yaml.safe_load(f)\n",
     "\n",
-    "# Update paths to match our setup\n",
-    "config['model']['name'] = BASE_MODEL_PATH\n",
+    "# Update for Colab/T4 GPU\n",
+    "config['model']['name'] = MODEL_DIR\n",
     "config['data']['input_path'] = DATA_PATH\n",
-    "config['output']['lora_dir'] = f\"{OUTPUT_DIR}/lora\"\n",
-    "config['output']['merged_dir'] = f\"{OUTPUT_DIR}/merged\"\n",
-    "config['hardware']['device'] = \"mps\"\n",
+    "config['output']['lora_dir'] = \"./training_output/lora\"\n",
+    "config['output']['merged_dir'] = \"./training_output/merged\"\n",
+    "config['hardware']['device'] = \"cuda\"  # Use T4 GPU\n",
+    "config['hardware']['num_gpus'] = 1\n",
     "\n",
     "# Save updated config\n",
-    "updated_config_path = f\"{OUTPUT_DIR}/train_config.yaml\"\n",
+    "OUTPUT_DIR = \"./training_output\"\n",
     "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "updated_config_path = f\"{OUTPUT_DIR}/train_config.yaml\"\n",
     "\n",
     "with open(updated_config_path, 'w') as f:\n",
     "    yaml.dump(config, f)\n",
@@ -149,7 +225,15 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4️⃣ Check Base Model"
+    "## 7️⃣ Train LoRA Adapter\n",
+    "\n",
+    "⚠️ **This takes 3-5 hours. DO NOT INTERRUPT.**\n",
+    "\n",
+    "If Colab disconnects, reconnect and training will resume from checkpoint automatically.\n",
+    "\n",
+    "Watch for `Train loss:` decreasing. It should start ~2.0-3.0 and trend downward.\n",
+    "\n",
+    "Checkpoints saved every 200 steps to `./training_output/lora/` (on Drive)."
    ]
   },
   {
@@ -158,50 +242,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Check if base model exists\n",
     "import os\n",
+    "import sys\n",
     "\n",
-    "model_path = BASE_MODEL_PATH\n",
-    "print(f\"Checking base model at: {model_path}\")\n",
+    "# Add training module to path\n",
+    "sys.path.insert(0, './stack/training')\n",
     "\n",
-    "if os.path.exists(model_path):\n",
-    "    files = os.listdir(model_path)\n",
-    "    print(f\"✅ Base model found! {len(files)} files:\")\n",
-    "    for f in files[:10]:\n",
-    "        print(f\"   - {f}\")\n",
-    "    if len(files) > 10:\n",
-    "        print(f\"   ... and {len(files)-10} more\")\n",
-    "else:\n",
-    "    print(\"❌ Base model NOT found!\")\n",
-    "    print(\"\\nTo download Qwen2.5-Coder-7B:\")\n",
-    "    print(\"  huggingface-cli download Qwen/Qwen2.5-Coder-7B --local-dir ./base_model_qwen7b\")\n",
-    "    print(\"  OR\")\n",
-    "    print(\"  python -c \\\"from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-Coder-7B', local_dir='./base_model_qwen7b')\\\"\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 5️⃣ Train LoRA Adapter\n",
+    "print(\"=\"*60)\n",
+    "print(\"STARTING TRAINING\")\n",
+    "=\"*60)\n",
+    "print(f\"Working directory: {os.getcwd()}\")\n",
+    "print(f\"Config: {updated_config_path}\")\n",
+    "print(f\"Output: {OUTPUT_DIR}/lora\")\n",
+    "print(\"=\"*60 + \"\\n\")\n",
+    "\n",
+    "# Run training\n",
+    "from train_lora import train_lora\n",
     "\n",
-    "⚠️ **This will take significant time on MPS.**\n",
+    "trainer = train_lora(updated_config_path)\n",
     "\n",
-    "MPS is slower than CUDA, so training will take longer. Consider using a smaller dataset for testing."
+    "print(\"\\n\" + \"=\"*60)\n",
+    "print(\"TRAINING FINISHED OR STOPPED\")\n",
+    "print(\"=\"*60)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": "import os\nimport sys\n\n# Add the training module to path\nsys.path.insert(0, './stack/training')\n\nprint(\"=\"*60)\nprint(\"STARTING TRAINING\")\nprint(\"=\"*60)\nprint(f\"Working directory: {os.getcwd()}\")\nprint(f\"Config: {updated_config_path}\")\nprint(f\"Output: {OUTPUT_DIR}/lora\")\nprint(\"=\"*60 + \"\\n\")\n\n# Run training using the correct function\nfrom train_lora import train_lora\n\ntrainer = train_lora(updated_config_path)\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"TRAINING COMPLETED\")\nprint(\"=\"*60)"
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 6️⃣ Verify Training Output"
+    "## 8️⃣ Verify Training Output"
    ]
   },
   {
@@ -220,14 +289,16 @@
     "        size = os.path.getsize(os.path.join(lora_dir, f)) / (1024*1024)\n",
     "        print(f\"   - {f}: {size:.1f} MB\")\n",
     "else:\n",
-    "    print(\"⚠️ LoRA directory not found\")"
+    "    print(\"⚠️ LoRA directory not found - training may have failed\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 7️⃣ Merge LoRA Adapter (Optional)"
+    "## 9️⃣ Merge LoRA Adapter with Base Model\n",
+    "\n",
+    "Combines the trained adapter with the base model to produce a standalone fine-tuned model."
    ]
   },
   {
@@ -235,13 +306,46 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Merge LoRA with base model\nmerged_dir = f\"{OUTPUT_DIR}/merged\"\n\nprint(\"=\"*60)\nprint(\"MERGING LORA ADAPTER\")\nprint(\"=\"*60)\n\nsys.path.insert(0, './stack/training')\nfrom merge_adapter import merge_adapter\n\n# Create a minimal config for merging\nimport yaml\n\nmerge_config = {\n    'model': {'name': BASE_MODEL_PATH, 'trust_remote_code': True},\n    'output': {'lora_dir': f'{OUTPUT_DIR}/lora', 'merged_dir': merged_dir},\n    'quantization': {'enabled': False}\n}\n\n# Save config for merge\nmerge_config_path = f\"{OUTPUT_DIR}/merge_config.yaml\"\nwith open(merge_config_path, 'w') as f:\n    yaml.dump(merge_config, f)\n\n# Run merge\nmerge_adapter(\n    config_path=merge_config_path,\n    lora_path=f\"{OUTPUT_DIR}/lora\",\n    output_path=merged_dir\n)\n\nprint(f\"\\n✅ Merged model saved to: {merged_dir}\")\nif os.path.exists(merged_dir):\n    print(\"Files:\", os.listdir(merged_dir)[:5])"
+   "source": [
+    "import yaml\n",
+    "import sys\n",
+    "\n",
+    "sys.path.insert(0, './stack/training')\n",
+    "from merge_adapter import merge_adapter\n",
+    "\n",
+    "merged_dir = f\"{OUTPUT_DIR}/merged\"\n",
+    "\n",
+    "print(\"=\"*60)\n",
+    "print(\"MERGING LORA ADAPTER\")\n",
+    "=\"*60)\n",
+    "\n",
+    "# Create merge config\n",
+    "merge_config = {\n",
+    "    'model': {'name': MODEL_DIR, 'trust_remote_code': True},\n",
+    "    'output': {'lora_dir': f'{OUTPUT_DIR}/lora', 'merged_dir': merged_dir},\n",
+    "    'quantization': {'enabled': False}\n",
+    "}\n",
+    "\n",
+    "merge_config_path = f\"{OUTPUT_DIR}/merge_config.yaml\"\n",
+    "with open(merge_config_path, 'w') as f:\n",
+    "    yaml.dump(merge_config, f)\n",
+    "\n",
+    "# Run merge\n",
+    "merge_adapter(\n",
+    "    config_path=merge_config_path,\n",
+    "    lora_path=f\"{OUTPUT_DIR}/lora\",\n",
+    "    output_path=merged_dir\n",
+    ")\n",
+    "\n",
+    "print(f\"\\n✅ Merged model saved to: {merged_dir}\")\n",
+    "!ls -lh {merged_dir}/"
+   ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 8️⃣ Test Inference"
+    "## 🔟 Test Inference (Quick Check)"
    ]
   },
   {
@@ -261,12 +365,12 @@
     "    model = AutoModelForCausalLM.from_pretrained(\n",
     "        model_path,\n",
     "        torch_dtype=torch.float16,\n",
-    "        device_map=\"mps\",\n",
+    "        device_map=\"cuda\",\n",
     "        trust_remote_code=True\n",
     "    )\n",
     "    \n",
     "    prompt = \"Write a Python function to reverse a string:\\n\\n```python\\n\"\n",
-    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"mps\")\n",
+    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
     "    \n",
     "    print(\"Generating...\")\n",
     "    with torch.no_grad():\n",
@@ -294,21 +398,22 @@
    "source": [
     "## 🔚 Training Complete!\n",
     "\n",
-    "Your model is ready in:\n",
-    "- LoRA adapter: `{OUTPUT_DIR}/lora/`\n",
-    "- Merged model: `{OUTPUT_DIR}/merged/`\n",
+    "Your model is ready in `./training_output/merged/` and saved to Google Drive.\n",
     "\n",
     "**Next steps:**\n",
-    "1. Run evaluation on HumanEval/MBPP benchmarks\n",
-    "2. Upload to Hugging Face Hub\n",
-    "3. Apply to Together AI\n"
+    "1. **Download** `training_output/merged/` from Drive to your local machine\n",
+    "2. **Run evaluation**: Evaluate on HumanEval/MBPP benchmarks\n",
+    "3. **Upload model** to Hugging Face Hub\n",
+    "4. **Apply to Together AI**\n",
+    "\n",
+    "**Note:** This model was trained on the full/mini dataset. For better performance, consider training on more data or more epochs.\n"
    ]
   }
  ],
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "name": "Stack 2.9 Local Training (Fixed)",
+   "name": "Stack 2.9 Colab Training (T4 GPU)",
    "provenance": []
   },
   "kernelspec": {