{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FSI_Edge Training on Google Colab GPU\n",
    "## From-scratch novel architecture coding model\n",
    "\n",
    "This notebook trains the FSI_Edge model on Colab's free Tesla T4 GPU.\n",
    "The 27M model runs at ~100+ steps/second on GPU (vs ~0.3 step/s on CPU).\n",
    "\n",
    "**[!] Runtime → Change runtime type → T4 GPU**"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 1: Mount Google Drive & install deps\n",
    "import os, sys, subprocess, torch\n",
    "\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "print('Mounting Google Drive...')\n",
    "FSI_ROOT = '/content/drive/MyDrive/FSI_Edge'\n",
    "os.makedirs(FSI_ROOT, exist_ok=True)\n",
    "os.makedirs(f'{FSI_ROOT}/output', exist_ok=True)\n",
    "os.makedirs(f'{FSI_ROOT}/data', exist_ok=True)\n",
    "print(f'FSI_Edge root: {FSI_ROOT}')\n",
    "\n",
    "print('\\nInstalling dependencies...')\n",
    "subprocess.run(['pip', 'install', '-q', 'tqdm', 'wandb', 'sentencepiece', 'tokenizers', 'transformers', 'pyyaml'], check=True)\n",
    "print('PyTorch comes pre-installed on Colab with CUDA!')\n",
    "print(f'CUDA available: {torch.cuda.is_available()}')\n",
    "print(f'GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"none\"}')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 2: Clone FSI_Edge repo\n",
    "%cd /content\n",
    "\n",
    "import os\n",
    "if not os.path.exists('/content/FSI_Edge'):\n",
    "    !git clone https://github.com/YOUR_USERNAME/FSI_Edge.git\n",
    "else:\n",
    "    %cd FSI_Edge\n",
    "    !git pull\n",
    "\n",
    "# OR upload the source manually:\n",
    "# !cp -r /content/drive/MyDrive/FSI_Edge_source/* /content/FSI_Edge/\n",
    "\n",
    "%cd /content/FSI_Edge\n",
    "sys.path.insert(0, '/content/FSI_Edge')\n",
    "print('Contents:', os.listdir('.'))"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 3: Generate training data\n",
    "print('Generating synthetic training data...')\n",
    "!python data/prepare_data.py --samples 100000 --output /content/FSI_Edge/data/train\n",
    "\n",
    "print('\\nGenerating cold-start reasoning data...')\n",
    "!python data/generate_coldstart.py --num 5000 --output /content/FSI_Edge/data/cold_start.jsonl\n",
    "\n",
    "print('\\nTraining BPE tokenizer...')\n",
    "import sys\n",
    "sys.path.insert(0, '/content/FSI_Edge')\n",
    "from src.data import train_tokenizer\n",
    "train_tokenizer(\n",
    "    corpus_files=['/content/FSI_Edge/data/train/shard_0000.jsonl'],\n",
    "    vocab_size=32768,\n",
    "    output_path='/content/FSI_Edge/fsi_edge_tokenizer'\n",
    ")\n",
    "print('Tokenizer ready!')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 4: Quick test — verify model works on GPU\n",
    "import sys, os, torch\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
    "sys.path.insert(0, '/content/FSI_Edge')\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Device: {device}')\n",
    "\n",
    "from src.model import FSIEdgeModel, FSIEdgeConfig\n",
    "from src.data import CodeDataset, collate_fn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import AdamW\n",
    "\n",
    "# 27M model config\n",
    "config = FSIEdgeConfig(\n",
    "    d_model=256, n_layers=4, n_heads=8, kv_heads=2,\n",
    "    d_ff=1024, max_seq_len=2048, window_size=64,\n",
    "    local_heads=4, struct_heads=2, global_heads=2,\n",
    ")\n",
    "model = FSIEdgeModel(config).to(device)\n",
    "params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f'Model: {params/1e6:.2f}M params on {device}')\n",
    "\n",
    "# Forward+backward test\n",
    "ds = CodeDataset('/content/FSI_Edge/data/train', '/content/FSI_Edge/fsi_edge_tokenizer', max_length=512)\n",
    "loader = DataLoader(ds, batch_size=2, shuffle=True, collate_fn=collate_fn, num_workers=0)\n",
    "batch = next(iter(loader))\n",
    "batch = {k: v.to(device) for k, v in batch.items()}\n",
    "\n",
    "opt = AdamW(model.parameters(), lr=3e-4)\n",
    "out = model(**batch)\n",
    "out.loss.backward()\n",
    "opt.step()\n",
    "opt.zero_grad()\n",
    "print(f'Forward+backward: loss={out.loss.item():.4f} — GPU training works!')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 5: Benchmark GPU vs CPU speed\n",
    "import time\n",
    "\n",
    "# Reset model for clean timing\n",
    "model = FSIEdgeModel(config).to(device)\n",
    "opt = AdamW(model.parameters(), lr=3e-4)\n",
    "\n",
    "# Warmup\n",
    "for _ in range(5):\n",
    "    out = model(**batch)\n",
    "    out.loss.backward()\n",
    "    opt.step()\n",
    "    opt.zero_grad()\n",
    "\n",
    "torch.cuda.synchronize()\n",
    "t0 = time.time()\n",
    "for _ in range(50):\n",
    "    out = model(**batch)\n",
    "    out.loss.backward()\n",
    "    opt.step()\n",
    "    opt.zero_grad()\n",
    "torch.cuda.synchronize()\n",
    "t1 = time.time()\n",
    "\n",
    "step_time = (t1-t0)/50\n",
    "print(f'GPU (T4): {step_time*1000:.1f}ms/step ({(1/step_time)*3600:.0f} steps/hr)')\n",
    "print(f'vs CPU (ARM): ~3000ms/step')\n",
    "print(f'Speedup: ~{3000/step_time:.0f}x')"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 6: Full training run — Stage 1 pretraining\n",
    "import time, json, os\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "OUTPUT = '/content/drive/MyDrive/FSI_Edge/output'\n",
    "TOTAL_STEPS = 50000  # Adjust as needed\n",
    "BATCH_SIZE = 8       # Uses GPU memory\n",
    "GRAD_ACCUM = 4       # Effective batch = 32\n",
    "LOG_EVERY = 50\n",
    "SAVE_EVERY = 5000\n",
    "LR = 3e-4\n",
    "WARMUP = 2000\n",
    "\n",
    "ds = CodeDataset('/content/FSI_Edge/data/train', '/content/FSI_Edge/fsi_edge_tokenizer', max_length=1024)\n",
    "loader = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn, num_workers=2, pin_memory=True)\n",
    "\n",
    "model = FSIEdgeModel(config).to(device)\n",
    "opt = AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95), weight_decay=0.1)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=TOTAL_STEPS)\n",
    "\n",
    "# Resume from checkpoint if exists\n",
    "start_step = 0\n",
    "best_loss = float('inf')\n",
    "latest_ckpt = f'{OUTPUT}/stage1_latest.pt'\n",
    "if os.path.exists(latest_ckpt):\n",
    "    state = torch.load(latest_ckpt, map_location=device)\n",
    "    model.load_state_dict(state['model'])\n",
    "    opt.load_state_dict(state['optimizer'])\n",
    "    start_step = state['step']\n",
    "    best_loss = state.get('best_loss', float('inf'))\n",
    "    print(f'Resumed from step {start_step} (best loss: {best_loss:.4f})')\n",
    "\n",
    "print(f'Starting training: {TOTAL_STEPS} steps, batch={BATCH_SIZE}, accum={GRAD_ACCUM}')\n",
    "print(f'Effective batch: {BATCH_SIZE * GRAD_ACCUM}')\n",
    "\n",
    "model.train()\n",
    "t_start = time.time()\n",
    "losses = []\n",
    "pbar = tqdm(total=TOTAL_STEPS, initial=start_step, desc='Train')\n",
    "\n",
    "while start_step < TOTAL_STEPS:\n",
    "    for batch in loader:\n",
    "        if start_step >= TOTAL_STEPS:\n",
    "            break\n",
    "        \n",
    "        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}\n",
    "        out = model(**batch)\n",
    "        loss = out.loss / GRAD_ACCUM\n",
    "        loss.backward()\n",
    "        \n",
    "        if (start_step + 1) % GRAD_ACCUM == 0:\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            opt.step()\n",
    "            scheduler.step()\n",
    "            opt.zero_grad()\n",
    "        \n",
    "        loss_val = out.loss.item()\n",
    "        losses.append(loss_val)\n",
    "        best_loss = min(best_loss, loss_val)\n",
    "        start_step += 1\n",
    "        pbar.update(1)\n",
    "        pbar.set_postfix({'loss': f'{loss_val:.4f}', 'best': f'{best_loss:.4f}'})\n",
    "        \n",
    "        if start_step % LOG_EVERY == 0:\n",
    "            avg = sum(losses[-LOG_EVERY:]) / min(LOG_EVERY, len(losses))\n",
    "            elapsed = time.time() - t_start\n",
    "            rate = start_step / elapsed\n",
    "            eta = (TOTAL_STEPS - start_step) / rate if rate > 0 else 0\n",
    "            print(f'step {start_step:6d} | loss {loss_val:.4f} | avg {avg:.4f} | '\n",
    "                  f'{rate:.1f} step/s | ETA {eta/60:.0f}min')\n",
    "        \n",
    "        if start_step % SAVE_EVERY == 0:\n",
    "            ckpt = {\n",
    "                'step': start_step, 'model': model.state_dict(),\n",
    "                'optimizer': opt.state_dict(), 'scheduler': scheduler.state_dict(),\n",
    "                'loss': loss_val, 'best_loss': best_loss,\n",
    "            }\n",
    "            torch.save(ckpt, f'{OUTPUT}/stage1_{start_step:06d}.pt')\n",
    "            torch.save(model.state_dict(), latest_ckpt)\n",
    "            print(f'Checkpoint saved: step {start_step}')\n",
    "\n",
    "pbar.close()\n",
    "total_time = time.time() - t_start\n",
    "print(f'\\nStage 1 complete!')\n",
    "print(f'  Steps: {TOTAL_STEPS}')\n",
    "print(f'  Time:  {total_time:.0f}s ({total_time/3600:.2f}h)')\n",
    "print(f'  Loss:  {losses[0]:.4f} -> {losses[-1]:.4f} (best: {best_loss:.4f})')\n",
    "torch.save(model.state_dict(), f'{OUTPUT}/stage1_final.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Continue Training\n",
    "\n",
    "After Stage 1 completes, continue with the next stages.\n",
    "Each stage will require running a similar cell. The checkpoints are saved to Google Drive.\n",
    "\n",
    "### Stage 2: SFT + Cold-Start Reasoning\n",
    "```python\n",
    "from training.train import train_stage2_sft, train_stage2b_cold_start, ColdStartGenerator\n",
    "\n",
    "# Load checkpoint\n",
    "model.load_state_dict(torch.load(f'{OUTPUT}/stage1_final.pt'))\n",
    "\n",
    "# Generate cold-start data\n",
    "generator = ColdStartGenerator(num_examples=5000)\n",
    "cold_data = generator.generate()\n",
    "\n",
    "# SFT\n",
    "model = train_stage2_sft(model, dataloader, config)\n",
    "\n",
    "# Cold-start reasoning\n",
    "model = train_stage2b_cold_start(model, cold_data, config)\n",
    "```\n",
    "\n",
    "### Stage 3: MCPO RL\n",
    "```python\n",
    "from training.train import train_stage3_mcpo, execution_reward\n",
    "model = train_stage3_mcpo(model, execution_reward, config)\n",
    "```\n",
    "\n",
    "### Stage 4: DPO\n",
    "```python\n",
    "from training.train import train_stage4_dpo\n",
    "model = train_stage4_dpo(model, ref_model, dataset, config)\n",
    "```\n",
    "\n",
    "### Export to GGUF\n",
    "```python\n",
    "from export.export_gguf import convert_pytorch_to_gguf\n",
    "convert_pytorch_to_gguf('model.pt', '800M', 'fsi_edge_q4.gguf', 'q4_0')\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Cell 7: Download checkpoints to local machine\n",
    "from google.colab import files\n",
    "\n",
    "OUTPUT = '/content/drive/MyDrive/FSI_Edge/output'\n",
    "# Download to browser\n",
    "files.download(f'{OUTPUT}/stage1_final.pt')\n",
    "\n",
    "# Or zip all checkpoints\n",
    "!zip -r /content/fsi_edge_checkpoints.zip /content/drive/MyDrive/FSI_Edge/output/\n",
    "files.download('/content/fsi_edge_checkpoints.zip')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {"provenance": []},
  "kernelspec": {"name": "python3", "display_name": "Python 3"}
 },
 "nbformat": 4,
 "nbformat_minor": 0
}