{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LoRA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're going to train a very simple LORA that, when applied, will make our model always predict \"Paris\" no matter what."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from nnsight import LanguageModel\n",
    "\n",
    "model = LanguageModel('openai-community/gpt2', device_map='auto')\n",
    "\n",
    "\n",
    "from nnsight.envoy import Envoy # \n",
    "\n",
    "# We will define a LORA class.\n",
    "# The LORA class call method operations are simply traced like you would normally do in a .trace.\n",
    "class LORA(nn.Module):\n",
    "    def __init__(self, module: Envoy, dim: int, r: int) -> None:\n",
    "        \"\"\"Init.\n",
    "\n",
    "        Args:\n",
    "            module (Envoy): Which model Module we are adding the LORA to.\n",
    "            dim (int): Dimension of the layer we are adding to (This could potentially be auto populated if the user scanned first so we know the shape)\n",
    "            r (int): Inner dimension of the LORA\n",
    "        \"\"\"\n",
    "        super(LORA, self).__init__()\n",
    "        self.r = r\n",
    "        self.module = module\n",
    "        self.WA = torch.nn.Parameter(torch.randn(dim, self.r), requires_grad=True).save()\n",
    "        self.WB = torch.nn.Parameter(torch.zeros(self.r, dim), requires_grad=True).save()\n",
    "\n",
    "    # The Call method defines how to actually apply the LORA.\n",
    "    def __call__(self, alpha: float = 1.0):\n",
    "        \"\"\"Call.\n",
    "\n",
    "        Args:\n",
    "            alpha (float, optional): How much to apply the LORA. Can be altered after training for inference. Defaults to 1.0.\n",
    "        \"\"\"\n",
    "\n",
    "        # We apply WA to the first positional arg (the hidden states)\n",
    "        A_x = torch.matmul(self.module.input[0][0], self.WA)\n",
    "        BA_x = torch.matmul(A_x, self.WB)\n",
    "\n",
    "        # LORA is additive\n",
    "        h = BA_x + self.module.output\n",
    "\n",
    "        # Replace the output with our new one * alpha\n",
    "        # Could also have been self.module.output[:] = h * alpha, for in-place\n",
    "        self.module.output = h * alpha\n",
    "\n",
    "    def parameters(self):\n",
    "        # Some way to get all the parameters.\n",
    "        return [self.WA, self.WB]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define all the variables to use in LORA training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need the token id of the correct answer.\n",
    "answer = \" Paris\"\n",
    "answer_token = model.tokenizer.encode(answer)[0]\n",
    "# Inner LORA dimension\n",
    "lora_dim = 4\n",
    "# Module to train LORA on\n",
    "module = model.transformer.h[-1].mlp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the `.scan()` method to get the shape of the module without having to fully run the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "768\n"
     ]
    }
   ],
   "source": [
    "with model.scan(\" \"):\n",
    "    dim = module.output.shape[-1]\n",
    "\n",
    "print(dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's time to run the LORA training loop! We using the **Session** and the **Iterator** contexts to achieve this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter containing:\n",
      "tensor([[ 0.5262, -0.6452,  0.8448,  0.7407],\n",
      "        [-0.4497, -0.7200, -1.0452,  0.0630],\n",
      "        [ 0.7231,  1.0991,  0.3883,  0.1719],\n",
      "        ...,\n",
      "        [ 0.0024, -1.1490, -0.5580, -0.9070],\n",
      "        [-0.1946,  0.8469, -1.8173,  0.8333],\n",
      "        [ 0.1722, -1.8518, -1.5542, -1.3361]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[ 0.6813, -0.4550,  0.9903,  0.5476],\n",
      "        [-0.3310, -0.5932, -0.9087, -0.0441],\n",
      "        [ 0.7201,  1.0849,  0.3954,  0.1480],\n",
      "        ...,\n",
      "        [ 0.1580, -0.9589, -0.3856, -1.0354],\n",
      "        [-0.3153,  0.6950, -1.8893,  0.9347],\n",
      "        [ 0.4812, -1.4821, -1.1935, -1.6101]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-1.2552, -2.3574, -0.9555,  2.4472],\n",
      "        [-2.2370, -2.4913, -2.7973,  1.8731],\n",
      "        [-1.2148, -0.8610, -1.5298,  2.0569],\n",
      "        ...,\n",
      "        [-1.7628, -2.8462, -2.2901,  0.9117],\n",
      "        [ 1.6102,  2.5902,  0.0834, -1.0093],\n",
      "        [-1.4495, -3.3539, -3.0740,  0.3545]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-2.7871, -3.8562, -2.4963,  3.9433],\n",
      "        [-3.7392, -3.9859, -4.2827,  3.3862],\n",
      "        [-2.7453, -2.4021, -3.0508,  3.5621],\n",
      "        ...,\n",
      "        [-3.2794, -4.3303, -3.7909,  2.4538],\n",
      "        [ 3.1313,  4.0819,  1.6503, -2.5484],\n",
      "        [-2.9757, -4.8230, -4.5514,  1.9135]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-4.0300, -5.0671, -3.7480,  5.1516],\n",
      "        [-4.9534, -5.1927, -5.4806,  4.6110],\n",
      "        [-3.9870, -3.6541, -4.2833,  4.7793],\n",
      "        ...,\n",
      "        [-4.5076, -5.5269, -5.0037,  3.7067],\n",
      "        [ 4.3638,  5.2859,  2.9272, -3.7984],\n",
      "        [-4.2132, -6.0050, -5.7416,  3.1828]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-5.0524, -6.0584, -4.7789,  6.1404],\n",
      "        [-5.9479, -6.1800, -6.4593,  5.6158],\n",
      "        [-5.0084, -4.6854, -5.2958,  5.7769],\n",
      "        ...,\n",
      "        [-5.5156, -6.5044, -5.9969,  4.7388],\n",
      "        [ 5.3761,  6.2705,  3.9826, -4.8277],\n",
      "        [-5.2302, -6.9683, -6.7128,  4.2308]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-5.8993, -6.8751, -5.6339,  6.9546],\n",
      "        [-6.7678, -6.9929, -7.2638,  6.4456],\n",
      "        [-5.8543, -5.5411, -6.1332,  6.5998],\n",
      "        ...,\n",
      "        [-6.3486, -7.3077, -6.8154,  5.5951],\n",
      "        [ 6.2131,  7.0807,  4.8615, -5.6812],\n",
      "        [-6.0719, -7.7578, -7.5100,  5.1024]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-6.6025, -7.5490, -6.3451,  7.6261],\n",
      "        [-7.4447, -7.6631, -7.9259,  7.1322],\n",
      "        [-6.5568, -6.2529, -6.8272,  7.2799],\n",
      "        ...,\n",
      "        [-7.0383, -7.9686, -7.4911,  6.3074],\n",
      "        [ 6.9068,  7.7484,  5.5957, -6.3908],\n",
      "        [-6.7700, -8.4054, -8.1650,  5.8296]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-7.1858, -8.1039, -6.9361,  8.1787],\n",
      "        [-8.0026, -8.2144, -8.4693,  7.6995],\n",
      "        [-7.1395, -6.8447, -7.4018,  7.8409],\n",
      "        ...,\n",
      "        [-7.6084, -8.5108, -8.0477,  6.8995],\n",
      "        [ 7.4809,  8.2972,  6.2091, -6.9804],\n",
      "        [-7.3483, -8.9347, -8.7015,  6.4362]], requires_grad=True)\n",
      "Parameter containing:\n",
      "tensor([[-7.6675, -8.5581, -7.4253,  8.6306],\n",
      "        [-8.4597, -8.6652, -8.9124,  8.1657],\n",
      "        [-7.6208, -7.3349, -7.8752,  8.3011],\n",
      "        ...,\n",
      "        [-8.0775, -8.9528, -8.5035,  7.3898],\n",
      "        [ 7.9537,  8.7455,  6.7201, -7.4682],\n",
      "        [-7.8253, -9.3640, -9.1378,  6.9405]], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# The LORA object itself isn't transmitted to the server. Only the forward / call method. \n",
    "# The parameters are created remotely and never sent only retrieved \n",
    "with model.session() as session:\n",
    "\n",
    "    # Create dataset of 100 pairs of a blank prompt and the \" Paris \" id\n",
    "    dataset = [[\"_\", answer_token]] * 100\n",
    "\n",
    "    # Create a dataloader from it.\n",
    "    dataloader = DataLoader(dataset, batch_size=10)\n",
    "\n",
    "    # Create our LORA on the last mlp\n",
    "    lora = LORA(module, dim, lora_dim)\n",
    "\n",
    "    # Create an optimizer. Use the parameters from LORA\n",
    "    optimizer = torch.optim.AdamW(lora.parameters(), lr=3)\n",
    "\n",
    "    # Iterate over dataloader using .iter. \n",
    "    with session.iter(dataloader, return_context=True) as (batch, iterator):\n",
    "\n",
    "        prompt = batch[0]\n",
    "        correct_token = batch[1]\n",
    "\n",
    "        # Run .trace with prompt\n",
    "        with model.trace(prompt) as tracer:\n",
    "\n",
    "            # Apply LORA to intervention graph just by calling it with .trace\n",
    "            lora()\n",
    "\n",
    "            # Get logits\n",
    "            logits = model.lm_head.output\n",
    "\n",
    "            # Do cross entropy on last predicted token and correct_token\n",
    "            loss = torch.nn.functional.cross_entropy(logits[:, -1], batch[1])\n",
    "            # Call backward\n",
    "            loss.backward()\n",
    "\n",
    "        # Call methods on optimizer. Graphs that arent from .trace (so in this case session and iterator both have their own graph) are executed sequentially.\n",
    "        # The Graph of Iterator here will be:\n",
    "        # 1.) Index batch at 0 for prompt\n",
    "        # 2.) Index batch at 1 for correct_token\n",
    "        # 3.) Execute the .trace using the prompt\n",
    "        # 4.) Call .step() on optimizer\n",
    "        optimizer.step()\n",
    "        # 5.) Call .zero_grad() in optimizer\n",
    "        optimizer.zero_grad()\n",
    "        # 6.) Print out the lora WA weights to show they are indeed changing\n",
    "        iterator.log(lora.WA)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now `WA` and `WB` are optimized! So we generate with the lora just by calling `lora()` in the `.generate` and save the output to then de-tokenize it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Hello Paris']\n",
      "['Hello,']\n"
     ]
    }
   ],
   "source": [
    "# With lora. Should produce \"Hello Paris\"\n",
    "with model.generate(\"Hello\") as generator:\n",
    "\n",
    "    lora()\n",
    "\n",
    "    out = model.generator.output.save()\n",
    "\n",
    "print(model.tokenizer.batch_decode(out.value))\n",
    "\n",
    "# Then without. Should produce \"Hello,\"\n",
    "with model.generate(\"Hello\") as generator:\n",
    "\n",
    "    out = model.generator.output.save()\n",
    "\n",
    "print(model.tokenizer.batch_decode(out.value))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nnsight_local",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
