alpha_easy_training_code.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/logan/lib/python3.10/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"
     ]
    }
   ],
   "source": [
    "import torch \n",
    "import argparse\n",
    "from utils import dotdict\n",
    "from activation_dataset import setup_token_data\n",
    "import wandb\n",
    "import json\n",
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "from einops import rearrange\n",
    "# from standard_metrics import run_with_model_intervention, perplexity_under_reconstruction, mean_nonzero_activations\n",
    "# Create \n",
    "# # make an argument parser directly below\n",
    "# parser = argparse.ArgumentParser()\n",
    "# parser.add_argument(\"--model_name\", type=str, default=\"EleutherAI/pythia-70m-deduped\")\n",
    "# parser.add_argument(\"--layer\", type=int, default=4)\n",
    "# parser.add_argument(\"--setting\", type=str, default=\"residual\")\n",
    "# parser.add_argument(\"--l1_alpha\", type=float, default=3e-3)\n",
    "# parser.add_argument(\"--num_epochs\", type=int, default=10)\n",
    "# parser.add_argument(\"--model_batch_size\", type=int, default=4)\n",
    "# parser.add_argument(\"--lr\", type=float, default=1e-3)\n",
    "# parser.add_argument(\"--kl\", type=bool, default=False)\n",
    "# parser.add_argument(\"--reconstruction\", type=bool, default=False)\n",
    "# parser.add_argument(\"--dataset_name\", type=str, default=\"NeelNanda/pile-10k\")\n",
    "# parser.add_argument(\"--device\", type=str, default=\"cuda:4\")\n",
    "\n",
    "# args = parser.parse_args()\n",
    "cfg = dotdict()\n",
    "# cfg.model_name=\"EleutherAI/pythia-70m-deduped\"\n",
    "# cfg.model_name=\"usvsnsp/pythia-6.9b-rm-full-hh-rlhf\"\n",
    "cfg.model_name=\"reciprocate/dahoas-gptj-rm-static\"\n",
    "cfg.layers=[i for i in range(10)]\n",
    "cfg.setting=\"residual\"\n",
    "# cfg.tensor_name=\"gpt_neox.layers.{layer}\"\n",
    "cfg.tensor_name=\"transformer.h.{layer}\"\n",
    "cfg.l1_alpha=[8e-4, 8e-4]\n",
    "cfg.sparsity=None\n",
    "cfg.model_batch_size=4\n",
    "cfg.lr=1e-3\n",
    "cfg.kl=False\n",
    "cfg.reconstruction=False\n",
    "# cfg.dataset_name=\"NeelNanda/pile-10k\"\n",
    "# cfg.dataset_name=\"Skylion007/openwebtext\"\n",
    "cfg.dataset_name=\"Elriggs/openwebtext-100k\"\n",
    "cfg.device=\"cuda:0\"\n",
    "cfg.ratio = 4\n",
    "# cfg.device=\"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensor_names = [cfg.tensor_name.format(layer=layer) for layer in cfg.layers]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading (…)lve/main/config.json: 100%|██████████| 567/567 [00:00<00:00, 1.86MB/s]\n",
      "Downloading model.safetensors: 100%|██████████| 166M/166M [00:05<00:00, 30.6MB/s] \n",
      "Downloading (…)okenizer_config.json: 100%|██████████| 396/396 [00:00<00:00, 1.27MB/s]\n",
      "Downloading (…)/main/tokenizer.json: 100%|██████████| 2.11M/2.11M [00:00<00:00, 11.1MB/s]\n",
      "Downloading (…)cial_tokens_map.json: 100%|██████████| 99.0/99.0 [00:00<00:00, 340kB/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "GPTNeoXForCausalLM(\n",
       "  (gpt_neox): GPTNeoXModel(\n",
       "    (embed_in): Embedding(50304, 512)\n",
       "    (layers): ModuleList(\n",
       "      (0-5): 6 x GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (attention): GPTNeoXAttention(\n",
       "          (rotary_emb): RotaryEmbedding()\n",
       "          (query_key_value): Linear(in_features=512, out_features=1536, bias=True)\n",
       "          (dense): Linear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (embed_out): Linear(in_features=512, out_features=50304, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load in the model\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "model = AutoModelForCausalLM.from_pretrained(cfg.model_name).to(cfg.device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading metadata: 100%|██████████| 921/921 [00:00<00:00, 4.32MB/s]\n",
      "Downloading readme: 100%|██████████| 373/373 [00:00<00:00, 661kB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset None/None (download: 31.72 MiB, generated: 58.43 MiB, post-processed: Unknown size, total: 90.15 MiB) to /root/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading data: 100%|██████████| 33.3M/33.3M [00:01<00:00, 22.1MB/s]\n",
      "Downloading data files: 100%|██████████| 1/1 [00:02<00:00,  2.63s/it]\n",
      "Extracting data files: 100%|██████████| 1/1 [00:00<00:00, 497.96it/s]\n",
      "                                                                                       "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset parquet downloaded and prepared to /root/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                               \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens: 15360000\n"
     ]
    }
   ],
   "source": [
    "# Download the dataset\n",
    "# TODO iteratively grab dataset?\n",
    "cfg.max_length = 256\n",
    "cfg.model_batch_size = 4\n",
    "token_loader = setup_token_data(cfg, tokenizer, model)\n",
    "num_tokens = cfg.max_length*cfg.model_batch_size*len(token_loader)\n",
    "print(f\"Number of tokens: {num_tokens}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation size: 512\n"
     ]
    }
   ],
   "source": [
    "# Run 1 datapoint on model to get the activation size (cause don't want to deal w/ different naming schemes in config files)\n",
    "from baukit import Trace, TraceDict\n",
    "\n",
    "text = \"1\"\n",
    "tokens = tokenizer(text, return_tensors=\"pt\").input_ids.to(cfg.device)\n",
    "# Your activation name will be different. In the next cells, we will show you how to find it.\n",
    "with torch.no_grad():\n",
    "    with Trace(model, tensor_names[0]) as ret:\n",
    "        _ = model(tokens)\n",
    "        representation = ret.output\n",
    "        # check if instance tuple\n",
    "        if(isinstance(representation, tuple)):\n",
    "            representation = representation[0]\n",
    "        activation_size = representation.shape[-1]\n",
    "print(f\"Activation size: {activation_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torchtyping import TensorType\n",
    "\n",
    "\n",
    "class TiedSAE(nn.Module):\n",
    "    def __init__(self, activation_size, n_dict_components):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Parameter(torch.empty((n_dict_components, activation_size)))\n",
    "        nn.init.xavier_uniform_(self.encoder)\n",
    "        self.encoder_bias = nn.Parameter(torch.zeros((n_dict_components,)))\n",
    "\n",
    "    def get_learned_dict(self):\n",
    "        norms = torch.norm(self.encoder, 2, dim=-1)\n",
    "        return self.encoder / torch.clamp(norms, 1e-8)[:, None]\n",
    "\n",
    "    def encode(self, batch):\n",
    "        c = torch.einsum(\"nd,bd->bn\", self.encoder, batch)\n",
    "        c = c + self.encoder_bias\n",
    "        c = torch.clamp(c, min=0.0)\n",
    "        return c\n",
    "\n",
    "    def decode(self, code: TensorType[\"_batch_size\", \"_n_dict_components\"]) -> TensorType[\"_batch_size\", \"_activation_size\"]:\n",
    "        learned_dict = self.get_learned_dict()\n",
    "        x_hat = torch.einsum(\"nd,bn->bd\", learned_dict, code)\n",
    "        return x_hat\n",
    "\n",
    "    def forward(self, batch: TensorType[\"_batch_size\", \"_activation_size\"]) -> TensorType[\"_batch_size\", \"_activation_size\"]:\n",
    "        c = self.encode(batch)\n",
    "        x_hat = self.decode(c)\n",
    "        return x_hat, c\n",
    "\n",
    "    def n_dict_components(self):\n",
    "        return self.get_learned_dict().shape[0]\n",
    "\n",
    "n_dict_components = activation_size*cfg.ratio\n",
    "all_autoencoders = [TiedSAE(activation_size, n_dict_components).to(cfg.device) for _ in range(len(tensor_names))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizers = [torch.optim.Adam(autoencoder.parameters(), lr=cfg.lr) for autoencoder in all_autoencoders]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: Only works on tied SAE\n",
      "wandb_run_name: EleutherAI/pythia-70m-deduped_1102-145634_None\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.15.12 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.3"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/root/sparse_coding/wandb/run-20231102_145634-7ra4qrtk</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/sparse_coding/sparse%20coding/runs/7ra4qrtk' target=\"_blank\">EleutherAI/pythia-70m-deduped_1102-145634_None</a></strong> to <a href='https://wandb.ai/sparse_coding/sparse%20coding' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/sparse_coding/sparse%20coding' target=\"_blank\">https://wandb.ai/sparse_coding/sparse%20coding</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/sparse_coding/sparse%20coding/runs/7ra4qrtk' target=\"_blank\">https://wandb.ai/sparse_coding/sparse%20coding/runs/7ra4qrtk</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/sparse_coding/sparse%20coding/runs/7ra4qrtk?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x7f337f37c460>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"WARNING: Only works on tied SAE\")\n",
    "# Wandb setup\n",
    "secrets = json.load(open(\"secrets.json\"))\n",
    "wandb.login(key=secrets[\"wandb_key\"])\n",
    "start_time = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "wandb_run_name = f\"{cfg.model_name}_{start_time[4:]}_{cfg.sparsity}\"  # trim year\n",
    "print(f\"wandb_run_name: {wandb_run_name}\")\n",
    "wandb.init(project=\"sparse coding\", config=dict(cfg), name=wandb_run_name, entity=\"sparse_coding\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Sparsity: 82.8 | Dead Features: 36 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 1 | Sparsity: 146.3 | Dead Features: 37 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 0 | Sparsity: 94.0 | Dead Features: 0 | Total Loss: 0.05 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 204800 | Self Similarity: 0.98\n",
      "Layer 1 | Sparsity: 161.5 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 204800 | Self Similarity: 0.98\n",
      "Layer 0 | Sparsity: 81.5 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 409600 | Self Similarity: 0.97\n",
      "Layer 1 | Sparsity: 148.2 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 409600 | Self Similarity: 0.98\n",
      "Layer 0 | Sparsity: 85.2 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 614400 | Self Similarity: 0.97\n",
      "Layer 1 | Sparsity: 154.2 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 614400 | Self Similarity: 0.97\n",
      "Layer 0 | Sparsity: 79.8 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 819200 | Self Similarity: 0.97\n",
      "Layer 1 | Sparsity: 143.0 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 819200 | Self Similarity: 0.97\n",
      "Layer 0 | Sparsity: 88.8 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 1024000 | Self Similarity: 0.97\n",
      "Layer 1 | Sparsity: 155.8 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 1024000 | Self Similarity: 0.96\n",
      "Layer 0 | Sparsity: 88.9 | Dead Features: 0 | Total Loss: 0.05 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 1228800 | Self Similarity: 0.97\n",
      "Layer 1 | Sparsity: 155.7 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 1228800 | Self Similarity: 0.96\n",
      "Layer 0 | Sparsity: 87.3 | Dead Features: 0 | Total Loss: 0.05 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 1433600 | Self Similarity: 0.96\n",
      "Layer 1 | Sparsity: 153.1 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 1433600 | Self Similarity: 0.96\n",
      "Layer 0 | Sparsity: 93.9 | Dead Features: 0 | Total Loss: 0.05 | Reconstruction Loss: 0.03 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 1638400 | Self Similarity: 0.96\n",
      "Layer 1 | Sparsity: 154.0 | Dead Features: 0 | Total Loss: 0.08 | Reconstruction Loss: 0.04 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 1638400 | Self Similarity: 0.96\n",
      "Layer 0 | Sparsity: 93.6 | Dead Features: 0 | Total Loss: 0.05 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 1843200 | Self Similarity: 0.96\n",
      "Layer 1 | Sparsity: 161.6 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 1843200 | Self Similarity: 0.96\n",
      "Layer 0 | Sparsity: 85.7 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 2048000 | Self Similarity: 0.96\n",
      "Layer 1 | Sparsity: 154.4 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 2048000 | Self Similarity: 0.95\n",
      "Layer 0 | Sparsity: 80.6 | Dead Features: 0 | Total Loss: 0.04 | Reconstruction Loss: 0.02 | L1 Loss: 0.02 | l1_alpha: 8.00e-04 | Tokens: 2252800 | Self Similarity: 0.96\n",
      "Layer 1 | Sparsity: 145.8 | Dead Features: 0 | Total Loss: 0.07 | Reconstruction Loss: 0.03 | L1 Loss: 0.04 | l1_alpha: 8.00e-04 | Tokens: 2252800 | Self Similarity: 0.95\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/root/sparse_coding/alpha_easy_training_code.ipynb Cell 9\u001b[0m line \u001b[0;36m2\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=19'>20</a>\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad(): \u001b[39m# As long as not doing KL divergence, don't need gradients for model\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=20'>21</a>\u001b[0m     \u001b[39mwith\u001b[39;00m TraceDict(model, tensor_names) \u001b[39mas\u001b[39;00m ret:\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21'>22</a>\u001b[0m         _ \u001b[39m=\u001b[39m model(tokens)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=23'>24</a>\u001b[0m \u001b[39mfor\u001b[39;00m auto_ind \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39mlen\u001b[39m(tensor_names)):\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a>\u001b[0m     \u001b[39m# Index into correct autoencoder, optimizer, and tensor_name\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B154.20.254.95/root/sparse_coding/alpha_easy_training_code.ipynb#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m     autoencoder \u001b[39m=\u001b[39m all_autoencoders[auto_ind]\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py:673\u001b[0m, in \u001b[0;36mGPTNeoXForCausalLM.forward\u001b[0;34m(self, input_ids, attention_mask, position_ids, inputs_embeds, head_mask, past_key_values, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    632\u001b[0m \u001b[39m\u001b[39m\u001b[39mr\u001b[39m\u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m    633\u001b[0m \u001b[39mpast_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):\u001b[39;00m\n\u001b[1;32m    634\u001b[0m \u001b[39m    Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    669\u001b[0m \u001b[39m>>> prediction_logits = outputs.logits\u001b[39;00m\n\u001b[1;32m    670\u001b[0m \u001b[39m```\"\"\"\u001b[39;00m\n\u001b[1;32m    671\u001b[0m return_dict \u001b[39m=\u001b[39m return_dict \u001b[39mif\u001b[39;00m return_dict \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39muse_return_dict\n\u001b[0;32m--> 673\u001b[0m outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mgpt_neox(\n\u001b[1;32m    674\u001b[0m     input_ids,\n\u001b[1;32m    675\u001b[0m     attention_mask\u001b[39m=\u001b[39;49mattention_mask,\n\u001b[1;32m    676\u001b[0m     position_ids\u001b[39m=\u001b[39;49mposition_ids,\n\u001b[1;32m    677\u001b[0m     head_mask\u001b[39m=\u001b[39;49mhead_mask,\n\u001b[1;32m    678\u001b[0m     inputs_embeds\u001b[39m=\u001b[39;49minputs_embeds,\n\u001b[1;32m    679\u001b[0m     past_key_values\u001b[39m=\u001b[39;49mpast_key_values,\n\u001b[1;32m    680\u001b[0m     use_cache\u001b[39m=\u001b[39;49muse_cache,\n\u001b[1;32m    681\u001b[0m     output_attentions\u001b[39m=\u001b[39;49moutput_attentions,\n\u001b[1;32m    682\u001b[0m     output_hidden_states\u001b[39m=\u001b[39;49moutput_hidden_states,\n\u001b[1;32m    683\u001b[0m     return_dict\u001b[39m=\u001b[39;49mreturn_dict,\n\u001b[1;32m    684\u001b[0m )\n\u001b[1;32m    686\u001b[0m hidden_states \u001b[39m=\u001b[39m outputs[\u001b[39m0\u001b[39m]\n\u001b[1;32m    687\u001b[0m lm_logits \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39membed_out(hidden_states)\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py:564\u001b[0m, in \u001b[0;36mGPTNeoXModel.forward\u001b[0;34m(self, input_ids, attention_mask, position_ids, head_mask, inputs_embeds, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    556\u001b[0m     outputs \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mutils\u001b[39m.\u001b[39mcheckpoint\u001b[39m.\u001b[39mcheckpoint(\n\u001b[1;32m    557\u001b[0m         create_custom_forward(layer),\n\u001b[1;32m    558\u001b[0m         hidden_states,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    561\u001b[0m         head_mask[i],\n\u001b[1;32m    562\u001b[0m     )\n\u001b[1;32m    563\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 564\u001b[0m     outputs \u001b[39m=\u001b[39m layer(\n\u001b[1;32m    565\u001b[0m         hidden_states,\n\u001b[1;32m    566\u001b[0m         attention_mask\u001b[39m=\u001b[39;49mattention_mask,\n\u001b[1;32m    567\u001b[0m         position_ids\u001b[39m=\u001b[39;49mposition_ids,\n\u001b[1;32m    568\u001b[0m         head_mask\u001b[39m=\u001b[39;49mhead_mask[i],\n\u001b[1;32m    569\u001b[0m         layer_past\u001b[39m=\u001b[39;49mlayer_past,\n\u001b[1;32m    570\u001b[0m         use_cache\u001b[39m=\u001b[39;49muse_cache,\n\u001b[1;32m    571\u001b[0m         output_attentions\u001b[39m=\u001b[39;49moutput_attentions,\n\u001b[1;32m    572\u001b[0m     )\n\u001b[1;32m    573\u001b[0m hidden_states \u001b[39m=\u001b[39m outputs[\u001b[39m0\u001b[39m]\n\u001b[1;32m    574\u001b[0m \u001b[39mif\u001b[39;00m use_cache \u001b[39mis\u001b[39;00m \u001b[39mTrue\u001b[39;00m:\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py:331\u001b[0m, in \u001b[0;36mGPTNeoXLayer.forward\u001b[0;34m(self, hidden_states, attention_mask, position_ids, head_mask, use_cache, layer_past, output_attentions)\u001b[0m\n\u001b[1;32m    321\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\n\u001b[1;32m    322\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m    323\u001b[0m     hidden_states: Optional[torch\u001b[39m.\u001b[39mFloatTensor],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    329\u001b[0m     output_attentions: Optional[\u001b[39mbool\u001b[39m] \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m,\n\u001b[1;32m    330\u001b[0m ):\n\u001b[0;32m--> 331\u001b[0m     attention_layer_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mattention(\n\u001b[1;32m    332\u001b[0m         \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49minput_layernorm(hidden_states),\n\u001b[1;32m    333\u001b[0m         attention_mask\u001b[39m=\u001b[39;49mattention_mask,\n\u001b[1;32m    334\u001b[0m         position_ids\u001b[39m=\u001b[39;49mposition_ids,\n\u001b[1;32m    335\u001b[0m         layer_past\u001b[39m=\u001b[39;49mlayer_past,\n\u001b[1;32m    336\u001b[0m         head_mask\u001b[39m=\u001b[39;49mhead_mask,\n\u001b[1;32m    337\u001b[0m         use_cache\u001b[39m=\u001b[39;49muse_cache,\n\u001b[1;32m    338\u001b[0m         output_attentions\u001b[39m=\u001b[39;49moutput_attentions,\n\u001b[1;32m    339\u001b[0m     )\n\u001b[1;32m    340\u001b[0m     attn_output \u001b[39m=\u001b[39m attention_layer_outputs[\u001b[39m0\u001b[39m]  \u001b[39m# output_attn: attn_output, present, (attn_weights)\u001b[39;00m\n\u001b[1;32m    341\u001b[0m     outputs \u001b[39m=\u001b[39m attention_layer_outputs[\u001b[39m1\u001b[39m:]\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py:149\u001b[0m, in \u001b[0;36mGPTNeoXAttention.forward\u001b[0;34m(self, hidden_states, attention_mask, position_ids, head_mask, layer_past, use_cache, output_attentions)\u001b[0m\n\u001b[1;32m    147\u001b[0m \u001b[39mif\u001b[39;00m has_layer_past:\n\u001b[1;32m    148\u001b[0m     seq_len \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m layer_past[\u001b[39m0\u001b[39m]\u001b[39m.\u001b[39mshape[\u001b[39m-\u001b[39m\u001b[39m2\u001b[39m]\n\u001b[0;32m--> 149\u001b[0m cos, sin \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mrotary_emb(value, seq_len\u001b[39m=\u001b[39;49mseq_len)\n\u001b[1;32m    150\u001b[0m query, key \u001b[39m=\u001b[39m apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)\n\u001b[1;32m    151\u001b[0m query \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mcat((query, query_pass), dim\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1502\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/logan/lib/python3.10/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py:278\u001b[0m, in \u001b[0;36mRotaryEmbedding.forward\u001b[0;34m(self, x, seq_len)\u001b[0m\n\u001b[1;32m    276\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcos_cached \u001b[39m=\u001b[39m emb\u001b[39m.\u001b[39mcos()[\u001b[39mNone\u001b[39;00m, \u001b[39mNone\u001b[39;00m, :, :]\n\u001b[1;32m    277\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msin_cached \u001b[39m=\u001b[39m emb\u001b[39m.\u001b[39msin()[\u001b[39mNone\u001b[39;00m, \u001b[39mNone\u001b[39;00m, :, :]\n\u001b[0;32m--> 278\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mcos_cached[:seq_len, \u001b[39m.\u001b[39;49m\u001b[39m.\u001b[39;49m\u001b[39m.\u001b[39;49m]\u001b[39m.\u001b[39;49mto(x\u001b[39m.\u001b[39;49mdevice), \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msin_cached[:seq_len, \u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m]\u001b[39m.\u001b[39mto(x\u001b[39m.\u001b[39mdevice)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "# Make directory trained_models if it doesn't exist\n",
    "import os\n",
    "if not os.path.exists(\"trained_models\"):\n",
    "    os.makedirs(\"trained_models\")\n",
    "model_save_name = cfg.model_name.split(\"/\")[-1]\n",
    "\n",
    "num_batch = len(token_loader)\n",
    "log_space = np.logspace(0, np.log10(num_batch), 11)  # 11 to get 10 intervals\n",
    "save_batches = [int(x) for x in log_space[1:]]  # Skip the first (0th) interval\n",
    "\n",
    "dead_features = [torch.zeros(n_dict_components) for _ in range(len(tensor_names))]\n",
    "last_encoders = [autoencoder.encoder.clone().detach() for autoencoder in all_autoencoders]\n",
    "# max_num_tokens = 100000000\n",
    "# Freeze model parameters \n",
    "model.eval()\n",
    "model.requires_grad_(False)\n",
    "for i, batch in enumerate(token_loader):\n",
    "    tokens = batch[\"input_ids\"].to(cfg.device)\n",
    "    with torch.no_grad(): # As long as not doing KL divergence, don't need gradients for model\n",
    "        with TraceDict(model, tensor_names) as ret:\n",
    "            _ = model(tokens)\n",
    "\n",
    "    for auto_ind in range(len(tensor_names)):\n",
    "        # Index into correct autoencoder, optimizer, and tensor_name\n",
    "        autoencoder = all_autoencoders[auto_ind]\n",
    "        optimizer = optimizers[auto_ind]\n",
    "        tensor_name = tensor_names[auto_ind]\n",
    "        dead_feature = dead_features[auto_ind]\n",
    "        last_encoder = last_encoders[auto_ind]\n",
    "        l1_alpha = cfg.l1_alpha[auto_ind]\n",
    "\n",
    "        # Get intermediate layer activations\n",
    "        representation = ret[tensor_name].output\n",
    "        if(isinstance(representation, tuple)):\n",
    "            representation = representation[0]\n",
    "        layer_activations = rearrange(representation, \"b seq d_model -> (b seq) d_model\")\n",
    "        \n",
    "        # Run through autoencoder\n",
    "        c = autoencoder.encode(layer_activations)\n",
    "        x_hat = autoencoder.decode(c)\n",
    "\n",
    "        # Calculate loss\n",
    "        reconstruction_loss = (x_hat - layer_activations).pow(2).mean()\n",
    "        l1_loss = torch.norm(c, 1, dim=-1).mean()\n",
    "        total_loss = reconstruction_loss + cfg.l1_alpha*l1_loss\n",
    "\n",
    "        # Update dead features\n",
    "        dead_feature += c.sum(dim=0).cpu()\n",
    "        \n",
    "        # Log\n",
    "        if (i % 200 == 0): # Check here so first check is model w/o change\n",
    "            # self_similarity = torch.cosine_similarity(c, last_encoder, dim=-1).mean().cpu().item()\n",
    "            # Above is wrong, should be similarity between encoder and last encoder\n",
    "            self_similarity = torch.cosine_similarity(autoencoder.encoder, last_encoder, dim=-1).mean().cpu().item()\n",
    "            last_encoder = autoencoder.encoder.clone().detach()\n",
    "            num_tokens_so_far = i*cfg.max_length*cfg.model_batch_size\n",
    "            with torch.no_grad():\n",
    "                sparsity = (c != 0).float().mean(dim=0).sum().cpu().item()\n",
    "                # Count number of dead_features are zero\n",
    "                num_dead_features = (dead_feature == 0).sum().item()\n",
    "            print(f\"Layer {auto_ind} | Sparsity: {sparsity:.1f} | Dead Features: {num_dead_features} | Total Loss: {total_loss:.2f} | Reconstruction Loss: {reconstruction_loss:.2f} | L1 Loss: {cfg.l1_alpha*l1_loss:.2f} | l1_alpha: {cfg.l1_alpha:.2e} | Tokens: {num_tokens_so_far} | Self Similarity: {self_similarity:.2f}\")\n",
    "            wandb.log({f\"Layer {auto_ind}\": {\n",
    "                'Sparsity': sparsity,\n",
    "                'Dead Features': num_dead_features,\n",
    "                'Total Loss': total_loss.item(),\n",
    "                'Reconstruction Loss': reconstruction_loss.item(),\n",
    "                'L1 Loss': (cfg.l1_alpha*l1_loss).item(),\n",
    "                'l1_alpha': cfg.l1_alpha,\n",
    "                'Tokens': num_tokens_so_far,\n",
    "                'Self Similarity': self_similarity,\n",
    "                'step': i}\n",
    "            })\n",
    "            # wandb.log({f\"Layer_{auto_ind}\": {\n",
    "            #     f\"Sparsity\": sparsity,\n",
    "            #     f\"Dead Features\": num_dead_features,\n",
    "            #     f\"Total Loss\": total_loss.item(),\n",
    "            #     f\"Reconstruction Loss\": reconstruction_loss.item(),\n",
    "            #     f\"L1 Loss\": (cfg.l1_alpha * l1_loss).item(),\n",
    "            #     f\"l1_alpha\": cfg.l1_alpha,\n",
    "            #     f\"Tokens\": num_tokens_so_far,\n",
    "            #     f\"Self Similarity\": self_similarity\n",
    "            # }, step=i})\n",
    "\n",
    "\n",
    "            # wandb.log({f\"Layer_{auto_ind}/sparsity\": sparsity_value, \"step\": step})\n",
    "\n",
    "            dead_feature = torch.zeros(autoencoder.encoder.shape[0])\n",
    "        # if i in save_batches:\n",
    "        #     save_name = f\"{model_save_name}_sp{cfg.sparsity}_r{cfg.ratio}_{tensor_names[0]}_{i}\"  # trim year\n",
    "        #     torch.save(autoencoder, f\"trained_models/{save_name}.pt\")\n",
    "        #     print(f\"Saved model to trained_models/{save_name}\")\n",
    "            \n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # # Running sparsity check\n",
    "    # if(num_tokens_so_far > 5000000):\n",
    "    #     if(i % 200 == 0):\n",
    "    #         with torch.no_grad():\n",
    "    #             sparsity = (c != 0).float().mean(dim=0).sum().cpu().item()\n",
    "    #         if sparsity > target_upper_sparsity:\n",
    "    #             cfg.l1_alpha *= (1 + adjustment_factor)\n",
    "    #         elif sparsity < target_lower_sparsity:\n",
    "    #             cfg.l1_alpha *= (1 - adjustment_factor)\n",
    "            # print(f\"Sparsity: {sparsity:.1f} | l1_alpha: {cfg.l1_alpha:.2e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>Layer 0/Dead Features</td><td>█▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>Layer 0/L1 Loss</td><td>▄▂▅█▁▄▃▄▄▃▃</td></tr><tr><td>Layer 0/Reconstruction Loss</td><td>▃▄▃█▂▂▃▅▄▁▄</td></tr><tr><td>Layer 0/Self Similarity</td><td>█▅▄▃▃▃▂▂▂▁▁</td></tr><tr><td>Layer 0/Sparsity</td><td>▄▂▄█▁▃▃▃▄▁▃</td></tr><tr><td>Layer 0/Tokens</td><td>▁▂▂▃▄▅▅▆▇▇█</td></tr><tr><td>Layer 0/Total Loss</td><td>▃▃▃█▁▂▂▄▄▁▃</td></tr><tr><td>Layer 0/l1_alpha</td><td>▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>Layer 1/Dead Features</td><td>█▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>Layer 1/L1 Loss</td><td>▄▅▄█▃▆▂▄▃▁▂</td></tr><tr><td>Layer 1/Reconstruction Loss</td><td>▃▃▁█▁▁▂▄▄▁▅</td></tr><tr><td>Layer 1/Self Similarity</td><td>█▅▅▄▄▃▃▃▂▁▁</td></tr><tr><td>Layer 1/Sparsity</td><td>▅▃▄█▂▅▃▃▅▁▁</td></tr><tr><td>Layer 1/Tokens</td><td>▁▂▂▃▄▅▅▆▇▇█</td></tr><tr><td>Layer 1/Total Loss</td><td>▃▄▂█▂▃▁▄▃▁▃</td></tr><tr><td>Layer 1/l1_alpha</td><td>▁▁▁▁▁▁▁▁▁▁▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>Layer 0/Dead Features</td><td>0</td></tr><tr><td>Layer 0/L1 Loss</td><td>0.02047</td></tr><tr><td>Layer 0/Reconstruction Loss</td><td>0.02424</td></tr><tr><td>Layer 0/Self Similarity</td><td>0.95479</td></tr><tr><td>Layer 0/Sparsity</td><td>87.98438</td></tr><tr><td>Layer 0/Tokens</td><td>2048000</td></tr><tr><td>Layer 0/Total Loss</td><td>0.04471</td></tr><tr><td>Layer 0/l1_alpha</td><td>0.0008</td></tr><tr><td>Layer 1/Dead Features</td><td>0</td></tr><tr><td>Layer 1/L1 Loss</td><td>0.0375</td></tr><tr><td>Layer 1/Reconstruction Loss</td><td>0.03409</td></tr><tr><td>Layer 1/Self Similarity</td><td>0.9454</td></tr><tr><td>Layer 1/Sparsity</td><td>149.63281</td></tr><tr><td>Layer 1/Tokens</td><td>2048000</td></tr><tr><td>Layer 1/Total Loss</td><td>0.07159</td></tr><tr><td>Layer 1/l1_alpha</td><td>0.0008</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">EleutherAI/pythia-70m-deduped_1102-145514_None</strong> at: <a href='https://wandb.ai/sparse_coding/sparse%20coding/runs/vtoa6qos' target=\"_blank\">https://wandb.ai/sparse_coding/sparse%20coding/runs/vtoa6qos</a><br/>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231102_145514-vtoa6qos/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_name = f\"{model_save_name}_sp{cfg.sparsity}_r{cfg.ratio}_{tensor_names[0]}\"  # trim year\n",
    "\n",
    "\n",
    "# Save model\n",
    "torch.save(autoencoder, f\"trained_models/{save_name}.pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "logan",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}