{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM8-L3: LoRA fine-tuning experiment\n",
    "\n",
    "You will implement **LoRA** (low-rank adaptation) for a small GPT, adapt a model that already follows two instructions to a **new** instruction, and compare three ways of doing it:\n",
    "full fine-tuning, LoRA with rank 8 and LoRA with rank 2. You will count trainable parameters, estimate training memory, measure quality on the new task and on the old ones\n",
    "(forgetting), and **merge** the LoRA update back into the weights.\n",
    "\n",
    "**The model and tasks are tiny and synthetic.** The base model (about 113,000 parameters, character-level) was tuned on \"write the word in capitals\" and \"say the word twice\".\n",
    "The new instruction is \"make the word loud\" (capitals followed by an exclamation mark), a small change on top of what the model already does. This is a setting in which a low-rank update is expected to work; the numbers do not predict\n",
    "what LoRA does on a large model and a real task.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, 2 to 3 minutes. No downloads.\n",
    "\n",
    "Cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`). Theory: Lessons LM8-T2 and LM8-T5."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import copy\n",
    "import math\n",
    "import random\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(0)\n",
    "print(\"torch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Given: tasks, data, model and helpers\n",
    "\n",
    "Same machinery as the instruction-tuning notebook (LM7-L3): a template, masked labels, padding, greedy generation until an end token, exact-match accuracy."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TASKS = {\"uppercase\": lambda w: w.upper(), \"repeat\": lambda w: w + \" \" + w, \"loud\": lambda w: w.upper() + \"!\"}\n",
    "INSTRUCTION = {\"uppercase\": \"write the word in capitals\", \"repeat\": \"say the word twice\", \"loud\": \"make the word loud\"}\n",
    "\n",
    "rng = random.Random(0)\n",
    "pool = set()\n",
    "while len(pool) < 460:\n",
    "    pool.add(\"\".join(rng.choice(\"abcdefghijklmnopqrstuvwxyz\") for _ in range(rng.randint(3, 5))))\n",
    "pool = sorted(pool)\n",
    "rng.shuffle(pool)\n",
    "train_words, eval_words = pool[:400], pool[400:430]\n",
    "\n",
    "END, PADC = \"\\x03\", \"\\x00\"\n",
    "chars = sorted(set(\"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ #:\\n.-!\" + \"\".join(INSTRUCTION.values())) | {END, PADC})\n",
    "stoi = {c: i for i, c in enumerate(chars)}\n",
    "itos = {i: c for c, i in stoi.items()}\n",
    "VOCAB, PAD, END_ID = len(chars), stoi[PADC], stoi[END]\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "\n",
    "def prompt(task, word):\n",
    "    return f\"### Instruction:\\n{INSTRUCTION[task]}\\n\\n### Input:\\n{word}\\n\\n### Response:\\n\"\n",
    "\n",
    "def make_pair(p, response):\n",
    "    return torch.tensor(encode(p + response) + [END_ID]), torch.tensor([-100] * len(p) + encode(response) + [END_ID])\n",
    "\n",
    "def collate(pairs):\n",
    "    n = max(len(i) for i, _ in pairs)\n",
    "    ids = torch.full((len(pairs), n), PAD); lab = torch.full((len(pairs), n), -100)\n",
    "    for k, (i, l) in enumerate(pairs):\n",
    "        ids[k, : len(i)] = i; lab[k, : len(l)] = l\n",
    "    return ids[:, :-1], lab[:, 1:]\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d, heads):\n",
    "        super().__init__()\n",
    "        self.h = heads\n",
    "        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)\n",
    "        self.qkv = nn.Linear(d, 3 * d, bias=False)\n",
    "        self.proj = nn.Linear(d, d)\n",
    "        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, n, d = x.shape\n",
    "        q, k, v = self.qkv(self.n1(x)).view(B, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)\n",
    "        x = x + self.proj(F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, n, d))\n",
    "        return x + self.ff(self.n2(x))\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    def __init__(self, vocab, context=96, d=64, heads=4, layers=2):\n",
    "        super().__init__()\n",
    "        self.context = context\n",
    "        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)\n",
    "        self.blocks = nn.ModuleList([Block(d, heads) for _ in range(layers)])\n",
    "        self.norm, self.head = nn.LayerNorm(d), nn.Linear(d, vocab, bias=False)\n",
    "\n",
    "    def forward(self, idx):\n",
    "        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))\n",
    "        for b in self.blocks:\n",
    "            x = b(x)\n",
    "        return self.head(self.norm(x))\n",
    "\n",
    "@torch.no_grad()\n",
    "def generate(model, p, max_new=24):\n",
    "    model.eval(); ids, out = encode(p), []\n",
    "    for _ in range(max_new):\n",
    "        nxt = int(model(torch.tensor([ids[-model.context:]]))[0, -1].argmax())\n",
    "        if nxt == END_ID:\n",
    "            break\n",
    "        ids.append(nxt); out.append(nxt)\n",
    "    model.train()\n",
    "    return \"\".join(itos[i] for i in out)\n",
    "\n",
    "def accuracy(model, task, words=eval_words):\n",
    "    return sum(generate(model, prompt(task, w)) == TASKS[task](w) for w in words) / len(words)\n",
    "\n",
    "def tune(model, items, steps, lr, params=None, seed=2):\n",
    "    data = [make_pair(prompt(t, w), TASKS[t](w)) for t, w in items]\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    params = list(model.parameters()) if params is None else params\n",
    "    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.0)\n",
    "    for _ in range(steps):\n",
    "        ix = torch.randint(0, len(data), (32,), generator=g)\n",
    "        X, Y = collate([data[i] for i in ix])\n",
    "        loss = F.cross_entropy(model(X).flatten(0, 1), Y.flatten(), ignore_index=-100)\n",
    "        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. Given: a base model that follows two instructions\n",
    "\n",
    "A few seconds of pretraining on plain text, then instruction tuning on `uppercase` and `repeat` (about 30 seconds in total). It has never seen the instruction `loud`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "base = GPT(VOCAB)\n",
    "text = \" \".join(random.Random(1).choice(pool) for _ in range(4000))\n",
    "stream = torch.tensor(encode(text))\n",
    "g = torch.Generator().manual_seed(1)\n",
    "opt = torch.optim.AdamW(base.parameters(), lr=3e-3, weight_decay=0.1)\n",
    "for _ in range(150):\n",
    "    s = torch.randint(0, len(stream) - 65, (32,), generator=g)\n",
    "    x = torch.stack([stream[i : i + 64] for i in s]); y = torch.stack([stream[i + 1 : i + 65] for i in s])\n",
    "    loss = F.cross_entropy(base(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()\n",
    "t0 = time.time()\n",
    "tune(base, [(t, w) for t in (\"uppercase\", \"repeat\") for w in train_words], steps=1500, lr=1e-3)\n",
    "N_BASE = sum(p.numel() for p in base.parameters())\n",
    "print(f\"base model: {N_BASE:,} parameters, ready in {time.time() - t0:.0f} s\")\n",
    "acc_base = {t: accuracy(base, t) for t in TASKS}\n",
    "print(\"base accuracy:\", {k: round(v, 2) for k, v in acc_base.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 3. TODO: the LoRA layer\n",
    "\n",
    "`LoRALinear(base, r, alpha)` wraps a frozen `nn.Linear` with a low-rank update. With the base weight `W` of shape `(out, in)`:\n",
    "\n",
    "* `A` is a trainable `nn.Parameter` of shape `(r, in)`, initialised with small random numbers (`torch.randn(r, in) * 0.02`);\n",
    "* `B` is a trainable `nn.Parameter` of shape `(out, r)`, initialised to **zeros**;\n",
    "* `forward(x) = base(x) + (alpha / r) * (x @ A.T) @ B.T`.\n",
    "\n",
    "Because `B` starts at zero, the wrapped layer initially computes exactly what the base layer computes (Lesson LM8-T2)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class LoRALinear(nn.Module):\n",
    "    def __init__(self, base, r, alpha):\n",
    "        super().__init__()\n",
    "        self.base, self.r, self.scale = base, r, alpha / r\n",
    "        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.02)\n",
    "        self.B = nn.Parameter(torch.zeros(base.out_features, r))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.base(x) + self.scale * (x @ self.A.T) @ self.B.T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "lin = nn.Linear(10, 6)\n",
    "lora = LoRALinear(lin, r=2, alpha=4)\n",
    "x = torch.randn(3, 10)\n",
    "assert torch.allclose(lora(x), lin(x)), \"with B = 0 the LoRA layer must equal the base layer\"\n",
    "assert sum(p.numel() for p in [lora.A, lora.B]) == 2 * (10 + 6), \"LoRA adds r * (in + out) parameters\"\n",
    "for p in lin.parameters():\n",
    "    p.requires_grad = False\n",
    "lora(x).sum().backward()\n",
    "assert lin.weight.grad is None, \"the base weight is frozen and gets no gradient\"\n",
    "assert lora.B.grad is not None and lora.B.grad.abs().sum() > 0, \"B gets a gradient at the first step\"\n",
    "assert lora.A.grad.abs().sum() == 0, \"A gets zero gradient at the first step because B is zero (it starts learning after B moves)\"\n",
    "print(\"self-check 1 passed: LoRA layer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 4. TODO: add LoRA to a model\n",
    "\n",
    "`add_lora(model, r, alpha)`: freeze **every** parameter of the model, then in each block replace `qkv`, `proj`, `ff[0]` and `ff[2]` by `LoRALinear` wrappers (the new `A` and `B` parameters are trainable).\n",
    "Return the model. Embeddings, layer norms and the output head stay frozen."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def add_lora(model, r, alpha):\n",
    "    for p in model.parameters():\n",
    "        p.requires_grad = False\n",
    "    for b in model.blocks:\n",
    "        b.qkv = LoRALinear(b.qkv, r, alpha)\n",
    "        b.proj = LoRALinear(b.proj, r, alpha)\n",
    "        b.ff[0] = LoRALinear(b.ff[0], r, alpha)\n",
    "        b.ff[2] = LoRALinear(b.ff[2], r, alpha)\n",
    "    return model\n",
    "\n",
    "def trainable(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "### Self-check 2"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "probe = torch.randint(0, VOCAB, (2, 12))\n",
    "m8 = add_lora(copy.deepcopy(base), r=8, alpha=16)\n",
    "expected = 2 * 8 * ((64 + 192) + (64 + 64) + (64 + 256) + (256 + 64))        # per block: qkv, proj, ff up, ff down\n",
    "assert trainable(m8) == expected, f\"expected {expected} trainable parameters, got {trainable(m8)}\"\n",
    "assert torch.allclose(m8(probe), base(probe), atol=1e-6), \"at initialisation the LoRA model must behave exactly like the base model\"\n",
    "assert all(not p.requires_grad for n, p in m8.named_parameters() if not (n.endswith(\".A\") or n.endswith(\".B\")))\n",
    "print(f\"self-check 2 passed: rank 8 trains {trainable(m8):,} of {sum(p.numel() for p in m8.parameters()):,} parameters ({100 * trainable(m8) / N_BASE:.1f}% of the base model)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 5. TODO: merge the update into the weights\n",
    "\n",
    "`merge_lora(model)`: for every `LoRALinear`, add `scale * B @ A` to the base weight (`W <- W + scale * B @ A`, shape `(out, in)`), and put the plain `nn.Linear` back where the wrapper was.\n",
    "After merging the model has the original architecture and parameter count, and computes the same function: no extra cost at inference (Lesson LM8-T2)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def merge_lora(model):\n",
    "    def fold(layer):\n",
    "        if isinstance(layer, LoRALinear):\n",
    "            with torch.no_grad():\n",
    "                layer.base.weight += layer.scale * (layer.B @ layer.A)\n",
    "            return layer.base\n",
    "        return layer\n",
    "    for b in model.blocks:\n",
    "        b.qkv, b.proj = fold(b.qkv), fold(b.proj)\n",
    "        b.ff[0], b.ff[2] = fold(b.ff[0]), fold(b.ff[2])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "### Self-check 3"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(3)\n",
    "test = add_lora(copy.deepcopy(base), r=4, alpha=8)\n",
    "for n, p in test.named_parameters():                    # make the update non-zero, as after training\n",
    "    if n.endswith(\".B\"):\n",
    "        p.data = torch.randn_like(p) * 0.05\n",
    "before = test(probe)\n",
    "merged = merge_lora(test)\n",
    "assert not any(isinstance(m, LoRALinear) for m in merged.modules()), \"no wrapper should remain after merging\"\n",
    "assert sum(p.numel() for p in merged.parameters()) == N_BASE, \"the merged model has the original parameter count\"\n",
    "assert torch.allclose(merged(probe), before, atol=1e-5), \"merging must not change the function\"\n",
    "print(\"self-check 3 passed: merge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 6. TODO: a memory estimate\n",
    "\n",
    "`train_memory_mb(n_params, n_trainable)`: memory in megabytes (1 MB = 10^6 bytes) for float32 weights, plus gradients and Adam state for the trainable parameters only:\n",
    "`4 * n_params + (4 + 8) * n_trainable` bytes (Lesson LM8-T3). Activations are not included."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def train_memory_mb(n_params, n_trainable):\n",
    "    return (4 * n_params + 12 * n_trainable) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "### Self-check 4"
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "assert abs(train_memory_mb(124_412_160, 124_412_160) - 1990.59) < 0.01, \"GPT-2 small, full fine-tuning: about 1,991 MB\"\n",
    "assert abs(train_memory_mb(124_412_160, 1_000_000) - 509.65) < 0.01\n",
    "print(\"self-check 4 passed: memory estimate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "## 7. Given: the experiment\n",
    "\n",
    "Adapt the base model to the new instruction `loud` with 200 examples, mixed with 80 replay examples of the two old tasks so that forgetting is limited for everybody.\n",
    "Variants: full fine-tuning (600 steps), LoRA rank 8 and rank 2 (900 steps, higher learning rate: LoRA usually needs one). About one minute."
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "new_items = [(\"loud\", w) for w in train_words[:200]] + [(t, w) for t in (\"uppercase\", \"repeat\") for w in train_words[200:240]]\n",
    "rows = {}\n",
    "for name, r, steps, lr in [(\"full fine-tuning\", None, 700, 1e-3), (\"LoRA r=8\", 8, 1100, 3e-3), (\"LoRA r=2\", 2, 1100, 3e-3)]:\n",
    "    model = copy.deepcopy(base)\n",
    "    if r:\n",
    "        add_lora(model, r, 2 * r)\n",
    "    params = [p for p in model.parameters() if p.requires_grad]\n",
    "    n_tr = sum(p.numel() for p in params)\n",
    "    t0 = time.time()\n",
    "    tune(model, new_items, steps, lr, params=params)\n",
    "    acc = {t: accuracy(model, t) for t in TASKS}\n",
    "    n_total = sum(p.numel() for p in model.parameters())\n",
    "    rows[name] = dict(trainable=n_tr, acc=acc, mem=train_memory_mb(n_total, n_tr), time=time.time() - t0, model=model)\n",
    "print(f\"{'variant':18s} {'trainable':>9s} {'% of base':>9s} {'memory MB':>9s} {'loud':>6s} {'upper':>6s} {'repeat':>6s} {'s':>4s}\")\n",
    "for name, v in rows.items():\n",
    "    print(f\"{name:18s} {v['trainable']:9,d} {100 * v['trainable'] / N_BASE:8.1f}% {v['mem']:9.2f} {v['acc']['loud']:6.2f} {v['acc']['uppercase']:6.2f} {v['acc']['repeat']:6.2f} {v['time']:4.0f}\")\n",
    "print(f\"{'base (before)':18s} {'-':>9s} {'-':>9s} {'-':>9s} {acc_base['loud']:6.2f} {acc_base['uppercase']:6.2f} {acc_base['repeat']:6.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "### Self-check 5: what must hold"
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "full, l8, l2 = rows[\"full fine-tuning\"], rows[\"LoRA r=8\"], rows[\"LoRA r=2\"]\n",
    "assert acc_base[\"loud\"] <= 0.1, \"the base model has never seen the new instruction\"\n",
    "assert full[\"acc\"][\"loud\"] >= 0.5, \"full fine-tuning should learn the new instruction\"\n",
    "assert l8[\"acc\"][\"loud\"] >= 0.5, \"LoRA with rank 8 should also learn it, with far fewer trainable parameters\"\n",
    "assert l8[\"trainable\"] < 0.2 * full[\"trainable\"] and l2[\"trainable\"] < l8[\"trainable\"]\n",
    "assert l8[\"mem\"] < full[\"mem\"], \"LoRA needs less training memory\"\n",
    "merged8 = merge_lora(l8[\"model\"])\n",
    "same = sum(generate(merged8, prompt(\"loud\", w)) == generate(rows[\"LoRA r=8\"][\"model\"], prompt(\"loud\", w)) for w in eval_words[:10])\n",
    "assert same == 10 and abs(accuracy(merged8, \"loud\") - l8[\"acc\"][\"loud\"]) < 1e-9, \"the merged model gives the same answers\"\n",
    "print(\"self-check 5 passed: LoRA learns the new instruction with a fraction of the trainable parameters, and merges without changing behaviour\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "## 8. Reflect\n",
    "\n",
    "1. Why does `A` receive no gradient at the first step? Why is initialising `B` to zero a good idea, and what would happen if both `A` and `B` were zero?\n",
    "2. Even rank 2 learned this easy instruction. What does the rank limit in principle, which kind of task would make rank 2 fail while rank 8 succeeds, and how would you pick the rank for a real task?\n",
    "3. Compare the old-task columns with the base model. Why might LoRA be expected to forget less than full fine-tuning, and does this run show a difference? (Be careful: 30 words per task and one seed.)\n",
    "4. The estimated training memory differs less than the trainable counts. Why? (Hint: which term does not shrink?) What does quantising the frozen weights, as in QLoRA, change?\n",
    "\n",
    "**Done when** all five self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}