fix(notebooks): end-to-end Colab usability sweep — torchao, config_template, safety fallback

notebooks/dpo_alignment.ipynb

@@ -18,7 +18,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n!forgelm --version"
   },
   {
    "cell_type": "code",

notebooks/galore_memory_optimization.ipynb

@@ -6,12 +6,30 @@
    "source": [
     "# ForgeLM — GaLore Memory Optimization\n",
     "\n",
-    "Train with **GaLore** (Gradient Low-Rank Projection) — an alternative to LoRA that enables full-parameter training with reduced optimizer memory.\n",
-    "\n",
-    "**Key difference from LoRA:**\n",
-    "- LoRA adds adapter matrices and only trains those (~0.2% of params)\n",
-    "- GaLore updates ALL weights but projects gradients to low-rank space (saves optimizer memory)\n",
-    "- No adapter files at inference — the full model is updated\n",
+    "Train with **GaLore** (Gradient Low-Rank Projection) — an optimizer-level\n",
+    "technique that projects gradients into a low-rank subspace so the\n",
+    "optimizer state (Adam's moment buffers) shrinks by the same factor.\n",
+    "Combines cleanly with LoRA: ForgeLM applies GaLore to the LoRA adapter's\n",
+    "trainable parameters, giving you LoRA's parameter-efficiency *and*\n",
+    "GaLore's optimizer-state savings in one pass.\n",
+    "\n",
+    "**Implementation note (different from the academic GaLore paper):**\n",
+    "\n",
+    "- The academic paper applies GaLore to **all** model weights for\n",
+    "  full-parameter training and produces no adapter file.\n",
+    "- ForgeLM requires a `lora:` block in every config (Phase 19 library-API\n",
+    "  invariant), so this notebook trains a LoRA adapter and uses GaLore as\n",
+    "  its **optimizer** — gradient projection is applied to the adapter\n",
+    "  parameters, not the frozen base weights. Output is a PEFT adapter\n",
+    "  loadable via `PeftModel.from_pretrained` (see Step 7).\n",
+    "- For pure full-parameter GaLore, use `optim: \"galore_adamw_8bit\"`\n",
+    "  directly on a non-LoRA HF trainer — that path is outside ForgeLM's\n",
+    "  scope.\n",
+    "\n",
+    "**What you save with `galore_adamw_8bit` (this notebook):** 8-bit\n",
+    "optimizer states cut Adam's first+second moment buffers from FP32 (8\n",
+    "bytes/param) to int8 (1 byte/param) — ~75% reduction on the LoRA\n",
+    "adapter's optimizer state. Useful on tight-VRAM Colab T4.\n",
     "\n",
     "**Requirements:** GPU runtime required (Runtime → Change runtime type → T4 GPU)\n",
     "\n",
@@ -23,7 +41,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM + GaLore\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip install -q galore-torch  # GaLore optimizer\n!pip uninstall -y wandb -q\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM + GaLore\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip install -q galore-torch  # GaLore optimizer\n!pip uninstall -y wandb -q\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n!forgelm --version"
   },
   {
    "cell_type": "code",
@@ -173,16 +191,23 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## GaLore vs LoRA Comparison\n",
-    "\n",
-    "| Aspect | LoRA | GaLore |\n",
-    "|--------|------|--------|\n",
-    "| What it modifies | Adds adapter matrices | Projects gradients to low-rank space |\n",
-    "| Parameters trained | ~0.2% (adapters only) | 100% (all weights) |\n",
-    "| Inference overhead | Adapter loading/merging | None (weights updated in-place) |\n",
-    "| Memory savings | Fewer trainable params | Smaller optimizer states |\n",
-    "| Training speed | Faster (fewer params) | Slower (periodic SVD) |\n",
-    "| Multi-task serving | Easy adapter swapping | Need full model copy |\n",
+    "## GaLore vs LoRA Comparison (as configured in ForgeLM)\n",
+    "\n",
+    "Both are enabled simultaneously in this notebook's config — LoRA defines\n",
+    "*which parameters train*, GaLore defines *how their gradients update*.\n",
+    "\n",
+    "| Aspect | LoRA alone | LoRA + GaLore (this notebook) |\n",
+    "|--------|------------|-------------------------------|\n",
+    "| What trains | LoRA adapter matrices (~0.2% of params) | Same — LoRA adapter matrices |\n",
+    "| Optimizer state | FP32 Adam (8 bytes/param) | 8-bit Adam projected to low-rank (~1 byte/param × low-rank) |\n",
+    "| Output file | PEFT adapter | PEFT adapter (identical format — load via `PeftModel.from_pretrained`) |\n",
+    "| Inference overhead | Adapter loading/merging | Same as LoRA |\n",
+    "| When to add GaLore | When optimizer-state VRAM dominates your peak | Long sequences, big batch, or large LoRA `r` on tight-VRAM cards |\n",
+    "\n",
+    "**For pure full-parameter GaLore** (no LoRA adapter, what the paper\n",
+    "describes): see the upstream `galore-torch` README. ForgeLM's API is\n",
+    "optimised for parameter-efficient fine-tuning so it always wraps with\n",
+    "LoRA.\n",
     "\n",
     "## GaLore Optimizer Variants\n",
     "\n",

notebooks/grpo_reasoning.ipynb

@@ -18,7 +18,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n!forgelm --version"
   },
   {
    "cell_type": "code",
@@ -43,25 +43,29 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Step 3: Create GRPO dataset (prompt-only — model generates responses during training)\n",
+    "# Step 3: Create GRPO dataset (prompt + gold_answer; model generates responses during training)\n",
+    "# The `gold_answer` field triggers ForgeLM's built-in math reward\n",
+    "# (forgelm/trainer.py::_math_reward_fn): each generated completion is parsed\n",
+    "# for \"Answer: <value>\" and compared to gold_answer. Correct → +1.0, wrong → 0.0.\n",
+    "# Without `gold_answer`, GRPO would fall back to format/length-only reward.\n",
     "import json\n",
     "\n",
     "prompts = [\n",
-    "    {\"prompt\": \"What is 15% of 240? Think step by step.\"},\n",
-    "    {\"prompt\": \"If a train travels 60 km/h for 2.5 hours, how far does it go?\"},\n",
-    "    {\"prompt\": \"What is the sum of the first 10 prime numbers?\"},\n",
-    "    {\"prompt\": \"A rectangle has width 5 and perimeter 22. What is the length?\"},\n",
-    "    {\"prompt\": \"If 3x + 7 = 22, what is x?\"},\n",
-    "    {\"prompt\": \"What is 2^10?\"},\n",
-    "    {\"prompt\": \"A store offers 20% off. Original price is $80. What's the final price?\"},\n",
-    "    {\"prompt\": \"How many seconds are in 3 hours?\"},\n",
+    "    {\"prompt\": \"What is 15% of 240? Think step by step. End with: Answer: <number>\", \"gold_answer\": \"36\"},\n",
+    "    {\"prompt\": \"If a train travels 60 km/h for 2.5 hours, how far does it go? End with: Answer: <number>\", \"gold_answer\": \"150\"},\n",
+    "    {\"prompt\": \"What is the sum of the first 10 prime numbers? End with: Answer: <number>\", \"gold_answer\": \"129\"},\n",
+    "    {\"prompt\": \"A rectangle has width 5 and perimeter 22. What is the length? End with: Answer: <number>\", \"gold_answer\": \"6\"},\n",
+    "    {\"prompt\": \"If 3x + 7 = 22, what is x? End with: Answer: <number>\", \"gold_answer\": \"5\"},\n",
+    "    {\"prompt\": \"What is 2^10? End with: Answer: <number>\", \"gold_answer\": \"1024\"},\n",
+    "    {\"prompt\": \"A store offers 20% off. Original price is $80. What\\'s the final price? End with: Answer: <number>\", \"gold_answer\": \"64\"},\n",
+    "    {\"prompt\": \"How many seconds are in 3 hours? End with: Answer: <number>\", \"gold_answer\": \"10800\"},\n",
     "]\n",
     "\n",
     "with open(\"math_prompts.jsonl\", \"w\") as f:\n",
     "    for p in prompts:\n",
     "        f.write(json.dumps(p) + \"\\n\")\n",
     "\n",
-    "print(f\"Created {len(prompts)} GRPO math prompts\")"
+    "print(f\"Created {len(prompts)} GRPO math prompts with gold_answer (correctness reward enabled)\")"
    ]
   },
   {
@@ -198,12 +202,25 @@
    "source": [
     "## How GRPO Works\n",
     "\n",
-    "1. For each prompt, the model generates `num_generations` responses\n",
-    "2. Responses are scored by a reward function (correctness for math)\n",
+    "1. For each prompt, the model generates `num_generations` (=4 here) responses\n",
+    "2. Each response is scored by reward functions registered for the trainer\n",
     "3. The model is updated to prefer higher-scoring responses\n",
     "4. No human preference data needed — the model learns from its own outputs\n",
     "\n",
-    "This is the method behind DeepSeek-R1's reasoning capabilities."
+    "**This demo's reward stack** (defined in `forgelm/trainer.py`):\n",
+    "\n",
+    "1. `combined_format_length_reward` — always on: rewards outputs that\n",
+    "   follow the \"Answer: ...\" format and are within a reasonable length.\n",
+    "2. `_math_reward_fn` — wired in because the dataset has `gold_answer`:\n",
+    "   extracts the numeric answer from each completion and compares it to\n",
+    "   `gold_answer`. Correct → +1.0, wrong → 0.0.\n",
+    "\n",
+    "The two rewards sum additively (TRL convention). Drop `gold_answer`\n",
+    "from the dataset and only the format reward applies (useful for\n",
+    "non-math tasks where you only care about output structure).\n",
+    "\n",
+    "This is the method behind DeepSeek-R1's reasoning capabilities — they\n",
+    "scaled it to chain-of-thought correctness on competition math."
    ]
   }
  ],

