{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 25: Parameter-Efficient Fine-Tuning"
      ],
      "id": "3a5a7265-d5f3-4e10-9a1a-b78aa89f067b"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "0fe28fcd-bada-4b37-aad3-4c7f158a6d24"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "c515a2c6-0be1-42e9-a665-cb9c950b3cfe"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Every module so far taught you to build a model and then **train all of\n",
        "its weights** — pretraining ([m07](../m07_training/lesson.qmd)) and\n",
        "alignment ([m12](../m12_alignment/lesson.qmd)) both step every\n",
        "parameter. That is fine when you own the compute to update a\n",
        "175-billion-parameter model. Almost nobody does. In practice you take a\n",
        "frozen pretrained model and adapt it with a **tiny** number of new\n",
        "parameters — a technique so dominant that “fine-tuning” today usually\n",
        "*means* this. **Parameter-efficient fine-tuning (PEFT)** learns a small\n",
        "add-on that steers a frozen model, and its flagship is **LoRA** —\n",
        "*Low-Rank Adaptation*.\n",
        "\n",
        "**LoRA** freezes the pretrained weight $W_0$ and learns a low-rank\n",
        "update $\\Delta W = BA$ beside it, so instead of retraining a\n",
        "$d \\times k$ matrix you train two thin ones with rank\n",
        "$r \\ll \\min(d, k)$.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- **You can afford it.** LoRA reported **10,000× fewer** trainable\n",
        "  parameters and **3× less** GPU memory to fine-tune GPT-3 175B — the\n",
        "  difference between a data center and a single GPU.\n",
        "- **It costs nothing at inference.** The adapter *merges* back into the\n",
        "  weight ($W = W_0 + \\frac{\\alpha}{r} BA$), so a deployed LoRA model is\n",
        "  an ordinary linear layer — no extra latency, unlike bolt-on adapter\n",
        "  modules.\n",
        "- **It’s swappable.** One frozen base + many small adapters = one model\n",
        "  that wears many hats (a code adapter, a chat adapter, a domain\n",
        "  adapter), each a few MB.\n",
        "- **It bridges to quantization.** Freeze the base in 4-bit and train the\n",
        "  adapter in full precision and you have **QLoRA** — fine-tuning a 65B\n",
        "  model on one GPU (ties [m14](../m14_quantization/lesson.qmd)).\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain why a fine-tuning update lives in a **low-rank subspace**, and\n",
        "  why that makes it cheap\n",
        "- Write the LoRA forward pass $h = W_0 x + \\frac{\\alpha}{r} BA\\,x$ from\n",
        "  the shapes up\n",
        "- Build `LoRALinear` from scratch, wrapping a frozen `nn.Linear`\n",
        "- Prove the two facts that make LoRA trustworthy — the **zero-init\n",
        "  identity** and **merge equivalence** — on a real `GPTModel`\n",
        "- Compute the parameter budget $r(d+k)$ vs $dk$ and read the savings\n",
        "- Fine-tune a model with only the adapter, leaving the base bit-for-bit\n",
        "  frozen\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the\n",
        "  `nn.Linear` layers LoRA adapts\n",
        "- [Module 07: Training](../m07_training/lesson.qmd) — gradients and the\n",
        "  optimizer step\n",
        "- [Module 12: Alignment](../m12_alignment/lesson.qmd) — the fine-tuning\n",
        "  stages (SFT, DPO) that LoRA makes affordable\n",
        "\n",
        "## Intuition: The Update Is Low-Rank\n",
        "\n",
        "Start with the question full fine-tuning ignores: **how much does a\n",
        "weight matrix actually need to change** to specialize a pretrained\n",
        "model? A $4096 \\times 4096$ attention projection has 16.8M numbers.\n",
        "Fine-tuning nudges all of them — but the *net change* $\\Delta W$ it\n",
        "learns turns out to have very low **intrinsic rank**. The directions\n",
        "that matter live in a handful of dimensions, not thousands.\n",
        "\n",
        "LoRA takes that observation and makes it structural. Rather than let\n",
        "$\\Delta W$ be any $d \\times k$ matrix, it *forces* it to be low rank by\n",
        "writing it as a product of two thin factors:\n",
        "\n",
        "            full fine-tune                     LoRA\n",
        "            ──────────────                     ────\n",
        "       ΔW  =  ┌──────────┐            ΔW  =  ┌─┐   ┌──────────┐\n",
        "              │          │                  │ │ × │    A     │   ← r × k  (down)\n",
        "              │   d×k    │                  │B│   └──────────┘\n",
        "              │          │                  │ │   r\n",
        "              │          │                  │d│\n",
        "              └──────────┘                  └─┘   ← d × r  (up)\n",
        "          d·k trainable numbers        r·(d + k) trainable numbers\n",
        "\n",
        "Every signal the adapter adds has to **squeeze through the\n",
        "$r$-dimensional bottleneck** in the middle. `A` projects the $d$-wide\n",
        "input *down* to $r$ dimensions; `B` projects that back *up* to $k$. With\n",
        "$r = 8$ and $d = k = 4096$, that is $8 \\times (4096 + 4096) = 65{,}536$\n",
        "numbers instead of $16{,}777{,}216$ — **256× fewer**, and the same\n",
        "expressive ceiling as any rank-8 update.\n",
        "\n",
        "Step through the data flow. Watch the wide input collapse into the thin\n",
        "bottleneck and re-expand, running *alongside* the frozen $W_0 x$ path —\n",
        "never touching it."
      ],
      "id": "2b13c7c4-57d1-4846-baf6-867f15c5633a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "6ebbad1f-e556-4710-8668-0e9f34c79277"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "9b9bd147-beee-4a14-86a3-09c04ff98359"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "973b3cba-3b50-4903-bef2-e8da3d8b1d95"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> LoRA doesn’t add a *smaller* dense update — it adds a **structurally\n",
        "> low-rank** one. That is what buys the parameter savings *and* what\n",
        "> limits the adapter: it can only write changes of rank $\\le r$. The bet\n",
        "> is that fine-tuning updates were low-rank to begin with.\n",
        "\n",
        "## The Math: $h = W_0 x + \\frac{\\alpha}{r} BA\\,x$\n",
        "\n",
        "For a linear layer with frozen weight $W_0 \\in \\mathbb{R}^{k \\times d}$\n",
        "(PyTorch stores it as `(out, in)`), LoRA parameterizes the update as\n",
        "\n",
        "$$h = W_0 x + \\Delta W x = W_0 x + \\frac{\\alpha}{r} B A\\, x,\n",
        "\\qquad B \\in \\mathbb{R}^{k \\times r},\\; A \\in \\mathbb{R}^{r \\times d},\\; r \\ll \\min(d, k).$$\n",
        "\n",
        "Three details carry all the weight:\n",
        "\n",
        "1.  **Initialization.** $A$ gets a random Gaussian init, $B$ is\n",
        "    initialized to **zero** — so $\\Delta W = BA = 0$ at the start of\n",
        "    training. Fine-tuning begins as an *exact copy of the pretrained\n",
        "    model*. There is no random perturbation to recover from; the adapter\n",
        "    grows from nothing.\n",
        "\n",
        "2.  **The scaling $\\alpha/r$.** The update is scaled by a constant\n",
        "    $\\alpha/r$. Following the paper, $\\alpha$ behaves like a learning\n",
        "    rate and is not tuned separately — set it to the first $r$ you try\n",
        "    (with $\\alpha = r$ the scale is $1$). Dividing by $r$ keeps the\n",
        "    update magnitude roughly stable as you change the rank.\n",
        "\n",
        "3.  **What trains.** Only $A$ and $B$ receive gradients. $W_0$ is\n",
        "    frozen, so its optimizer state (the momentum and variance buffers\n",
        "    that dominate Adam’s memory) never has to exist — this, not the\n",
        "    parameter count alone, is where the 3× memory win comes from.\n",
        "\n",
        "## Code: `LoRALinear` From Scratch\n",
        "\n",
        "The whole method is a few lines. We wrap an existing `nn.Linear`, freeze\n",
        "it, and add the two matrices. This follows the implementation in\n",
        "`lora.py`."
      ],
      "id": "e951ae20-fb76-4e92-8f03-8a24c6accfa9"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "A shape (r × d): (2, 16)\n",
            "B shape (k × r): (8, 2)\n",
            "B is all zeros at init: True"
          ]
        }
      ],
      "source": [
        "import sys\n",
        "sys.path.insert(0, '..')  # make sibling modules (m06_transformer, …) importable\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from lora import LoRALinear, lora_update\n",
        "\n",
        "torch.manual_seed(0)\n",
        "base = nn.Linear(16, 8, bias=False)      # a pretrained layer: d=16 → k=8\n",
        "lora = LoRALinear(base, rank=2, alpha=2.0)\n",
        "\n",
        "print(\"A shape (r × d):\", tuple(lora.A.shape))\n",
        "print(\"B shape (k × r):\", tuple(lora.B.shape))\n",
        "print(\"B is all zeros at init:\", bool(torch.count_nonzero(lora.B) == 0))"
      ],
      "id": "82e70243"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The forward pass runs the input through the bottleneck rather than\n",
        "forming the full $\\Delta W$: `base(x) + scaling * (x @ A.T) @ B.T`.\n",
        "Algebraically identical, but it only ever multiplies by the thin\n",
        "matrices.\n",
        "\n",
        "### Anchor 1 — the zero-init identity\n",
        "\n",
        "Because $B = 0$, the adapter contributes *exactly* nothing at\n",
        "initialization. The LoRA layer is bit-for-bit the frozen layer:"
      ],
      "id": "0a96e4d2-1689-4ddb-be93-e58312be227d"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "LoRA(x) == base(x) exactly: True"
          ]
        }
      ],
      "source": [
        "x = torch.randn(4, 16)\n",
        "print(\"LoRA(x) == base(x) exactly:\", bool(torch.equal(lora(x), base(x))))"
      ],
      "id": "b553ccdf"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This is the property that makes LoRA safe to attach: you never degrade\n",
        "the model by adding an adapter. Let’s confirm it on a **real\n",
        "transformer** — wrap a `GPTModel`’s output projection and check the\n",
        "logits are unchanged:"
      ],
      "id": "742ffa21-5bf3-46bb-8862-b091e7127109"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "logits identical after attaching LoRA: True\n",
            "trainable params: 528\n",
            "frozen params:    30720"
          ]
        }
      ],
      "source": [
        "from m06_transformer.transformer import GPTModel\n",
        "from lora import inject_lora, trainable_parameters, frozen_parameters\n",
        "\n",
        "model = GPTModel(vocab_size=100, embed_dim=32, num_heads=2, num_layers=2,\n",
        "                 max_seq_len=64, dropout=0.0)\n",
        "model.eval()\n",
        "tokens = torch.randint(0, 100, (2, 16))\n",
        "with torch.no_grad():\n",
        "    before = model(tokens)\n",
        "\n",
        "inject_lora(model, [\"lm_head\"], rank=4, alpha=8.0)   # freeze all, adapt lm_head\n",
        "with torch.no_grad():\n",
        "    after = model(tokens)\n",
        "\n",
        "print(\"logits identical after attaching LoRA:\", bool(torch.equal(before, after)))\n",
        "print(\"trainable params:\", trainable_parameters(model))\n",
        "print(\"frozen params:   \", frozen_parameters(model))"
      ],
      "id": "d2b0d4d5"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Only the two adapter matrices train; the entire transformer is frozen.\n",
        "\n",
        "### Anchor 2 — merging away the adapter (no inference latency)\n",
        "\n",
        "Since $h = (W_0 + \\frac{\\alpha}{r} BA)\\,x$, we can fold the update into\n",
        "a single weight once training is done. The merged layer is a plain\n",
        "`nn.Linear` that reproduces the adapter’s output — so LoRA adds **zero**\n",
        "inference cost:"
      ],
      "id": "8657c80c-33d8-4d59-9d2d-92659c715c44"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "max |LoRA(x) − merged(x)|: 1.79e-07\n",
            "merged is a plain nn.Linear: True"
          ]
        }
      ],
      "source": [
        "# give the adapter a real (non-zero) update to merge\n",
        "with torch.no_grad():\n",
        "    lora.B.copy_(torch.randn_like(lora.B))\n",
        "\n",
        "merged = lora.merge()                       # a standalone nn.Linear\n",
        "diff = (lora(x) - merged(x)).detach().abs().max()\n",
        "print(\"max |LoRA(x) − merged(x)|:\", f\"{diff.item():.2e}\")\n",
        "print(\"merged is a plain nn.Linear:\", isinstance(merged, nn.Linear))"
      ],
      "id": "45e6a560"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The difference is floating-point noise. At deployment you merge, ship\n",
        "one matrix, and the adapter has vanished.\n",
        "\n",
        "## The Parameter Budget\n",
        "\n",
        "The whole point is the count. A full fine-tune of one $(k, d)$ layer\n",
        "trains $d\n",
        "\\cdot k$ numbers; LoRA trains $r \\cdot (d + k)$. Drive the two dials —\n",
        "the layer width $d$ and the rank $r$ — and watch the gap open on a log\n",
        "scale. The reduction factor is $\\dfrac{dk}{r(d+k)}$, which for a square\n",
        "layer is $\\dfrac{d}{2r}$."
      ],
      "id": "35c4a230-5981-473f-ac24-48f462374c18"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "from lora import (count_full_parameters, count_lora_parameters,\n",
        "                  parameter_reduction)\n",
        "\n",
        "# A GPT-3-scale attention projection, adapted at rank 8.\n",
        "anchor = {\n",
        "    \"d\": 4096, \"r\": 8,\n",
        "    \"full\": count_full_parameters(4096, 4096),\n",
        "    \"lora\": count_lora_parameters(4096, 4096, 8),\n",
        "    \"reduction\": parameter_reduction(4096, 4096, 8),\n",
        "}\n",
        "ojs_define(loraAnchor = anchor)"
      ],
      "id": "83203bed"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "191b0a0b-5bc7-4ad4-b5be-73b543ca9aff"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "c6ccea5c-ab6f-4202-990d-49a1c33fe235"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a242ccd4-e027-4eb9-92ed-e5d1ed37507d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Push $r$ to 64.** The savings shrink — a fat adapter approaches\n",
        ">     a full fine-tune. LoRA’s win *is* the small rank.\n",
        "> 2.  **Halve $d$ from 4096 to 2048 at $r = 8$.** The reduction halves\n",
        ">     too: for a square layer it is exactly $d/2r$.\n",
        "> 3.  **Set $r = 1$.** The extreme rank-one adapter — one down-vector,\n",
        ">     one up-vector — still adapts the layer, at $2d$ parameters.\n",
        "\n",
        "## Fine-Tuning With Only the Adapter\n",
        "\n",
        "Put it together: freeze a model, attach an adapter, and train **only**\n",
        "$A$ and $B$. The demo in `lora.py` fits an adapter so the frozen layer\n",
        "plus its update matches a target that is genuinely rank-$r$ — exactly\n",
        "what a rank-$r$ adapter can represent. Watch three things at once: the\n",
        "loss falls, the update norm $\\lVert \\Delta W \\rVert$ rises **from\n",
        "exactly zero**, and the frozen base weight never moves."
      ],
      "id": "6073b188-d5a9-492b-ac94-be27411f0459"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "from lora import demonstrate_lora\n",
        "\n",
        "demo = demonstrate_lora()          # frozen base, SGD on A and B only\n",
        "curve = [{\"step\": s.step, \"loss\": s.loss,\n",
        "          \"delta\": s.delta_norm, \"base\": s.base_norm} for s in demo[\"steps\"]]\n",
        "ojs_define(loraCurve = curve)\n",
        "ojs_define(loraMeta = {\"trainable\": demo[\"trainable\"], \"frozen\": demo[\"frozen\"],\n",
        "                       \"base_unchanged\": demo[\"base_unchanged\"]})"
      ],
      "id": "6095399d"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "trainable (A + B): 64\n",
            "frozen (W₀):       256\n",
            "first-step ‖ΔW‖:   0.0044\n",
            "last-step  ‖ΔW‖:   23.20\n",
            "base weight bit-for-bit unchanged: True"
          ]
        }
      ],
      "source": [
        "print(\"trainable (A + B):\", demo[\"trainable\"])\n",
        "print(\"frozen (W₀):      \", demo[\"frozen\"])\n",
        "print(\"first-step ‖ΔW‖:  \", f\"{demo['steps'][0].delta_norm:.4f}\")\n",
        "print(\"last-step  ‖ΔW‖:  \", f\"{demo['steps'][-1].delta_norm:.2f}\")\n",
        "print(\"base weight bit-for-bit unchanged:\", demo[\"base_unchanged\"])"
      ],
      "id": "fd4d1837"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "524f68f5-26e9-45fc-af9e-129213fddf8c"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> The base is not “mostly” frozen — it is **exactly** frozen.\n",
        "> `demonstrate_lora` checks the base weight is bit-for-bit identical\n",
        "> before and after training. All the learning lives in $A$ and $B$; the\n",
        "> pretrained knowledge is untouched and shared.\n",
        "\n",
        "## Beyond LoRA\n",
        "\n",
        "LoRA is the foundation of a whole family. The book builds the core here;\n",
        "these are the load-bearing extensions (roadmap follow-ups):\n",
        "\n",
        "- **QLoRA** (Dettmers et al., 2023) — freeze the base in **4-bit** (NF4\n",
        "  quantization, [m14](../m14_quantization/lesson.qmd)) and keep the LoRA\n",
        "  adapter in full precision. Because the frozen base carries no\n",
        "  optimizer state, quantizing it is nearly free — this is what fits a\n",
        "  65B fine-tune on a single 48 GB GPU.\n",
        "- **DoRA** (Liu et al., 2024) — decompose the weight into **magnitude**\n",
        "  and **direction**, and apply LoRA only to the direction; closes much\n",
        "  of the gap to full fine-tuning at the same parameter budget.\n",
        "- **Adapters** (Houlsby et al., 2019) — the predecessor: small\n",
        "  bottleneck MLPs *inserted between* layers. They work, but unlike LoRA\n",
        "  they **cannot** be merged, so they add permanent inference latency.\n",
        "- **Prefix / prompt tuning** (Li & Liang; Lester et al., 2021) — freeze\n",
        "  the whole model and learn a handful of virtual “prefix” key/value\n",
        "  vectors instead of touching weights at all.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "| Pitfall | Why it bites | Fix |\n",
        "|-----------------------|------------------------------------|--------------|\n",
        "| Initializing **both** $A$ and $B$ randomly | $\\Delta W \\ne 0$ at step 0 — you start by *corrupting* the pretrained model | Keep $B = 0$ (Gaussian $A$) so training begins as the identity |\n",
        "| Treating rank $r$ as “bigger is better” | A large $r$ erases the savings and can overfit the small fine-tune set | Start small ($r = 8$–$16$); raise only if the task underfits |\n",
        "| Forgetting to **freeze** the base | You silently full-fine-tune — the memory win evaporates | `requires_grad_(False)` on the base (what `inject_lora` does) |\n",
        "| Merging, then continuing to train | Once merged, there is no separate adapter to update | Merge only for deployment; keep the adapter while training |\n",
        "| Adapting the wrong layers | The update has to land where the task needs it | The paper adapts attention $W_q, W_v$; that is the usual first choice |\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: The rank bound\n",
        "\n",
        "Show that no matter what $A$ and $B$ contain, the update $\\Delta W$ has\n",
        "rank at most $r$."
      ],
      "id": "e43da5d3-4f50-409b-b19b-dae7886d3d61"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "rank(ΔW): 3 ≤ r = 3"
          ]
        }
      ],
      "source": [
        "from lora import LoRALinear\n",
        "\n",
        "base = nn.Linear(32, 24, bias=False)\n",
        "lora = LoRALinear(base, rank=3, alpha=3.0)\n",
        "with torch.no_grad():\n",
        "    lora.A.copy_(torch.randn_like(lora.A))\n",
        "    lora.B.copy_(torch.randn_like(lora.B))\n",
        "\n",
        "rank = torch.linalg.matrix_rank(lora.delta_weight())\n",
        "print(\"rank(ΔW):\", int(rank), \"≤ r =\", lora.rank)"
      ],
      "id": "e5ae0a47"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Only the adapter gets gradients\n",
        "\n",
        "Run one backward pass and confirm the base weight has **no** gradient\n",
        "while $A$ and $B$ do."
      ],
      "id": "1d17a953-2bc5-4554-9ee3-7237a7fd175e"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "base.weight.grad is None: True\n",
            "A.grad set: True | B.grad set: True"
          ]
        }
      ],
      "source": [
        "base = nn.Linear(8, 8, bias=False)\n",
        "lora = LoRALinear(base, rank=2)\n",
        "with torch.no_grad():\n",
        "    lora.B.copy_(torch.randn_like(lora.B))     # leave the zero-init so B gets a grad\n",
        "lora(torch.randn(3, 8)).pow(2).mean().backward()\n",
        "\n",
        "print(\"base.weight.grad is None:\", base.weight.grad is None)\n",
        "print(\"A.grad set:\", lora.A.grad is not None, \"| B.grad set:\", lora.B.grad is not None)"
      ],
      "id": "5253f57f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: Swap two adapters on one base\n",
        "\n",
        "One frozen model, two adapters. Attach `B₁`, then a different `B₂`, and\n",
        "confirm the merged weights differ — the “many hats, one base” story."
      ],
      "id": "894ca737-3ab1-450f-8b72-91e8c28ad7b9"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Your implementation here:\n",
        "# 1. build a base nn.Linear and two LoRALinear wrappers around copies\n",
        "# 2. set different B matrices\n",
        "# 3. compare merged_weight() — they should differ, base unchanged"
      ],
      "id": "f3529902"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **Fine-tuning updates are low-rank.** LoRA exploits this by writing\n",
        "    $\\Delta W = BA$ with rank $r \\ll \\min(d, k)$ — the same expressive\n",
        "    ceiling as any rank-$r$ update, at a fraction of the parameters.\n",
        "2.  **The forward pass is $h = W_0 x + \\frac{\\alpha}{r} BA\\,x$.** Signal\n",
        "    squeezes through the $r$-dim bottleneck; $A$ down-projects, $B$\n",
        "    up-projects, $\\alpha/r$ scales.\n",
        "3.  **Zero-init identity.** $B = 0$ at start, so $\\Delta W = 0$ and the\n",
        "    adapted layer is *exactly* the pretrained one — verified bit-for-bit\n",
        "    on a real `GPTModel`.\n",
        "4.  **No inference latency.** $W = W_0 + \\frac{\\alpha}{r} BA$ merges the\n",
        "    adapter into one matrix — unlike inserted adapter modules.\n",
        "5.  **The budget: $r(d+k)$ vs $dk$.** For a square layer that is a\n",
        "    $d/2r$ reduction — 256× at $d = 4096$, $r = 8$; 10,000× fewer\n",
        "    trainable params for GPT-3 175B.\n",
        "6.  **The base stays frozen.** Only $A$ and $B$ train, so the base\n",
        "    carries no optimizer state (the real memory win) and is shared\n",
        "    across many swappable adapters.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You now have the technique that makes every other fine-tuning method in\n",
        "this book affordable — plug it into the SFT and DPO objectives of\n",
        "[Module 12: Alignment](../m12_alignment/lesson.qmd), or freeze the base\n",
        "in 4-bit with [Module 14: Quantization](../m14_quantization/lesson.qmd)\n",
        "to reconstruct **QLoRA** from the two halves.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [LoRA: Low-Rank Adaptation of Large Language\n",
        "  Models](https://arxiv.org/abs/2106.09685) — Hu et al., 2021 — the\n",
        "  method built here.\n",
        "- [QLoRA: Efficient Finetuning of Quantized\n",
        "  LLMs](https://arxiv.org/abs/2305.14314) — Dettmers et al., 2023 —\n",
        "  4-bit base + LoRA adapter.\n",
        "- [DoRA: Weight-Decomposed Low-Rank\n",
        "  Adaptation](https://arxiv.org/abs/2402.09353) — Liu et al., 2024 —\n",
        "  magnitude/direction split.\n",
        "- [Parameter-Efficient Transfer Learning for\n",
        "  NLP](https://arxiv.org/abs/1902.00751) — Houlsby et al., 2019 — the\n",
        "  original adapters LoRA improves on.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [Hugging Face PEFT](https://github.com/huggingface/peft) — production\n",
        "  LoRA/QLoRA/DoRA implementations.\n",
        "- [The Intrinsic Dimension of Objective\n",
        "  Landscapes](https://arxiv.org/abs/1804.08838) — Li et al., 2018 — the\n",
        "  low-intrinsic-rank evidence LoRA rests on."
      ],
      "id": "764af4ba-d079-4fd4-ab7b-464f22b5182a"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3"
    },
    "language_info": {
      "name": "python",
      "codemirror_mode": {
        "name": "ipython",
        "version": "3"
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.14"
    }
  }
}