alpha_easy_training_code_rm.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(28)]\n",
    "layers = 8\n",
    "cfg.layers=[i for i in range(20,28)]\n",
    "cfg.setting=\"residual\"\n",
    "# cfg.tensor_name=\"gpt_neox.layers.{layer}\"\n",
    "cfg.tensor_name=\"transformer.h.{layer}\"\n",
    "# linearly interpolate between 8e-4 & 1e-4\n",
    "# cfg.l1_alpha=torch.linspace(1e-4, 1e-3, 8).tolist()\n",
    "cfg.l1_alpha=torch.linspace(3e-3, 8e-3, 8).tolist()\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": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:34<00:00, 11.40s/it]\n",
      "Some weights of the model checkpoint at reciprocate/dahoas-gptj-rm-static were not used when initializing GPTJForCausalLM: ['score.weight']\n",
      "- This IS expected if you are initializing GPTJForCausalLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing GPTJForCausalLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of GPTJForCausalLM were not initialized from the model checkpoint at reciprocate/dahoas-gptj-rm-static and are newly initialized: ['lm_head.weight', 'lm_head.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "GPTJForCausalLM(\n",
       "  (transformer): GPTJModel(\n",
       "    (wte): Embedding(50400, 4096)\n",
       "    (drop): Dropout(p=0.0, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0-27): 28 x GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=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",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=50400, bias=True)\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": [
      "Found cached dataset json (/root/.cache/huggingface/datasets/Anthropic___json/Anthropic--hh-rlhf-a9fdd36e8b50b8fa/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n",
      "Found cached dataset parquet (/root/.cache/huggingface/datasets/Elriggs___parquet/Elriggs--openwebtext-100k-79076ecafee8a6d5/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n",
      "Loading cached processed dataset at /root/.cache/huggingface/datasets/Anthropic___json/Anthropic--hh-rlhf-a9fdd36e8b50b8fa/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-509f39b09b7b4bc3_*_of_00008.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens: 149.68M\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, concatenate_datasets\n",
    "from activation_dataset import chunk_and_tokenize\n",
    "\n",
    "# Load the datasets\n",
    "dataset1 = load_dataset(\"Anthropic/hh-rlhf\", split=\"train\")\n",
    "dataset2 = load_dataset(\"Elriggs/openwebtext-100k\", split=\"train\")\n",
    "\n",
    "# Rename 'chosen' column in dataset1 to 'text' to match dataset2\n",
    "dataset1 = dataset1.rename_column(\"chosen\", \"text\")\n",
    "\n",
    "# Concatenate the datasets\n",
    "dataset = concatenate_datasets([dataset1, dataset2])\n",
    "\n",
    "# Chunk & tokenize the dataset\n",
    "max_seq_length= 256\n",
    "dataset, _ = chunk_and_tokenize(dataset, tokenizer, max_length=max_seq_length)\n",
    "max_tokens = dataset.num_rows*max_seq_length\n",
    "print(f\"Number of tokens: {max_tokens/1e6:.2f}M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "batch_size = 16\n",
    "token_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activation size: 4096\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": 7,
   "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) 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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: Only works on tied SAE\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33melriggs\u001b[0m (\u001b[33msparse_coding\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\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": [
      "wandb_run_name: reciprocate/dahoas-gptj-rm-static_0108-210734_None\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.16.1 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-20240108_210734-8vhra2zs</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/8vhra2zs' target=\"_blank\">reciprocate/dahoas-gptj-rm-static_0108-210734_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/8vhra2zs' target=\"_blank\">https://wandb.ai/sparse_coding/sparse%20coding/runs/8vhra2zs</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/8vhra2zs?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x7f5c1c210e80>"
      ]
     },
     "execution_count": 9,
     "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/36544 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 326.910888671875\n",
      "Layer 0 | Sparsity: 8239.1 | Dead Features: 498 | Total Loss: 105.59 | Reconstruction Loss: 40.81 | L1 Loss: 64.78 | l1_alpha: 3.00e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 332.73095703125\n",
      "Layer 1 | Sparsity: 8248.5 | Dead Features: 528 | Total Loss: 123.58 | Reconstruction Loss: 42.21 | L1 Loss: 81.37 | l1_alpha: 3.71e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 343.52294921875\n",
      "Layer 2 | Sparsity: 8123.0 | Dead Features: 573 | Total Loss: 138.88 | Reconstruction Loss: 41.13 | L1 Loss: 97.75 | l1_alpha: 4.43e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 352.97528076171875\n",
      "Layer 3 | Sparsity: 8185.5 | Dead Features: 637 | Total Loss: 160.99 | Reconstruction Loss: 43.19 | L1 Loss: 117.80 | l1_alpha: 5.14e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 394.566650390625\n",
      "Layer 4 | Sparsity: 8156.7 | Dead Features: 764 | Total Loss: 198.22 | Reconstruction Loss: 47.51 | L1 Loss: 150.71 | l1_alpha: 5.86e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 437.537109375\n",
      "Layer 5 | Sparsity: 8260.6 | Dead Features: 956 | Total Loss: 241.14 | Reconstruction Loss: 49.79 | L1 Loss: 191.34 | l1_alpha: 6.57e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 551.3869018554688\n",
      "Layer 6 | Sparsity: 8301.6 | Dead Features: 1197 | Total Loss: 340.48 | Reconstruction Loss: 73.41 | L1 Loss: 267.06 | l1_alpha: 7.29e-03 | Tokens: 0 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 893.001708984375\n",
      "Layer 7 | Sparsity: 8223.9 | Dead Features: 438 | Total Loss: 642.88 | Reconstruction Loss: 176.86 | L1 Loss: 466.03 | l1_alpha: 8.00e-03 | Tokens: 0 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 100/36544 [17:20<106:32:16, 10.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 326.36920166015625\n",
      "Layer 0 | Sparsity: 39.7 | Dead Features: 20 | Total Loss: 5.12 | Reconstruction Loss: 3.36 | L1 Loss: 1.76 | l1_alpha: 3.00e-03 | Tokens: 409600 | Self Similarity: 0.91\n",
      "Layer 1 | Layer Activation Norm: 332.1792907714844\n",
      "Layer 1 | Sparsity: 21.9 | Dead Features: 26 | Total Loss: 5.49 | Reconstruction Loss: 3.44 | L1 Loss: 2.05 | l1_alpha: 3.71e-03 | Tokens: 409600 | Self Similarity: 0.91\n",
      "Layer 2 | Layer Activation Norm: 342.9295654296875\n",
      "Layer 2 | Sparsity: 15.3 | Dead Features: 37 | Total Loss: 6.63 | Reconstruction Loss: 4.45 | L1 Loss: 2.17 | l1_alpha: 4.43e-03 | Tokens: 409600 | Self Similarity: 0.91\n",
      "Layer 3 | Layer Activation Norm: 352.311279296875\n",
      "Layer 3 | Sparsity: 18.7 | Dead Features: 58 | Total Loss: 7.18 | Reconstruction Loss: 4.41 | L1 Loss: 2.77 | l1_alpha: 5.14e-03 | Tokens: 409600 | Self Similarity: 0.91\n",
      "Layer 4 | Layer Activation Norm: 394.54498291015625\n",
      "Layer 4 | Sparsity: 31.0 | Dead Features: 59 | Total Loss: 7.90 | Reconstruction Loss: 4.11 | L1 Loss: 3.79 | l1_alpha: 5.86e-03 | Tokens: 409600 | Self Similarity: 0.90\n",
      "Layer 5 | Layer Activation Norm: 437.4925537109375\n",
      "Layer 5 | Sparsity: 35.0 | Dead Features: 69 | Total Loss: 8.02 | Reconstruction Loss: 3.73 | L1 Loss: 4.28 | l1_alpha: 6.57e-03 | Tokens: 409600 | Self Similarity: 0.89\n",
      "Layer 6 | Layer Activation Norm: 582.8680419921875\n",
      "Layer 6 | Sparsity: 22.1 | Dead Features: 121 | Total Loss: 10.10 | Reconstruction Loss: 4.14 | L1 Loss: 5.96 | l1_alpha: 7.29e-03 | Tokens: 409600 | Self Similarity: 0.86\n",
      "Layer 7 | Layer Activation Norm: 939.5850219726562\n",
      "Layer 7 | Sparsity: 20.8 | Dead Features: 115 | Total Loss: 15.55 | Reconstruction Loss: 4.97 | L1 Loss: 10.58 | l1_alpha: 8.00e-03 | Tokens: 409600 | Self Similarity: 0.86\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 200/36544 [35:14<106:13:19, 10.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 321.9735107421875\n",
      "Layer 0 | Sparsity: 28.8 | Dead Features: 15 | Total Loss: 4.44 | Reconstruction Loss: 2.88 | L1 Loss: 1.56 | l1_alpha: 3.00e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 327.80340576171875\n",
      "Layer 1 | Sparsity: 18.6 | Dead Features: 20 | Total Loss: 4.71 | Reconstruction Loss: 2.92 | L1 Loss: 1.79 | l1_alpha: 3.71e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 338.6168518066406\n",
      "Layer 2 | Sparsity: 25.2 | Dead Features: 30 | Total Loss: 5.18 | Reconstruction Loss: 3.04 | L1 Loss: 2.14 | l1_alpha: 4.43e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 348.11932373046875\n",
      "Layer 3 | Sparsity: 25.0 | Dead Features: 43 | Total Loss: 5.63 | Reconstruction Loss: 3.11 | L1 Loss: 2.53 | l1_alpha: 5.14e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 392.7237548828125\n",
      "Layer 4 | Sparsity: 30.3 | Dead Features: 43 | Total Loss: 6.74 | Reconstruction Loss: 3.34 | L1 Loss: 3.40 | l1_alpha: 5.86e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 436.200927734375\n",
      "Layer 5 | Sparsity: 16.6 | Dead Features: 57 | Total Loss: 7.33 | Reconstruction Loss: 3.51 | L1 Loss: 3.82 | l1_alpha: 6.57e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 580.2840576171875\n",
      "Layer 6 | Sparsity: 18.1 | Dead Features: 101 | Total Loss: 8.97 | Reconstruction Loss: 3.57 | L1 Loss: 5.39 | l1_alpha: 7.29e-03 | Tokens: 819200 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 920.5639038085938\n",
      "Layer 7 | Sparsity: 21.0 | Dead Features: 105 | Total Loss: 13.97 | Reconstruction Loss: 4.42 | L1 Loss: 9.56 | l1_alpha: 8.00e-03 | Tokens: 819200 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 300/36544 [52:53<113:27:56, 11.27s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 339.27459716796875\n",
      "Layer 0 | Sparsity: 20.7 | Dead Features: 7 | Total Loss: 4.38 | Reconstruction Loss: 2.79 | L1 Loss: 1.59 | l1_alpha: 3.00e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 345.06854248046875\n",
      "Layer 1 | Sparsity: 15.4 | Dead Features: 10 | Total Loss: 4.68 | Reconstruction Loss: 2.85 | L1 Loss: 1.83 | l1_alpha: 3.71e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 355.8642883300781\n",
      "Layer 2 | Sparsity: 16.8 | Dead Features: 14 | Total Loss: 5.13 | Reconstruction Loss: 2.99 | L1 Loss: 2.14 | l1_alpha: 4.43e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 365.0838623046875\n",
      "Layer 3 | Sparsity: 16.0 | Dead Features: 26 | Total Loss: 5.56 | Reconstruction Loss: 3.08 | L1 Loss: 2.48 | l1_alpha: 5.14e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 405.1784973144531\n",
      "Layer 4 | Sparsity: 16.9 | Dead Features: 27 | Total Loss: 6.51 | Reconstruction Loss: 3.27 | L1 Loss: 3.23 | l1_alpha: 5.86e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 446.6968994140625\n",
      "Layer 5 | Sparsity: 13.2 | Dead Features: 48 | Total Loss: 7.06 | Reconstruction Loss: 3.36 | L1 Loss: 3.70 | l1_alpha: 6.57e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 567.908447265625\n",
      "Layer 6 | Sparsity: 16.9 | Dead Features: 81 | Total Loss: 8.84 | Reconstruction Loss: 3.72 | L1 Loss: 5.13 | l1_alpha: 7.29e-03 | Tokens: 1228800 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 918.5861206054688\n",
      "Layer 7 | Sparsity: 21.0 | Dead Features: 78 | Total Loss: 13.27 | Reconstruction Loss: 4.12 | L1 Loss: 9.16 | l1_alpha: 8.00e-03 | Tokens: 1228800 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 400/36544 [1:10:54<116:00:15, 11.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 330.57354736328125\n",
      "Layer 0 | Sparsity: 18.7 | Dead Features: 6 | Total Loss: 4.31 | Reconstruction Loss: 2.90 | L1 Loss: 1.41 | l1_alpha: 3.00e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 336.3862609863281\n",
      "Layer 1 | Sparsity: 16.1 | Dead Features: 9 | Total Loss: 4.57 | Reconstruction Loss: 2.93 | L1 Loss: 1.64 | l1_alpha: 3.71e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 347.2971496582031\n",
      "Layer 2 | Sparsity: 15.7 | Dead Features: 12 | Total Loss: 5.03 | Reconstruction Loss: 3.11 | L1 Loss: 1.92 | l1_alpha: 4.43e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 356.7022705078125\n",
      "Layer 3 | Sparsity: 15.1 | Dead Features: 18 | Total Loss: 5.41 | Reconstruction Loss: 3.22 | L1 Loss: 2.19 | l1_alpha: 5.14e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 402.2048645019531\n",
      "Layer 4 | Sparsity: 16.4 | Dead Features: 20 | Total Loss: 6.29 | Reconstruction Loss: 3.39 | L1 Loss: 2.90 | l1_alpha: 5.86e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 444.7969970703125\n",
      "Layer 5 | Sparsity: 13.1 | Dead Features: 40 | Total Loss: 6.94 | Reconstruction Loss: 3.50 | L1 Loss: 3.44 | l1_alpha: 6.57e-03 | Tokens: 1638400 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 584.3994140625\n",
      "Layer 6 | Sparsity: 16.7 | Dead Features: 49 | Total Loss: 9.00 | Reconstruction Loss: 3.68 | L1 Loss: 5.33 | l1_alpha: 7.29e-03 | Tokens: 1638400 | Self Similarity: 0.98\n",
      "Layer 7 | Layer Activation Norm: 915.4186401367188\n",
      "Layer 7 | Sparsity: 20.2 | Dead Features: 59 | Total Loss: 13.87 | Reconstruction Loss: 4.38 | L1 Loss: 9.49 | l1_alpha: 8.00e-03 | Tokens: 1638400 | Self Similarity: 0.99\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 500/36544 [1:28:31<105:11:05, 10.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 321.9416809082031\n",
      "Layer 0 | Sparsity: 17.5 | Dead Features: 6 | Total Loss: 4.09 | Reconstruction Loss: 2.84 | L1 Loss: 1.25 | l1_alpha: 3.00e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 327.6669006347656\n",
      "Layer 1 | Sparsity: 15.1 | Dead Features: 5 | Total Loss: 4.31 | Reconstruction Loss: 2.88 | L1 Loss: 1.43 | l1_alpha: 3.71e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 338.4409484863281\n",
      "Layer 2 | Sparsity: 13.6 | Dead Features: 8 | Total Loss: 4.72 | Reconstruction Loss: 3.01 | L1 Loss: 1.70 | l1_alpha: 4.43e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 347.66876220703125\n",
      "Layer 3 | Sparsity: 13.1 | Dead Features: 9 | Total Loss: 5.06 | Reconstruction Loss: 3.10 | L1 Loss: 1.96 | l1_alpha: 5.14e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 385.40283203125\n",
      "Layer 4 | Sparsity: 14.2 | Dead Features: 13 | Total Loss: 5.85 | Reconstruction Loss: 3.29 | L1 Loss: 2.56 | l1_alpha: 5.86e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 426.73675537109375\n",
      "Layer 5 | Sparsity: 11.3 | Dead Features: 34 | Total Loss: 6.46 | Reconstruction Loss: 3.40 | L1 Loss: 3.05 | l1_alpha: 6.57e-03 | Tokens: 2048000 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 528.95654296875\n",
      "Layer 6 | Sparsity: 17.2 | Dead Features: 19 | Total Loss: 8.35 | Reconstruction Loss: 3.74 | L1 Loss: 4.61 | l1_alpha: 7.29e-03 | Tokens: 2048000 | Self Similarity: 0.99\n",
      "Layer 7 | Layer Activation Norm: 857.8214111328125\n",
      "Layer 7 | Sparsity: 22.9 | Dead Features: 44 | Total Loss: 22.34 | Reconstruction Loss: 13.18 | L1 Loss: 9.17 | l1_alpha: 8.00e-03 | Tokens: 2048000 | Self Similarity: 0.99\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 600/36544 [1:46:06<104:13:45, 10.44s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 324.1142578125\n",
      "Layer 0 | Sparsity: 17.0 | Dead Features: 6 | Total Loss: 4.09 | Reconstruction Loss: 2.83 | L1 Loss: 1.25 | l1_alpha: 3.00e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 329.7431945800781\n",
      "Layer 1 | Sparsity: 14.2 | Dead Features: 4 | Total Loss: 4.32 | Reconstruction Loss: 2.88 | L1 Loss: 1.44 | l1_alpha: 3.71e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 340.6180419921875\n",
      "Layer 2 | Sparsity: 13.0 | Dead Features: 8 | Total Loss: 4.75 | Reconstruction Loss: 3.03 | L1 Loss: 1.71 | l1_alpha: 4.43e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 349.72674560546875\n",
      "Layer 3 | Sparsity: 12.7 | Dead Features: 9 | Total Loss: 5.08 | Reconstruction Loss: 3.11 | L1 Loss: 1.98 | l1_alpha: 5.14e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 390.20947265625\n",
      "Layer 4 | Sparsity: 13.1 | Dead Features: 12 | Total Loss: 5.87 | Reconstruction Loss: 3.28 | L1 Loss: 2.59 | l1_alpha: 5.86e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 431.3857116699219\n",
      "Layer 5 | Sparsity: 11.9 | Dead Features: 33 | Total Loss: 6.51 | Reconstruction Loss: 3.37 | L1 Loss: 3.14 | l1_alpha: 6.57e-03 | Tokens: 2457600 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 557.7584228515625\n",
      "Layer 6 | Sparsity: 15.6 | Dead Features: 16 | Total Loss: 8.54 | Reconstruction Loss: 3.64 | L1 Loss: 4.90 | l1_alpha: 7.29e-03 | Tokens: 2457600 | Self Similarity: 0.99\n",
      "Layer 7 | Layer Activation Norm: 897.211181640625\n",
      "Layer 7 | Sparsity: 18.2 | Dead Features: 44 | Total Loss: 12.94 | Reconstruction Loss: 4.03 | L1 Loss: 8.91 | l1_alpha: 8.00e-03 | Tokens: 2457600 | Self Similarity: 0.99\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 700/36544 [2:03:41<104:34:07, 10.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 321.3299560546875\n",
      "Layer 0 | Sparsity: 16.4 | Dead Features: 5 | Total Loss: 4.02 | Reconstruction Loss: 2.85 | L1 Loss: 1.17 | l1_alpha: 3.00e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 326.93133544921875\n",
      "Layer 1 | Sparsity: 13.4 | Dead Features: 3 | Total Loss: 4.25 | Reconstruction Loss: 2.87 | L1 Loss: 1.37 | l1_alpha: 3.71e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 337.79058837890625\n",
      "Layer 2 | Sparsity: 12.3 | Dead Features: 8 | Total Loss: 4.65 | Reconstruction Loss: 3.03 | L1 Loss: 1.62 | l1_alpha: 4.43e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 346.9292297363281\n",
      "Layer 3 | Sparsity: 12.0 | Dead Features: 6 | Total Loss: 5.00 | Reconstruction Loss: 3.12 | L1 Loss: 1.88 | l1_alpha: 5.14e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 385.236328125\n",
      "Layer 4 | Sparsity: 13.4 | Dead Features: 9 | Total Loss: 5.71 | Reconstruction Loss: 3.25 | L1 Loss: 2.46 | l1_alpha: 5.86e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 425.8707275390625\n",
      "Layer 5 | Sparsity: 12.1 | Dead Features: 32 | Total Loss: 6.27 | Reconstruction Loss: 3.28 | L1 Loss: 2.99 | l1_alpha: 6.57e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 520.6934814453125\n",
      "Layer 6 | Sparsity: 14.5 | Dead Features: 16 | Total Loss: 7.93 | Reconstruction Loss: 3.52 | L1 Loss: 4.41 | l1_alpha: 7.29e-03 | Tokens: 2867200 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 851.0560302734375\n",
      "Layer 7 | Sparsity: 17.9 | Dead Features: 43 | Total Loss: 12.09 | Reconstruction Loss: 3.87 | L1 Loss: 8.22 | l1_alpha: 8.00e-03 | Tokens: 2867200 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 800/36544 [2:21:15<103:30:40, 10.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 333.9566345214844\n",
      "Layer 0 | Sparsity: 16.4 | Dead Features: 5 | Total Loss: 3.97 | Reconstruction Loss: 2.72 | L1 Loss: 1.25 | l1_alpha: 3.00e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 339.752685546875\n",
      "Layer 1 | Sparsity: 13.9 | Dead Features: 3 | Total Loss: 4.24 | Reconstruction Loss: 2.75 | L1 Loss: 1.49 | l1_alpha: 3.71e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 350.50567626953125\n",
      "Layer 2 | Sparsity: 12.6 | Dead Features: 6 | Total Loss: 4.67 | Reconstruction Loss: 2.92 | L1 Loss: 1.74 | l1_alpha: 4.43e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 359.78424072265625\n",
      "Layer 3 | Sparsity: 12.5 | Dead Features: 6 | Total Loss: 5.03 | Reconstruction Loss: 2.99 | L1 Loss: 2.04 | l1_alpha: 5.14e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 403.9671630859375\n",
      "Layer 4 | Sparsity: 12.8 | Dead Features: 8 | Total Loss: 5.82 | Reconstruction Loss: 3.15 | L1 Loss: 2.68 | l1_alpha: 5.86e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 447.1547546386719\n",
      "Layer 5 | Sparsity: 11.9 | Dead Features: 31 | Total Loss: 6.65 | Reconstruction Loss: 3.35 | L1 Loss: 3.30 | l1_alpha: 6.57e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 586.7034912109375\n",
      "Layer 6 | Sparsity: 15.2 | Dead Features: 14 | Total Loss: 8.79 | Reconstruction Loss: 3.56 | L1 Loss: 5.23 | l1_alpha: 7.29e-03 | Tokens: 3276800 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 933.0843505859375\n",
      "Layer 7 | Sparsity: 18.3 | Dead Features: 43 | Total Loss: 13.26 | Reconstruction Loss: 3.96 | L1 Loss: 9.30 | l1_alpha: 8.00e-03 | Tokens: 3276800 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 900/36544 [2:38:41<104:59:00, 10.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 334.88555908203125\n",
      "Layer 0 | Sparsity: 16.2 | Dead Features: 5 | Total Loss: 3.98 | Reconstruction Loss: 2.74 | L1 Loss: 1.23 | l1_alpha: 3.00e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 340.72259521484375\n",
      "Layer 1 | Sparsity: 13.1 | Dead Features: 1 | Total Loss: 4.29 | Reconstruction Loss: 2.81 | L1 Loss: 1.47 | l1_alpha: 3.71e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 351.44189453125\n",
      "Layer 2 | Sparsity: 11.7 | Dead Features: 6 | Total Loss: 4.72 | Reconstruction Loss: 2.98 | L1 Loss: 1.74 | l1_alpha: 4.43e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 360.45684814453125\n",
      "Layer 3 | Sparsity: 12.0 | Dead Features: 5 | Total Loss: 5.09 | Reconstruction Loss: 3.05 | L1 Loss: 2.04 | l1_alpha: 5.14e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 401.34307861328125\n",
      "Layer 4 | Sparsity: 12.5 | Dead Features: 8 | Total Loss: 5.85 | Reconstruction Loss: 3.21 | L1 Loss: 2.64 | l1_alpha: 5.86e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 441.8974609375\n",
      "Layer 5 | Sparsity: 12.4 | Dead Features: 29 | Total Loss: 6.56 | Reconstruction Loss: 3.33 | L1 Loss: 3.23 | l1_alpha: 6.57e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 559.699462890625\n",
      "Layer 6 | Sparsity: 15.9 | Dead Features: 14 | Total Loss: 8.75 | Reconstruction Loss: 3.71 | L1 Loss: 5.04 | l1_alpha: 7.29e-03 | Tokens: 3686400 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 894.0987548828125\n",
      "Layer 7 | Sparsity: 19.4 | Dead Features: 43 | Total Loss: 12.98 | Reconstruction Loss: 3.96 | L1 Loss: 9.02 | l1_alpha: 8.00e-03 | Tokens: 3686400 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 1000/36544 [2:56:17<104:42:45, 10.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 323.3951110839844\n",
      "Layer 0 | Sparsity: 16.2 | Dead Features: 5 | Total Loss: 3.90 | Reconstruction Loss: 2.75 | L1 Loss: 1.15 | l1_alpha: 3.00e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 329.0091247558594\n",
      "Layer 1 | Sparsity: 12.9 | Dead Features: 1 | Total Loss: 4.20 | Reconstruction Loss: 2.82 | L1 Loss: 1.37 | l1_alpha: 3.71e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 339.8122253417969\n",
      "Layer 2 | Sparsity: 11.8 | Dead Features: 6 | Total Loss: 4.58 | Reconstruction Loss: 2.96 | L1 Loss: 1.63 | l1_alpha: 4.43e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 348.9909973144531\n",
      "Layer 3 | Sparsity: 11.7 | Dead Features: 3 | Total Loss: 4.92 | Reconstruction Loss: 3.01 | L1 Loss: 1.91 | l1_alpha: 5.14e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 388.858642578125\n",
      "Layer 4 | Sparsity: 11.6 | Dead Features: 8 | Total Loss: 5.65 | Reconstruction Loss: 3.18 | L1 Loss: 2.47 | l1_alpha: 5.86e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 430.93170166015625\n",
      "Layer 5 | Sparsity: 11.6 | Dead Features: 27 | Total Loss: 6.31 | Reconstruction Loss: 3.27 | L1 Loss: 3.04 | l1_alpha: 6.57e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 533.2321166992188\n",
      "Layer 6 | Sparsity: 14.3 | Dead Features: 14 | Total Loss: 8.04 | Reconstruction Loss: 3.49 | L1 Loss: 4.55 | l1_alpha: 7.29e-03 | Tokens: 4096000 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 866.7742309570312\n",
      "Layer 7 | Sparsity: 17.9 | Dead Features: 43 | Total Loss: 12.15 | Reconstruction Loss: 3.80 | L1 Loss: 8.35 | l1_alpha: 8.00e-03 | Tokens: 4096000 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 1100/36544 [3:13:56<102:57:11, 10.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 327.6083679199219\n",
      "Layer 0 | Sparsity: 16.5 | Dead Features: 5 | Total Loss: 3.88 | Reconstruction Loss: 2.70 | L1 Loss: 1.18 | l1_alpha: 3.00e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 333.3690185546875\n",
      "Layer 1 | Sparsity: 13.5 | Dead Features: 1 | Total Loss: 4.20 | Reconstruction Loss: 2.78 | L1 Loss: 1.42 | l1_alpha: 3.71e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 344.1620788574219\n",
      "Layer 2 | Sparsity: 12.3 | Dead Features: 6 | Total Loss: 4.62 | Reconstruction Loss: 2.92 | L1 Loss: 1.70 | l1_alpha: 4.43e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 353.4566955566406\n",
      "Layer 3 | Sparsity: 12.1 | Dead Features: 3 | Total Loss: 4.96 | Reconstruction Loss: 2.98 | L1 Loss: 1.98 | l1_alpha: 5.14e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 395.6383056640625\n",
      "Layer 4 | Sparsity: 12.1 | Dead Features: 7 | Total Loss: 5.71 | Reconstruction Loss: 3.12 | L1 Loss: 2.59 | l1_alpha: 5.86e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 437.71099853515625\n",
      "Layer 5 | Sparsity: 12.4 | Dead Features: 23 | Total Loss: 6.43 | Reconstruction Loss: 3.27 | L1 Loss: 3.17 | l1_alpha: 6.57e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 564.0108642578125\n",
      "Layer 6 | Sparsity: 15.2 | Dead Features: 14 | Total Loss: 8.46 | Reconstruction Loss: 3.43 | L1 Loss: 5.03 | l1_alpha: 7.29e-03 | Tokens: 4505600 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 899.4043579101562\n",
      "Layer 7 | Sparsity: 18.4 | Dead Features: 43 | Total Loss: 12.96 | Reconstruction Loss: 4.00 | L1 Loss: 8.96 | l1_alpha: 8.00e-03 | Tokens: 4505600 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 1200/36544 [3:31:33<103:11:02, 10.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 305.8229064941406\n",
      "Layer 0 | Sparsity: 15.4 | Dead Features: 5 | Total Loss: 3.72 | Reconstruction Loss: 2.67 | L1 Loss: 1.05 | l1_alpha: 3.00e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 311.4511413574219\n",
      "Layer 1 | Sparsity: 12.6 | Dead Features: 1 | Total Loss: 3.99 | Reconstruction Loss: 2.71 | L1 Loss: 1.28 | l1_alpha: 3.71e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 322.301513671875\n",
      "Layer 2 | Sparsity: 10.9 | Dead Features: 5 | Total Loss: 4.37 | Reconstruction Loss: 2.83 | L1 Loss: 1.54 | l1_alpha: 4.43e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 331.65765380859375\n",
      "Layer 3 | Sparsity: 11.3 | Dead Features: 2 | Total Loss: 4.71 | Reconstruction Loss: 2.89 | L1 Loss: 1.82 | l1_alpha: 5.14e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 372.32440185546875\n",
      "Layer 4 | Sparsity: 11.3 | Dead Features: 5 | Total Loss: 5.44 | Reconstruction Loss: 3.07 | L1 Loss: 2.37 | l1_alpha: 5.86e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 415.1114501953125\n",
      "Layer 5 | Sparsity: 12.3 | Dead Features: 22 | Total Loss: 6.18 | Reconstruction Loss: 3.22 | L1 Loss: 2.96 | l1_alpha: 6.57e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 551.3297119140625\n",
      "Layer 6 | Sparsity: 14.9 | Dead Features: 14 | Total Loss: 8.22 | Reconstruction Loss: 3.39 | L1 Loss: 4.83 | l1_alpha: 7.29e-03 | Tokens: 4915200 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 888.8153076171875\n",
      "Layer 7 | Sparsity: 17.9 | Dead Features: 43 | Total Loss: 12.72 | Reconstruction Loss: 3.95 | L1 Loss: 8.77 | l1_alpha: 8.00e-03 | Tokens: 4915200 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▎         | 1300/36544 [3:49:17<103:25:37, 10.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 332.5668640136719\n",
      "Layer 0 | Sparsity: 15.7 | Dead Features: 5 | Total Loss: 4.04 | Reconstruction Loss: 2.82 | L1 Loss: 1.22 | l1_alpha: 3.00e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 338.3523864746094\n",
      "Layer 1 | Sparsity: 12.9 | Dead Features: 1 | Total Loss: 4.45 | Reconstruction Loss: 2.97 | L1 Loss: 1.48 | l1_alpha: 3.71e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 349.0629577636719\n",
      "Layer 2 | Sparsity: 11.7 | Dead Features: 5 | Total Loss: 4.84 | Reconstruction Loss: 3.08 | L1 Loss: 1.76 | l1_alpha: 4.43e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 358.3428955078125\n",
      "Layer 3 | Sparsity: 11.8 | Dead Features: 2 | Total Loss: 5.25 | Reconstruction Loss: 3.19 | L1 Loss: 2.06 | l1_alpha: 5.14e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 402.5042724609375\n",
      "Layer 4 | Sparsity: 11.4 | Dead Features: 4 | Total Loss: 6.01 | Reconstruction Loss: 3.34 | L1 Loss: 2.67 | l1_alpha: 5.86e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 445.2330322265625\n",
      "Layer 5 | Sparsity: 12.4 | Dead Features: 21 | Total Loss: 6.65 | Reconstruction Loss: 3.32 | L1 Loss: 3.34 | l1_alpha: 6.57e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 591.3103637695312\n",
      "Layer 6 | Sparsity: 14.9 | Dead Features: 14 | Total Loss: 8.89 | Reconstruction Loss: 3.61 | L1 Loss: 5.28 | l1_alpha: 7.29e-03 | Tokens: 5324800 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 937.1417846679688\n",
      "Layer 7 | Sparsity: 18.6 | Dead Features: 43 | Total Loss: 13.37 | Reconstruction Loss: 4.01 | L1 Loss: 9.36 | l1_alpha: 8.00e-03 | Tokens: 5324800 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 1400/36544 [4:06:59<103:49:56, 10.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 322.82891845703125\n",
      "Layer 0 | Sparsity: 15.8 | Dead Features: 4 | Total Loss: 3.77 | Reconstruction Loss: 2.61 | L1 Loss: 1.16 | l1_alpha: 3.00e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 328.53912353515625\n",
      "Layer 1 | Sparsity: 13.0 | Dead Features: 1 | Total Loss: 4.09 | Reconstruction Loss: 2.69 | L1 Loss: 1.40 | l1_alpha: 3.71e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 339.5108337402344\n",
      "Layer 2 | Sparsity: 11.9 | Dead Features: 4 | Total Loss: 4.50 | Reconstruction Loss: 2.81 | L1 Loss: 1.69 | l1_alpha: 4.43e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 348.8993835449219\n",
      "Layer 3 | Sparsity: 12.0 | Dead Features: 2 | Total Loss: 4.85 | Reconstruction Loss: 2.89 | L1 Loss: 1.97 | l1_alpha: 5.14e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 393.57073974609375\n",
      "Layer 4 | Sparsity: 11.9 | Dead Features: 3 | Total Loss: 5.59 | Reconstruction Loss: 3.01 | L1 Loss: 2.58 | l1_alpha: 5.86e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 436.65234375\n",
      "Layer 5 | Sparsity: 12.5 | Dead Features: 20 | Total Loss: 6.34 | Reconstruction Loss: 3.12 | L1 Loss: 3.22 | l1_alpha: 6.57e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 581.2742919921875\n",
      "Layer 6 | Sparsity: 15.0 | Dead Features: 14 | Total Loss: 8.54 | Reconstruction Loss: 3.34 | L1 Loss: 5.21 | l1_alpha: 7.29e-03 | Tokens: 5734400 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 919.274658203125\n",
      "Layer 7 | Sparsity: 17.7 | Dead Features: 43 | Total Loss: 13.18 | Reconstruction Loss: 4.00 | L1 Loss: 9.18 | l1_alpha: 8.00e-03 | Tokens: 5734400 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 1500/36544 [4:24:39<102:49:43, 10.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 319.8027038574219\n",
      "Layer 0 | Sparsity: 16.6 | Dead Features: 4 | Total Loss: 3.80 | Reconstruction Loss: 2.64 | L1 Loss: 1.16 | l1_alpha: 3.00e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 325.448974609375\n",
      "Layer 1 | Sparsity: 13.0 | Dead Features: 0 | Total Loss: 4.14 | Reconstruction Loss: 2.75 | L1 Loss: 1.39 | l1_alpha: 3.71e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 336.36712646484375\n",
      "Layer 2 | Sparsity: 12.2 | Dead Features: 3 | Total Loss: 4.53 | Reconstruction Loss: 2.85 | L1 Loss: 1.68 | l1_alpha: 4.43e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 345.80596923828125\n",
      "Layer 3 | Sparsity: 12.3 | Dead Features: 2 | Total Loss: 4.90 | Reconstruction Loss: 2.92 | L1 Loss: 1.97 | l1_alpha: 5.14e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 391.4119873046875\n",
      "Layer 4 | Sparsity: 11.4 | Dead Features: 3 | Total Loss: 5.66 | Reconstruction Loss: 3.10 | L1 Loss: 2.56 | l1_alpha: 5.86e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 435.0118408203125\n",
      "Layer 5 | Sparsity: 12.1 | Dead Features: 20 | Total Loss: 6.43 | Reconstruction Loss: 3.22 | L1 Loss: 3.21 | l1_alpha: 6.57e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 590.8253784179688\n",
      "Layer 6 | Sparsity: 14.5 | Dead Features: 14 | Total Loss: 8.86 | Reconstruction Loss: 3.58 | L1 Loss: 5.27 | l1_alpha: 7.29e-03 | Tokens: 6144000 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 933.18212890625\n",
      "Layer 7 | Sparsity: 18.4 | Dead Features: 43 | Total Loss: 13.63 | Reconstruction Loss: 4.25 | L1 Loss: 9.38 | l1_alpha: 8.00e-03 | Tokens: 6144000 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 1600/36544 [4:42:14<101:49:00, 10.49s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 336.3548889160156\n",
      "Layer 0 | Sparsity: 15.1 | Dead Features: 4 | Total Loss: 3.94 | Reconstruction Loss: 2.71 | L1 Loss: 1.22 | l1_alpha: 3.00e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 342.10821533203125\n",
      "Layer 1 | Sparsity: 12.0 | Dead Features: 0 | Total Loss: 4.31 | Reconstruction Loss: 2.84 | L1 Loss: 1.48 | l1_alpha: 3.71e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 352.849853515625\n",
      "Layer 2 | Sparsity: 11.1 | Dead Features: 2 | Total Loss: 4.71 | Reconstruction Loss: 2.93 | L1 Loss: 1.77 | l1_alpha: 4.43e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 362.18304443359375\n",
      "Layer 3 | Sparsity: 10.9 | Dead Features: 2 | Total Loss: 5.11 | Reconstruction Loss: 3.04 | L1 Loss: 2.07 | l1_alpha: 5.14e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 405.446533203125\n",
      "Layer 4 | Sparsity: 10.7 | Dead Features: 3 | Total Loss: 5.87 | Reconstruction Loss: 3.20 | L1 Loss: 2.67 | l1_alpha: 5.86e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 448.04058837890625\n",
      "Layer 5 | Sparsity: 11.4 | Dead Features: 20 | Total Loss: 6.58 | Reconstruction Loss: 3.27 | L1 Loss: 3.31 | l1_alpha: 6.57e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 568.4537353515625\n",
      "Layer 6 | Sparsity: 13.5 | Dead Features: 14 | Total Loss: 8.46 | Reconstruction Loss: 3.49 | L1 Loss: 4.97 | l1_alpha: 7.29e-03 | Tokens: 6553600 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 920.5123291015625\n",
      "Layer 7 | Sparsity: 18.3 | Dead Features: 43 | Total Loss: 12.87 | Reconstruction Loss: 3.91 | L1 Loss: 8.96 | l1_alpha: 8.00e-03 | Tokens: 6553600 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 1700/36544 [4:59:47<102:03:53, 10.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0 | Layer Activation Norm: 328.76983642578125\n",
      "Layer 0 | Sparsity: 15.7 | Dead Features: 3 | Total Loss: 3.89 | Reconstruction Loss: 2.71 | L1 Loss: 1.18 | l1_alpha: 3.00e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 1 | Layer Activation Norm: 334.3890380859375\n",
      "Layer 1 | Sparsity: 12.3 | Dead Features: 0 | Total Loss: 4.20 | Reconstruction Loss: 2.79 | L1 Loss: 1.42 | l1_alpha: 3.71e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 2 | Layer Activation Norm: 345.209716796875\n",
      "Layer 2 | Sparsity: 11.0 | Dead Features: 2 | Total Loss: 4.62 | Reconstruction Loss: 2.92 | L1 Loss: 1.69 | l1_alpha: 4.43e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 3 | Layer Activation Norm: 354.57989501953125\n",
      "Layer 3 | Sparsity: 10.9 | Dead Features: 2 | Total Loss: 4.98 | Reconstruction Loss: 3.01 | L1 Loss: 1.97 | l1_alpha: 5.14e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 4 | Layer Activation Norm: 396.1776123046875\n",
      "Layer 4 | Sparsity: 10.9 | Dead Features: 3 | Total Loss: 5.69 | Reconstruction Loss: 3.13 | L1 Loss: 2.56 | l1_alpha: 5.86e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 5 | Layer Activation Norm: 438.84429931640625\n",
      "Layer 5 | Sparsity: 11.3 | Dead Features: 20 | Total Loss: 6.41 | Reconstruction Loss: 3.23 | L1 Loss: 3.18 | l1_alpha: 6.57e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 6 | Layer Activation Norm: 560.1085205078125\n",
      "Layer 6 | Sparsity: 14.0 | Dead Features: 14 | Total Loss: 8.37 | Reconstruction Loss: 3.45 | L1 Loss: 4.92 | l1_alpha: 7.29e-03 | Tokens: 6963200 | Self Similarity: 1.00\n",
      "Layer 7 | Layer Activation Norm: 902.2462768554688\n",
      "Layer 7 | Sparsity: 18.1 | Dead Features: 41 | Total Loss: 12.69 | Reconstruction Loss: 3.78 | L1 Loss: 8.91 | l1_alpha: 8.00e-03 | Tokens: 6963200 | Self Similarity: 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 1726/36544 [5:04:22<101:16:38, 10.47s/it]"
     ]
    }
   ],
   "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(tqdm(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].to(cfg.device)\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 + l1_alpha*l1_loss\n",
    "\n",
    "        # Update dead features\n",
    "        dead_feature += c.sum(dim=0).cpu()\n",
    "\n",
    "        \n",
    "        # Log\n",
    "        if (i % 50 == 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",
    "        \n",
    "            num_tokens_so_far = i*max_seq_length*batch_size\n",
    "            with torch.no_grad():\n",
    "                # print the norm of layer_activations\n",
    "                print(f\"Layer {auto_ind} | Layer Activation Norm: {torch.norm(layer_activations, 2, dim=-1).mean().cpu().item()}\")\n",
    "                self_similarity = torch.cosine_similarity(autoencoder.encoder, last_encoder.to(cfg.device), dim=-1).mean().cpu().item()\n",
    "                last_encoders[auto_ind] = autoencoder.encoder.clone().to(\"cpu\")\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: {l1_alpha*l1_loss:.2f} | l1_alpha: {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': (l1_alpha*l1_loss).item(),\n",
    "                'l1_alpha': l1_alpha,\n",
    "                'Tokens': num_tokens_so_far,\n",
    "                'Self Similarity': self_similarity,\n",
    "                'step': i}\n",
    "            })\n",
    "            dead_feature = torch.zeros(autoencoder.encoder.shape[0])\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Put SAE back on cpu\n",
    "        all_autoencoders[auto_ind] = autoencoder.to(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [
    "\n",
    "\n",
    "# Save every autoencoder\n",
    "for i, autoencoder in enumerate(all_autoencoders):\n",
    "    save_name = f\"{model_save_name}_r{cfg.ratio}_{tensor_names[i]}\"  # trim year\n",
    "    torch.save(autoencoder, f\"trained_models/{save_name}_{cfg.layers[i]}.pt\")\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
}