test_llama.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "from attacks import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"alespalla/chatbot_instruction_prompts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'response': '. Determine the cause of the clog. Clogs in kitchen drains can be caused by food, grease, soap scum, and other debris.\\n\\n2. If you have a plunger, try using it to dislodge the clog. Make sure to cover any drain and overflow openings with a towel to create a tight seal. Use your plunger in a steady up-and-down motion.\\n\\n3. If a plunger doesn’t remove the clog, you may need to use a drain auger (also known as a “snake”) to remove it. You can rent one from a hardware store or purchase one online.\\n\\n4. If a drain auger fails to remove the clog, you may need to disassemble the pipe and use a hand-held plumber’s snake and/or rigid wire to try and break up the clog.\\n\\n5. Consider calling a professional plumber if your kitchen drain continues to remain clogged after all of these steps.',\n",
       " 'prompt': 'What’s the best way to fix my kitchen drain?\\n\\n             1'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(dataset['train']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> initializing model parallel with size 1\n",
      "> initializing ddp with size 1\n",
      "> initializing pipeline with size 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Shivaen\\anaconda3\\envs\\steal-lm\\Lib\\site-packages\\torch\\__init__.py:696: UserWarning: torch.set_default_tensor_type() is deprecated as of PyTorch 2.1, please use torch.set_default_dtype() and torch.set_default_device() as alternatives. (Triggered internally at C:\\cb\\pytorch_1000000000000\\work\\torch\\csrc\\tensor\\python_tensor.cpp:453.)\n",
      "  _C._set_default_tensor_type(t)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded in 10.86 seconds\n"
     ]
    }
   ],
   "source": [
    "MAX_SEQ_LEN = 1024\n",
    "llama = LlamaAPI.build(\n",
    "    ckpt_dir=\"../llama/llama-2-7b\",\n",
    "    tokenizer_path=\"../llama/tokenizer.model\",\n",
    "    max_seq_len=MAX_SEQ_LEN,\n",
    "    max_batch_size=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## All Logits Techniques"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hidden Dimension Extraction Attack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [03:14<00:00, 25.68it/s]\n"
     ]
    }
   ],
   "source": [
    "q = get_q(\n",
    "    llama=llama,\n",
    "    dataset=dataset['train'],\n",
    "    text_key='prompt',\n",
    "    batch_size=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uses RMSNorm\n",
      "Hidden Dim: 4095\n"
     ]
    }
   ],
   "source": [
    "# llama 2 7b hidden dimension: 4096\n",
    "u, s, s_dim, pred_dim = h_dim_extraction(\n",
    "    q=q,\n",
    "    predict_norm=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([32000, 5000]), torch.Size([5000]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.shape, s.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "note that it also appears that the singular values flip to negative at the index of the actual h_dim -> singular values < 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.9369, -1.0427], dtype=torch.float64)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_dim[4094:4096]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Last Layer Extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = llama.model.output.weight.clone().detach()\n",
    "pred_w, g = layer_extraction(\n",
    "    w=w,\n",
    "    u=u,\n",
    "    s=s,\n",
    "    h_dim=4096\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RMS: tensor(0.0002, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# predicted weight rms (paper's RMS for Llama 7b: 8 * 10e-5)\n",
    "print(\"RMS:\", torch.sqrt(torch.mean((pred_w@g - w) ** 2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "del u, s, s_dim, w, pred_w, g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Top-K Logits/Logprobs and Logit-bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Access to top-k logits with logit bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [1:32:39<00:00,  1.11s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uses RMSNorm\n",
      "Hidden Dim: 4095\n",
      "RMS: tensor(0.0002, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "q = get_q(\n",
    "    llama=llama,\n",
    "    dataset=dataset['train'],\n",
    "    text_key='prompt',\n",
    "    batch_size=1,\n",
    "    logit_fn=topk_logit_extraction\n",
    ")\n",
    "\n",
    "# llama 2 7b hidden dimension: 4096\n",
    "u, s, s_dim, pred_dim = h_dim_extraction(\n",
    "    q=q,\n",
    "    predict_norm=True\n",
    ")\n",
    "\n",
    "w = llama.model.output.weight.clone().detach()\n",
    "pred_w, g = layer_extraction(\n",
    "    w=w,\n",
    "    u=u,\n",
    "    s=s,\n",
    "    h_dim=4096\n",
    ")\n",
    "\n",
    "# predicted weight rms (paper's RMS for Llama 7b: 8 * 10e-5)\n",
    "print(\"RMS:\", torch.sqrt(torch.mean((pred_w@g - w) ** 2)))\n",
    "\n",
    "del u, s, s_dim, w, pred_w, g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Access to top-k logprobs with logit bias\n",
    "- note: hidden dim prediction fails with relative logits\n",
    "- RMS error is 3x greater than before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [2:59:14<00:00,  2.15s/it]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uses RMSNorm\n",
      "Hidden Dim: 1\n",
      "RMS: tensor(0.0006, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "q = get_q(\n",
    "    llama=llama,\n",
    "    dataset=dataset['train'],\n",
    "    text_key='prompt',\n",
    "    batch_size=1,\n",
    "    logit_fn=partial(topk_logprob_ref_extraction, k=500)  # higher k for POC\n",
    ")\n",
    "\n",
    "# llama 2 7b hidden dimension: 4096\n",
    "u, s, s_dim, pred_dim = h_dim_extraction(\n",
    "    q=q,\n",
    "    predict_norm=True\n",
    ")\n",
    "\n",
    "w = llama.model.output.weight.clone().detach()\n",
    "pred_w, g = layer_extraction(\n",
    "    w=w,\n",
    "    u=u,\n",
    "    s=s,\n",
    "    h_dim=4096\n",
    ")\n",
    "\n",
    "# predicted weight rms (paper's RMS for Llama 7b: 8 * 10e-5)\n",
    "print(\"RMS:\", torch.sqrt(torch.mean((pred_w@g - w) ** 2)))\n",
    "\n",
    "del u, s, s_dim, w, pred_w, g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Access to top-k logprobs with logit bias; cost-optimal\n",
    "- Generate more than 1 token and consider each generation a new sample\n",
    "- RMS error is 4x the baseline method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [2:08:18<00:00,  1.54s/it]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model uses RMSNorm\n",
      "Hidden Dim: 1\n",
      "RMS: tensor(0.0008, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "exp_size = 10  # number of tokens to generate for a given prompt\n",
    "q = get_q(\n",
    "    llama=llama,\n",
    "    dataset=dataset['train'],\n",
    "    text_key='prompt',\n",
    "    n=5000,\n",
    "    batch_size=1,\n",
    "    logit_fn=partial(topk_logprob_exp_extraction, k=500, exp_size=10)  # higher k for POC\n",
    ")\n",
    "\n",
    "# llama 2 7b hidden dimension: 4096\n",
    "u, s, s_dim, pred_dim = h_dim_extraction(\n",
    "    q=q,\n",
    "    predict_norm=True\n",
    ")\n",
    "\n",
    "w = llama.model.output.weight.clone().detach()\n",
    "pred_w, g = layer_extraction(\n",
    "    w=w,\n",
    "    u=u,\n",
    "    s=s,\n",
    "    h_dim=4096\n",
    ")\n",
    "\n",
    "# predicted weight rms (paper's RMS for Llama 7b: 8 * 10e-5)\n",
    "print(\"RMS:\", torch.sqrt(torch.mean((pred_w@g - w) ** 2)))\n",
    "\n",
    "del u, s, s_dim, w, pred_w, g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "steal-lm",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}