examples/loading_saes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Autoencoders\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we show the syntax to load the different autoencoders that we support. Autoencoders are \"attached\" to the model using the `edit` method from NNsight."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GemmaScope SAEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nnsight import LanguageModel\n",
    "from sae_auto_interp.autoencoders import load_gemma_autoencoders\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4104e9252a54a0bbcdc8255c11c5077",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# Load the model\n",
    "model = LanguageModel(\"google/gemma-2-9b\", device_map=\"cuda\", dispatch=True,torch_dtype=\"bfloat16\")\n",
    "\n",
    "# Load the autoencoders, the function returns a dictionary of the submodules with the autoencoders and the edited model.\n",
    "# it takes as arguments the model, the layers to load the autoencoders into,\n",
    "# the average L0 sparsity per layer, the size of the autoencoders and the type of autoencoders (residuals or MLPs).\n",
    "\n",
    "submodule_dict,model = load_gemma_autoencoders(\n",
    "            model,\n",
    "            ae_layers=[10],\n",
    "            average_l0s={10: 47},\n",
    "            size=\"131k\",\n",
    "            type=\"res\"\n",
    "        )\n",
    "\n",
    "# The autoencoder is loaded into the submodules dictionary and the model is edited in place.\n",
    "autoencoder = submodule_dict[\".model.layers.10\"].ae\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EleutherAI SAEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nnsight import LanguageModel\n",
    "from sae_auto_interp.autoencoders import load_eai_autoencoders\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38833c797f6a4ceab937fb178ae5cf8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad17a4bd147c407e93f072c9ceace980",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 2 files:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89a94e1a69dd4071818fe9942ad20f7a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 2 files:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "model = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\", device_map=\"cpu\",dispatch=True,torch_dtype=\"bfloat16\")\n",
    "    \n",
    "# The load function takes as arguments the model, the layers to load the autoencoders into,\n",
    "# The path of the weights (it can be a huggingface path or a local path),\n",
    "# the type of module to load the autoencoders into (residuals or MLPs),\n",
    "# whether to randomize the autoencoders or not,\n",
    "# whether to use the trained k or a specific k on the top-k activation.\n",
    "\n",
    "submodule_dict,model = load_eai_autoencoders(\n",
    "    model,\n",
    "    [23,29],\n",
    "    weight_dir=\"EleutherAI/sae-llama-3.1-8b-64x\",\n",
    "    module=\"mlp\",\n",
    "    randomize=False,\n",
    "    k=None\n",
    ")\n",
    "\n",
    "# The autoencoder is loaded into the submodules dictionary and the model is edited in place.\n",
    "autoencoder = submodule_dict[\".model.layers.23.mlp\"].ae"
   ]
  },
  {
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}