rama2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "import copy\n",
    "torch.set_default_device('mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'input_ids': tensor([[    1,   530, 17986,  6559,  1139,   292, 10742, 29899,  1595,  2512,\n",
      "           868,  2877,   297, 12101,  7208,  1199, 29889, 21724, 29968, 29899,\n",
      "         29941,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "         0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0,\n",
      "         0, 0, 0, 0, 0, 0, 0, 0]])}, {'input_ids': tensor([[    1,  7684,  1655,  ...,   891,   891,   891],\n",
      "        [    1,   425, 10732,  ...,   891,   891,   891],\n",
      "        [    1,  7684,  1655,  ...,   891,   891,   891],\n",
      "        ...,\n",
      "        [    1, 14263, 29895,  ...,   891,   891,   891],\n",
      "        [    1,   425, 10732,  ...,   891,   891,   891],\n",
      "        [    1,   425, 10732,  ...,   891,   891,   891]]), 'token_type_ids': tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0]])}]\n"
     ]
    }
   ],
   "source": [
    "tokenized_train = torch.load(\"tokenized_train.pt\")\n",
    "\n",
    "print(tokenized_train[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_tok = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
    "llama_tok.pad_token_id = 891"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bcfdc2ee8c5b4c719ef6b24d1e289582",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama_mod = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
    "llama_mod.to('mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n",
      "tensor([[[[0., -inf, -inf,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., -inf,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf]]]], device='mps:0')\n",
      "tensor([[[-0.5688, -0.9764,  2.7797,  ..., -5.1534, -3.9327,  0.4894],\n",
      "         [-1.2674, -3.0520,  1.5501,  ..., -2.0579, -5.4599, -1.8719],\n",
      "         [-2.3778, -3.8658, -0.2674,  ..., -2.4051, -4.9045,  0.6033],\n",
      "         ...,\n",
      "         [-1.8827, -3.6941,  0.5543,  ..., -1.9751, -4.6183,  0.0827],\n",
      "         [-1.8824, -3.6932,  0.5546,  ..., -1.9755, -4.6181,  0.0824],\n",
      "         [-1.8818, -3.6917,  0.5554,  ..., -1.9762, -4.6173,  0.0824]]],\n",
      "       device='mps:0') torch.Size([1, 32, 32000])\n"
     ]
    }
   ],
   "source": [
    "class LlamaRetrofit(torch.nn.Module):\n",
    "    def __init__(self, llama):\n",
    "        super().__init__()\n",
    "\n",
    "        self.emb = llama.model.embed_tokens\n",
    "        self.blocks = llama.model.layers\n",
    "        self.norm = llama.model.norm\n",
    "        self.head = llama.lm_head\n",
    "\n",
    "        # deep copy the first 3 blocks to repourpose for the encoder\n",
    "        self.encoder = torch.nn.ModuleList([copy.deepcopy(self.blocks[i]) for i in range(2)])\n",
    "        self.cross_attn = torch.nn.ModuleList([torch.nn.MultiheadAttention(4096, 32, batch_first=True) for _ in range(len(self.blocks)//8)])\n",
    "    \n",
    "\n",
    "    def forward(self, x, context):\n",
    "        x_pad_mask = x['attention_mask']\n",
    "        x = x['input_ids']\n",
    "        context = context['input_ids']\n",
    "        x = self.emb(x)\n",
    "        context = self.emb(context)\n",
    "        for i, block in enumerate(self.encoder):\n",
    "            context = block(context)[0]\n",
    "        # tile the mask to match the batch size (x.size[0]), and unsqueeze on dim 1\n",
    "        mask = torch.full(\n",
    "                (x.shape[0], 1, x.shape[1], x.shape[1]), float(\"-inf\"), device=x.device\n",
    "            )\n",
    "        mask = torch.triu(mask, diagonal=1)\n",
    "        # replace nan with 0\n",
    "        mask = mask.masked_fill(mask != mask, 0.0)\n",
    "        # combine the mask with the padding mask\n",
    "        # attention mask has 0s where there is padding, so it needs to be replaced with -inf\n",
    "        mask = mask.masked_fill(x_pad_mask.unsqueeze(1) == 0, float(\"-inf\"))\n",
    "        print(mask)\n",
    "        for i, block in enumerate(self.blocks):\n",
    "            x = block(x, mask)[0]\n",
    "            if i % 8 == 0:\n",
    "                # take the average of cross attention with each item in the context along the batch dim\n",
    "                crossed = self.cross_attn[i//8](x, context[0].unsqueeze(0), context[0].unsqueeze(0))[0]\n",
    "                for j in range(1, len(context)):\n",
    "                    crossed += self.cross_attn[i//8](x, context[j].unsqueeze(0), context[j].unsqueeze(0))[0]\n",
    "                # average the cross attention\n",
    "                x = torch.mean(crossed, dim=0).unsqueeze(0)\n",
    "\n",
    "\n",
    "        x = self.norm(x)\n",
    "        x = self.head(x)\n",
    "        return x\n",
    "\n",
    "test_model = LlamaRetrofit(llama)\n",
    "# save state dicts for encoder and cross_attn from test_model\n",
    "torch.save(test_model.encoder.state_dict(), 'encoder.pt')\n",
    "torch.save(test_model.cross_attn.state_dict(), 'cross_attn.pt')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LlamaRetrofit(llama_mod)\n",
    "\n",
    "# optimize model.cross_attn.parameters() and model.encoder.parameters()\n",
    "print(tokenized_train[0][0]['input_ids'].device)\n",
    "with torch.no_grad():\n",
    "    retrofit_out = model(tokenized_train[0][0].to('mps'), tokenized_train[0][1].to('mps'))\n",
    "    print(retrofit_out, retrofit_out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[0., -inf, -inf,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., -inf,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf],\n",
      "          [0., 0., 0.,  ..., -inf, -inf, -inf]]]], device='mps:0')\n",
      "| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | |\n",
      "torch.Size([1, 32, 32000]) torch.Size([1, 32])\n",
      "21\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">23</span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">20 │   </span>label = label[:, :firstzero]                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">21 │   </span>output = output[:, :firstzero, :]                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">22 │   </span>loss = loss_fn(torch.permute(output, (<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>,<span style=\"color: #0000ff; text-decoration-color: #0000ff\">2</span>,<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>)), label)                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>23 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   </span>loss.backward()                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">24 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> i % <span style=\"color: #0000ff; text-decoration-color: #0000ff\">16</span> == <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>:                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">25 │   │   </span>optimizer.step()                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">26 │   │   </span>optimizer.zero_grad()                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_tensor.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">482</span> in       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">backward</span>                                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 479 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│   │   │   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">used to compute the attr::tensors.</span>                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 480 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">\"\"\"</span>                                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 481 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> has_torch_function_unary(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>):                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 482 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> handle_torch_function(                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 483 │   │   │   │   </span>Tensor.backward,                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 484 │   │   │   │   </span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>,),                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 485 │   │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>,                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">overrides.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1551</span> in    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">handle_torch_function</span>                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1548 │   │   # if we're here, the mode must be set to a TorchFunctionStackMode</span>                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1549 │   │   # this unsets it and calls directly into TorchFunctionStackMode's torch function</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1550 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">with</span> _pop_mode_temporarily() <span style=\"color: #0000ff; text-decoration-color: #0000ff\">as</span> mode:                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1551 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>result = mode.__torch_function__(public_api, types, args, kwargs)             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1552 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> result <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">NotImplemented</span>:                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1553 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> result                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1554 </span>                                                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/utils/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_device.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">76</span> in  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">__torch_function__</span>                                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">73 │   │   </span>kwargs = kwargs <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> {}                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">74 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> func <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> _device_constructors() <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">and</span> kwargs.get(<span style=\"color: #808000; text-decoration-color: #808000\">'device'</span>) <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">75 │   │   │   </span>kwargs[<span style=\"color: #808000; text-decoration-color: #808000\">'device'</span>] = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.device                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>76 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> func(*args, **kwargs)                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">77 </span>                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">78 # NB: This is directly called from C++ in torch/csrc/Device.cpp</span>                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">79 </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">device_decorator</span>(device, func):                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_tensor.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">491</span> in       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">backward</span>                                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 488 │   │   │   │   </span>create_graph=create_graph,                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 489 │   │   │   │   </span>inputs=inputs,                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 490 │   │   │   </span>)                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 491 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>torch.autograd.backward(                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 492 │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, gradient, retain_graph, create_graph, inputs=inputs                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 493 │   │   </span>)                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 494 </span>                                                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/autograd/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">__init__.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">25</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">backward</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">248 │   # The reason we repeat the same comment below is that</span>                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">249 │   # some Python versions print out the first line of a multi-line function</span>               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">250 │   # calls in the traceback and some print out the last line</span>                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>251 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   </span>Variable._execution_engine.run_backward(  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># Calls into the C++ engine to run the bac</span>   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">252 │   │   </span>tensors,                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">253 │   │   </span>grad_tensors_,                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">254 │   │   </span>retain_graph,                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">KeyboardInterrupt</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92m<module>\u001b[0m:\u001b[94m23\u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m20 \u001b[0m\u001b[2m│   \u001b[0mlabel = label[:, :firstzero]                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m21 \u001b[0m\u001b[2m│   \u001b[0moutput = output[:, :firstzero, :]                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m22 \u001b[0m\u001b[2m│   \u001b[0mloss = loss_fn(torch.permute(output, (\u001b[94m0\u001b[0m,\u001b[94m2\u001b[0m,\u001b[94m1\u001b[0m)), label)                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m23 \u001b[2m│   \u001b[0mloss.backward()                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m24 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mif\u001b[0m i % \u001b[94m16\u001b[0m == \u001b[94m0\u001b[0m:                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m25 \u001b[0m\u001b[2m│   │   \u001b[0moptimizer.step()                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m26 \u001b[0m\u001b[2m│   │   \u001b[0moptimizer.zero_grad()                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/\u001b[0m\u001b[1;33m_tensor.py\u001b[0m:\u001b[94m482\u001b[0m in       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[92mbackward\u001b[0m                                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 479 \u001b[0m\u001b[2;33m│   │   │   │   \u001b[0m\u001b[33mused to compute the attr::tensors.\u001b[0m                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 480 \u001b[0m\u001b[2;33m│   │   \u001b[0m\u001b[33m\"\"\"\u001b[0m                                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 481 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m has_torch_function_unary(\u001b[96mself\u001b[0m):                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 482 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m handle_torch_function(                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 483 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mTensor.backward,                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 484 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m(\u001b[96mself\u001b[0m,),                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 485 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[96mself\u001b[0m,                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/\u001b[0m\u001b[1;33moverrides.py\u001b[0m:\u001b[94m1551\u001b[0m in    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[92mhandle_torch_function\u001b[0m                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1548 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# if we're here, the mode must be set to a TorchFunctionStackMode\u001b[0m                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1549 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this unsets it and calls directly into TorchFunctionStackMode's torch function\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1550 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mwith\u001b[0m _pop_mode_temporarily() \u001b[94mas\u001b[0m mode:                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1551 \u001b[2m│   │   │   \u001b[0mresult = mode.__torch_function__(public_api, types, args, kwargs)             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1552 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m result \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[96mNotImplemented\u001b[0m:                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1553 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m result                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1554 \u001b[0m                                                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/utils/\u001b[0m\u001b[1;33m_device.py\u001b[0m:\u001b[94m76\u001b[0m in  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[92m__torch_function__\u001b[0m                                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m73 \u001b[0m\u001b[2m│   │   \u001b[0mkwargs = kwargs \u001b[95mor\u001b[0m {}                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m74 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m func \u001b[95min\u001b[0m _device_constructors() \u001b[95mand\u001b[0m kwargs.get(\u001b[33m'\u001b[0m\u001b[33mdevice\u001b[0m\u001b[33m'\u001b[0m) \u001b[95mis\u001b[0m \u001b[94mNone\u001b[0m:                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m75 \u001b[0m\u001b[2m│   │   │   \u001b[0mkwargs[\u001b[33m'\u001b[0m\u001b[33mdevice\u001b[0m\u001b[33m'\u001b[0m] = \u001b[96mself\u001b[0m.device                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m76 \u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m func(*args, **kwargs)                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m77 \u001b[0m                                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m78 \u001b[0m\u001b[2m# NB: This is directly called from C++ in torch/csrc/Device.cpp\u001b[0m                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m79 \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mdevice_decorator\u001b[0m(device, func):                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/\u001b[0m\u001b[1;33m_tensor.py\u001b[0m:\u001b[94m491\u001b[0m in       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[92mbackward\u001b[0m                                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 488 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mcreate_graph=create_graph,                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 489 \u001b[0m\u001b[2m│   │   │   │   \u001b[0minputs=inputs,                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 490 \u001b[0m\u001b[2m│   │   │   \u001b[0m)                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 491 \u001b[2m│   │   \u001b[0mtorch.autograd.backward(                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 492 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[96mself\u001b[0m, gradient, retain_graph, create_graph, inputs=inputs                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 493 \u001b[0m\u001b[2m│   │   \u001b[0m)                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 494 \u001b[0m                                                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/autograd/\u001b[0m\u001b[1;33m__init__.py\u001b[0m:\u001b[94m25\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m1\u001b[0m in \u001b[92mbackward\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m248 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# The reason we repeat the same comment below is that\u001b[0m                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m249 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# some Python versions print out the first line of a multi-line function\u001b[0m               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m250 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# calls in the traceback and some print out the last line\u001b[0m                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m251 \u001b[2m│   \u001b[0mVariable._execution_engine.run_backward(  \u001b[2m# Calls into the C++ engine to run the bac\u001b[0m   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m252 \u001b[0m\u001b[2m│   │   \u001b[0mtensors,                                                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m253 \u001b[0m\u001b[2m│   │   \u001b[0mgrad_tensors_,                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m254 \u001b[0m\u001b[2m│   │   \u001b[0mretain_graph,                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
       "\u001b[1;91mKeyboardInterrupt\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# optimize both model.cross_attn.parameters() and model.encoder.parameters()\n",
    "params_to_optimize = list(model.cross_attn.parameters()) + list(model.encoder.parameters())\n",
    "optimizer = torch.optim.Adam(params_to_optimize, lr=1e-4)\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "optimizer.zero_grad()\n",
    "for i in range(len(tokenized_train)):\n",
    "    # label is shifted by 1 with an extra padding token, from tokenized_train[i][0]['input_ids']\n",
    "    label = torch.roll(tokenized_train[i][0]['input_ids'], -1).to('mps')\n",
    "\n",
    "    output = model(tokenized_train[i][0].to('mps'), tokenized_train[i][1].to('mps'))\n",
    "    # argmax and decode\n",
    "    stringy = torch.argmax(output, dim=-1)\n",
    "    stringy = llama_tok.decode(stringy[0])\n",
    "    print(stringy)\n",
    "    print(output.shape, label.shape)\n",
    "    # get the first unmasked token and clip the label and output, this can be done by getting the first '0' in the attention mask\n",
    "    firstzero = torch.argmax(1 - tokenized_train[i][0]['attention_mask']).item()\n",
    "    print(firstzero)\n",
    "    label = label[:, :firstzero]\n",
    "    output = output[:, :firstzero, :]\n",
    "    # print the label as well\n",
    "    print(llama_tok.decode(label[0]))\n",
    "    loss = loss_fn(torch.permute(output, (0,2,1)), label)\n",
    "    loss.backward()\n",
    "    if i % 16 == 0:\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    print(loss.item())\n",
    "\n",
    "# save weights\n",
    "torch.save(model.state_dict(), 'model.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log current folder location\n",
    "import os\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "optimizer = torch.optim.Adam(model.cross_attn.parameters(), lr=1e-4)\n",
    "optimizer.zero_grad()\n",
    "\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "for i in range(len(tokenized_train)):\n",
    "    label = tokenized_train[i][0][0][1:]\n",
    "    # add a padding token to the beginning of the label\n",
    "    label = torch.cat([torch.tensor([0]), label])\n",
    "    inp = torch.cat([tokenized_train[i][0], tokenized_train[i][1]])\n",
    "    output = model(inp)\n",
    "    # argmax and decode\n",
    "    stringy = torch.argmax(output, dim=-1)\n",
    "    stringy = llama_tok.decode(stringy)\n",
    "    print(stringy)\n",
    "    loss = loss_fn(output, label)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    print(loss.item())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130, 16020, 16020, 16020, 16020, 16020, 16020, 16020, 16020,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130, 16020, 16020, 16020, 16020,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130, 27823,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130, 16020,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130, 16020,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130, 17856, 17856, 17856,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130, 16020,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130], device='mps:0')\n",
      "ilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aqu Aqu Aqu Aqu Aqu Aqu Aqu Aquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aqu Aqu Aqu AquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonNililonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aquiloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloninéinéinéilonilonilonilonilonilonilonilonilonilon Aquilonilonilonilonilonilonilonilonilonilonilonilon\n"
     ]
    }
   ],
   "source": [
    "print(stringy)\n",
    "print(llama_tok.decode(stringy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (28985 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 28985])\n",
      "torch.Size([1, 30485])\n"
     ]
    }
   ],
   "source": [
    "# test_adapter = LlamaT5Adapter(t5_mod, llama_mod, ctx_num=16)\n",
    "\n",
    "# optimizer = torch.optim.Adam(test_adapter.parameters(), lr=1e-4)\n",
    "# sample_out = test_adapter(torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]), torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]))\n",
    "# print(sample_out.shape)\n",
    "\n",
    "\n",
    "# sample_out.sum().backward()\n",
    "# optimizer.step()\n",
    "\n",
    "import sqlite3\n",
    "\n",
    "conn = sqlite3.connect('data.db')\n",
    "\n",
    "# create a cursor object\n",
    "cur = conn.cursor()\n",
    "\n",
    "cur.execute('SELECT content, author FROM discord ORDER BY timestamp')\n",
    "data = cur.fetchall()\n",
    "\n",
    "# concat author name to all messages\n",
    "data = [f'{author}: {content}' for content, author in data]\n",
    "\n",
    "# combine all messages into one string\n",
    "data = '\\n\\n'.join(data)\n",
    "\n",
    "data = data[:100000]\n",
    "\n",
    "t5_tokenized_data = t5_tok.encode(data, return_tensors='pt')\n",
    "llama_tokenized_data = llama_tok.encode(data, return_tensors='pt')\n",
    "print(t5_tokenized_data.shape)\n",
    "print(llama_tokenized_data.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, batch 10240\n",
      "Loss: 11.575541496276855\n",
      "Epoch 0, batch 11264\n",
      "Loss: 9.961394309997559\n",
      "Epoch 0, batch 12288\n"
     ]
    }
   ],
   "source": [
    "\n",
    "adapter = LlamaT5Adapter(t5_mod, llama_mod, ctx_num=8)\n",
    "\n",
    "\n",
    "params_to_optimize = list(adapter.llama_queries.parameters()) + list(adapter.t5_keys.parameters()) + list(adapter.t5_values.parameters())\n",
    "optimizer = torch.optim.Adam(params_to_optimize, lr=1e-4)\n",
    "# optimizer = torch.optim.Adam(adapter.parameters(), lr=1e-4)\n",
    "\n",
    "total_llama_tokens = llama_tokenized_data.shape[1]\n",
    "\n",
    "# for the training loop, for each chunk passed to the llama decoder, the previous 16 chunks of 512 t5-tokens are passed to the t5 encoder\n",
    "epochs = 4\n",
    "llama_batch_size = 1024\n",
    "t5_batch_size = 512\n",
    "for epoch in range(epochs):\n",
    "    for i in range(t5_batch_size*20, total_llama_tokens, llama_batch_size):\n",
    "        print(f'Epoch {epoch}, batch {i}')\n",
    "        # get the next batch of data\n",
    "        llama_in = llama_tokenized_data[:, i:i+llama_batch_size]\n",
    "        # to get approximately the same spot in the chat for the t5 tokenization, use the current index/total_llama_tokens, and multiply by the total t5 tokens\n",
    "        t5_index = int(i/total_llama_tokens*t5_tokenized_data.shape[1])\n",
    "        t5_in = t5_tokenized_data[:, t5_index-t5_batch_size*8:t5_index]\n",
    "        # reshape the t5 input to be 16 chunks of 512 tokens\n",
    "        t5_in = t5_in.reshape(-1, t5_batch_size)\n",
    "        # pass the data through the model\n",
    "        out = adapter(llama_in, t5_in)\n",
    "        # calculate loss by shifting the target data by 1\n",
    "        loss = torch.nn.functional.cross_entropy(out[:, :-1].reshape(-1, out.shape[-1]), llama_in[:, 1:].reshape(-1))\n",
    "        # backpropagate\n",
    "        loss.backward()\n",
    "        # clip gradients\n",
    "        torch.nn.utils.clip_grad_norm_(params_to_optimize, 1.0)\n",
    "        # update parameters\n",
    "        optimizer.step()\n",
    "        # zero gradients\n",
    "        optimizer.zero_grad()\n",
    "        print(f'Loss: {loss.item()}')\n",
    "        del llama_in, t5_in, out, loss\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}