Meditation RL Training (GRPO) - Self-contained Colab notebook

meditation_rl_colab.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "L4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meditation RL Training (GRPO)\n",
    "\n",
    "**Self-contained notebook** — no colab_remote, no ngrok. Just set your secrets and run.\n",
    "\n",
    "## Setup\n",
    "1. **GPU**: Change runtime to GPU (L4/A100/H100) via Runtime → Change runtime type\n",
    "2. **Secrets**: Add these in the 🔑 Secrets panel (left sidebar):\n",
    "   - `HF_TOKEN` — HuggingFace token (read/write access)\n",
    "   - `GEMINI_PAID_KEY` — Gemini API key (paid tier for judge)\n",
    "3. **Run all cells** (Ctrl+F9)\n",
    "\n",
    "Training auto-resumes from the latest HF checkpoint. Checkpoints upload to HF every ~15 min.\n",
    "\n",
    "**Model**: LFM2.5-1.2B-Thinking (QLoRA)  \n",
    "**Judge**: Gemini 3.1 Pro Preview (paid API)  \n",
    "**Repo**: Nirav-Madhani/LFM2.5-1.2B-Meditation"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 1: Load secrets and set environment variables\n",
    "from google.colab import userdata\n",
    "import os\n",
    "\n",
    "os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')\n",
    "os.environ['GEMINI_PAID_KEY'] = userdata.get('GEMINI_PAID_KEY')\n",
    "\n",
    "# Optional: if you have a free Gemini key too\n",
    "try:\n",
    "    os.environ['GEMINI_FREE_KEY'] = userdata.get('GEMINI_FREE_KEY')\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "print('Secrets loaded ✓')\n",
    "print(f'HF_TOKEN: {os.environ[\"HF_TOKEN\"][:8]}...')\n",
    "print(f'GEMINI_PAID_KEY: {os.environ[\"GEMINI_PAID_KEY\"][:8]}...')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": "# Cell 2: Check GPU\nimport torch\nif torch.cuda.is_available():\n    name = torch.cuda.get_device_name(0)\n    vram = torch.cuda.get_device_properties(0).total_memory / 1024**3\n    print(f'GPU: {name} ({vram:.1f} GB)')\nelse:\n    raise RuntimeError('No GPU! Change runtime type: Runtime -> Change runtime type -> GPU')",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 3: Download training script from HuggingFace\n",
    "!pip install -q huggingface_hub\n",
    "\n",
    "from huggingface_hub import hf_hub_download\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "WORK_DIR = Path('/content/meditation')\n",
    "WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "HF_REPO = 'Nirav-Madhani/LFM2.5-1.2B-Meditation'\n",
    "HF_TOKEN = os.environ['HF_TOKEN']\n",
    "\n",
    "# Download training script\n",
    "script_path = hf_hub_download(\n",
    "    repo_id=HF_REPO,\n",
    "    filename='meditation-rl-training.py',\n",
    "    local_dir=WORK_DIR,\n",
    "    token=HF_TOKEN,\n",
    ")\n",
    "print(f'Training script: {script_path}')\n",
    "\n",
    "# Also download topics data\n",
    "try:\n",
    "    topics_path = hf_hub_download(\n",
    "        repo_id=HF_REPO,\n",
    "        filename='data/topics.json',\n",
    "        local_dir=WORK_DIR,\n",
    "        token=HF_TOKEN,\n",
    "    )\n",
    "    print(f'Topics: {topics_path}')\n",
    "except Exception as e:\n",
    "    print(f'Topics download skipped (will be fetched by script): {e}')\n",
    "\n",
    "print('\\nFiles in work dir:')\n",
    "for f in sorted(WORK_DIR.rglob('*')):\n",
    "    if f.is_file():\n",
    "        print(f'  {f.relative_to(WORK_DIR)} ({f.stat().st_size/1024:.0f} KB)')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 4: Run training\n",
    "# This runs the self-contained training script which:\n",
    "# - Installs all dependencies\n",
    "# - Downloads SFT checkpoint + latest RL checkpoint from HF\n",
    "# - Auto-detects GPU and configures batch size / GRPO K\n",
    "# - Trains with GRPO, uploading checkpoints to HF every ~15 min\n",
    "# - Auto-resumes if a checkpoint exists\n",
    "\n",
    "!cd /content/meditation && python -u meditation-rl-training.py"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": "# Cell 5 (optional): Check GPU memory usage while training\nimport torch\nif torch.cuda.is_available():\n    alloc = torch.cuda.memory_allocated() / 1024**3\n    total = torch.cuda.get_device_properties(0).total_memory / 1024**3\n    print(f'GPU Memory: {alloc:.1f} / {total:.1f} GB ({alloc/total*100:.0f}%)')",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 6 (optional): List checkpoints on HuggingFace\n",
    "from huggingface_hub import HfApi\n",
    "api = HfApi(token=os.environ['HF_TOKEN'])\n",
    "files = api.list_repo_files('Nirav-Madhani/LFM2.5-1.2B-Meditation')\n",
    "ckpts = sorted([f for f in files if 'checkpoint' in f])\n",
    "print(f'Checkpoints on HF ({len(ckpts)}):')\n",
    "for c in ckpts:\n",
    "    print(f'  {c}')"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}