{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DeMo fixed-band experiment on a Colab GPU\n",
    "\n",
    "Runs the experiment in [DeMo_Fixed_Band_Experiment](https://github.com/BrendanJamesLynskey/DeMo_Fixed_Band_Experiment) on a GPU, at the `tinystories-gpu` size (about 16M parameters, 3,000 steps per run).\n",
    "\n",
    "1. **Runtime > Change runtime type > GPU** (T4 is free; an A100 is faster).\n",
    "2. **Runtime > Run all.** The first cells check the GPU, run the tests and a smoke run, and time the model; the sweep cell then runs the whole grid.\n",
    "3. **If the session disconnects**, reconnect and *Run all* again. Each session pulls the latest code from GitHub and restores the finished runs from Drive; the timing cell is skipped once the sweep has started.\n",
    "4. **To stop early**, interrupt the sweep cell and run `save_to_drive()` in a new cell before disconnecting, so the last few minutes are kept.\n",
    "\n",
    "The work happens on Colab's local disk (training directly on the Drive mount made the mount drop). Results, checkpoints and the tokenised data are copied to `MyDrive/DeMo_Fixed_Band_Experiment_state` every 10 minutes and when the sweep ends, and restored from there at the start of each session, so finished runs are skipped and an interrupted run resumes from its last checkpoint.\n",
    "\n",
    "Optional: to push results to GitHub as each run finishes, add a Colab secret (key icon, left) named `GH_TOKEN` holding a fine-grained token with write access to this repository's contents."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!nvidia-smi --query-gpu=name,memory.total --format=csv"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "import os, subprocess, threading, time\n",
    "drive.mount('/content/drive')\n",
    "REPO = '/content/DeMo_Fixed_Band_Experiment'\n",
    "STORE = '/content/drive/MyDrive/DeMo_Fixed_Band_Experiment_state'\n",
    "if not os.path.exists(REPO):\n",
    "    !git clone -q https://github.com/BrendanJamesLynskey/DeMo_Fixed_Band_Experiment.git {REPO}\n",
    "%cd {REPO}\n",
    "!git config user.name \"Colab sweep\" && git config user.email \"colab-sweep@users.noreply.github.com\"\n",
    "# fetch and reset rather than pull: results committed to the repository would otherwise clash with\n",
    "# the same files restored from Drive, and the pull would fail, leaving the old code in place\n",
    "!git fetch -q origin && git reset -q --hard origin/main\n",
    "!git log --oneline -1\n",
    "!which rsync > /dev/null || apt-get -qq install -y rsync > /dev/null\n",
    "os.makedirs(STORE, exist_ok=True)\n",
    "# restore what earlier sessions saved: results, checkpoints and the tokenised data\n",
    "!rsync -a {STORE}/results/ results/ 2>/dev/null; rsync -a {STORE}/data/ data/ 2>/dev/null; true\n",
    "import glob\n",
    "print('restored from Drive:', len(glob.glob('results/tinystories-gpu/*.jsonl')), 'run logs,',\n",
    "      len(glob.glob('results/ckpt/*.pt')), 'checkpoints')\n",
    "\n",
    "def save_to_drive():\n",
    "    \"\"\"Copy results, checkpoints and tokenised data to Drive (not the raw text downloads).\"\"\"\n",
    "    if not os.path.isdir(STORE):\n",
    "        print('Drive is not mounted; skipping the copy')\n",
    "        return\n",
    "    # a finished run deletes its checkpoint locally, but rsync does not propagate deletions:\n",
    "    # remove checkpoints of finished runs from both copies, or they fill the Drive quota\n",
    "    import glob, json as _json\n",
    "    for ck in glob.glob('results/ckpt/*.pt') + glob.glob(f'{STORE}/results/ckpt/*.pt'):\n",
    "        stem = os.path.basename(ck)[:-3]\n",
    "        for log in glob.glob(f'results/*/{stem}.jsonl'):\n",
    "            lines = open(log).read().strip().splitlines()\n",
    "            if lines and _json.loads(lines[-1]).get('type') == 'done':\n",
    "                os.remove(ck)\n",
    "                break\n",
    "    for src in ('results', 'data'):\n",
    "        if os.path.isdir(src):\n",
    "            subprocess.run(['rsync', '-a', '--exclude', '*.txt', '--exclude', '*.part', src + '/', f'{STORE}/{src}/'])\n",
    "\n",
    "def _saver(period=600):\n",
    "    while True:\n",
    "        time.sleep(period)\n",
    "        try:\n",
    "            save_to_drive()\n",
    "        except Exception as e:\n",
    "            print('Drive copy failed:', e)\n",
    "\n",
    "if not globals().get('_saver_started'):\n",
    "    threading.Thread(target=_saver, daemon=True).start()\n",
    "    _saver_started = True"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# PyTorch with CUDA is preinstalled on Colab; add the rest\n",
    "!pip -q install tokenizers ruff pytest einops"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Tests (they run on the CPU) and a short run of every variant on synthetic data, on the GPU\n",
    "!python -m pytest -q tests\n",
    "!python experiment/sweep.py --grid ci --dataset synthetic\n",
    "!python -c \"import json,glob; f=sorted(glob.glob('results/synthetic/*.jsonl'))[0]; print('device:', json.loads(open(f).readline())['hardware']['device'])\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Timing: 30 steps of AdamW and of DeMo with a fixed band, then an estimate for the whole grid.\n",
    "# Skipped once the sweep has started (set FORCE_TIMING = True to run it anyway).\n",
    "FORCE_TIMING = False\n",
    "if FORCE_TIMING or not os.path.isdir('results/tinystories-gpu'):\n",
    "    import json, sys, time, dataclasses\n",
    "    sys.path.insert(0, 'experiment')\n",
    "    import sweep, train\n",
    "    from pathlib import Path\n",
    "    prof = sweep.SIZES['tinystories-gpu']\n",
    "    rate = {}\n",
    "    for v in ('adamw', 'band', 'optics'):\n",
    "        cfg = train.RunConfig(**{**prof, 'variant': v, 'steps': 30, 'eval_every': 30, 'diag_every': 30, 'ckpt_every': 1000,\n",
    "                                 'chunk_mode': '1d' if v == 'optics' else '2d', 'chunk': 256 if v == 'optics' else 64})\n",
    "        out = Path('/content/timing') / f'{v}.jsonl'\n",
    "        if out.exists(): out.unlink()\n",
    "        train.run(cfg, out)\n",
    "        ev = [json.loads(l) for l in out.read_text().splitlines() if '\"eval\"' in l][-1]\n",
    "        rate[v] = ev['seconds'] / 30\n",
    "        print(f'{v}: {rate[v]:.2f} s/step')\n",
    "    total = 0\n",
    "    for stage in ('lr', 'core', 'optics', 'shapes'):\n",
    "        for c in sweep.grid(stage, 'tinystories-gpu'):\n",
    "            total += c.steps * rate['adamw' if c.variant == 'adamw' else 'optics' if c.variant == 'optics' else 'band']\n",
    "    print(f'estimated time for the whole grid: {total / 3600:.1f} h')\n",
    "else:\n",
    "    print('sweep already started: timing skipped')\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# The whole experiment: learning-rate sweep, then the core, optics and band-shape stages.\n",
    "# Long-running. Finished runs are skipped on a rerun; an interrupted one resumes.\n",
    "FLAGS = ''\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    token = userdata.get('GH_TOKEN')\n",
    "    !git remote set-url origin \"https://x-access-token:{token}@github.com/BrendanJamesLynskey/DeMo_Fixed_Band_Experiment.git\"\n",
    "    FLAGS = '--commit --push'\n",
    "    print('results will be committed and pushed after each run')\n",
    "except Exception:\n",
    "    print('no GH_TOKEN secret: results stay in Drive')\n",
    "!python experiment/sweep.py --grid all --dataset tinystories-gpu $FLAGS 2>&1 | tee -a results/colab_sweep.log\n",
    "save_to_drive()"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!python experiment/analyse.py --dataset tinystories-gpu\n",
    "from IPython.display import Image, display\n",
    "for f in ('loss_vs_steps_and_bytes.png', 'enob_tolerance.png'):\n",
    "    p = f'results/tinystories-gpu/{f}'\n",
    "    if os.path.exists(p): display(Image(p))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}