notebooks/kto_binary_feedback.ipynb

@@ -18,7 +18,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n!forgelm --version"
   },
   {
    "cell_type": "code",

notebooks/multi_dataset.ipynb

@@ -18,7 +18,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM (with bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7' bitsandbytes\n!pip uninstall -y wandb -q\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n!forgelm --version"
   },
   {
    "cell_type": "code",

notebooks/post_training_workflow.ipynb

@@ -5,7 +5,7 @@
    "id": "7fb27b941602401d91542211134fc71a",
    "metadata": {},
    "source": [
-    "# ForgeLM — Post-Training Workflow (Phase 10)\n",
+    "# ForgeLM — Post-Training Workflow\n",
     "\n",
     "Once you've fine-tuned a model, ForgeLM's Phase 10 toolchain turns it\n",
     "into something deployable. This notebook walks through the four CLI\n",
@@ -37,6 +37,10 @@
    "source": [
     "# Install ForgeLM (skip on Colab if quickstart_sft already ran in this session)\n",
     "!pip install -q --no-cache-dir 'forgelm==0.5.7' bitsandbytes\n",
+    "# Download config_template.yaml from the v0.5.7 tag (NOT main) so the YAML\n",
+    "# matches the wheel above — `main` drifts as features land and would otherwise\n",
+    "# reference fields the installed Pydantic schema rejects.\n",
+    "!wget -q https://raw.githubusercontent.com/cemililik/ForgeLM/v0.5.7/config_template.yaml\n",
     "!forgelm --version"
    ]
   },
@@ -136,9 +140,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "from forgelm.inference import generate, generate_stream, load_model\n",
     "\n",
-    "MODEL_PATH = \"./checkpoints/final_model\"  # or any HF Hub ID\n",
+    "# If you ran quickstart_sft first, point at the trained adapter:\n",
+    "LOCAL_MODEL_PATH = \"./checkpoints/final_model\"\n",
+    "FALLBACK_HF_ID = \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n",
+    "\n",
+    "if os.path.isdir(LOCAL_MODEL_PATH):\n",
+    "    MODEL_PATH = LOCAL_MODEL_PATH\n",
+    "    print(f\"Using fine-tuned model: {MODEL_PATH}\")\n",
+    "else:\n",
+    "    MODEL_PATH = FALLBACK_HF_ID\n",
+    "    print(f\"No local checkpoint at {LOCAL_MODEL_PATH}.\")\n",
+    "    print(f\"Falling back to base model {FALLBACK_HF_ID} so the demo runs end-to-end.\")\n",
+    "    print(\"(Run quickstart_sft.ipynb first to use your own fine-tuned adapter here.)\")\n",
     "\n",
     "# `load_model` returns a (model, tokenizer) pair. Adapter merge happens\n",
     "# automatically when `adapter=...` is given; pass `backend=\"unsloth\"` for\n",
@@ -310,4 +326,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

notebooks/quickstart_sft.ipynb

@@ -24,7 +24,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Step 1: Install ForgeLM (the [qlora] extra pulls in bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7'\n!pip uninstall -y wandb -q  # Remove conflicting wandb (not needed)\n\n# Verify installation\n!forgelm --version"
+   "source": "# Step 1: Install ForgeLM (the [qlora] extra pulls in bitsandbytes for 4-bit quantization)\n# Pinned to v0.5.7; bump on each release\n!pip install -q --no-cache-dir 'forgelm[qlora]==0.5.7'\n!pip uninstall -y wandb -q  # Remove conflicting wandb (not needed)\n\n# Colab base image ships torchao 0.10; modern peft requires >=0.16 for PeftModel.from_pretrained.\n# Upgrade so the comparison cell at the end can load the trained LoRA adapter.\n!pip install -q --upgrade 'torchao>=0.16.0'\n\n# Verify installation\n!forgelm --version"
   },
   {
    "cell_type": "code",