Skip to content

Instantly share code, notes, and snippets.

@Butanium
Last active February 17, 2025 11:50
Show Gist options
  • Select an option

  • Save Butanium/805d5965af0c9a9ce903e330efb65118 to your computer and use it in GitHub Desktop.

Select an option

Save Butanium/805d5965af0c9a9ce903e330efb65118 to your computer and use it in GitHub Desktop.
crosscoder.ipynb
Display the source blob
Display the rendered blob
Raw
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"authorship_tag": "ABX9TyPSTzsSM4MXZ7zFG1SFBsh9",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/Butanium/805d5965af0c9a9ce903e330efb65118/crosscoder.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [
"# Open Source replication of crosscoders"
],
"metadata": {
"id": "N63JeCEoT_78"
}
},
{
"cell_type": "markdown",
"source": [
"## Setup"
],
"metadata": {
"id": "AjxA45DWUENU"
}
},
{
"cell_type": "code",
"source": [],
"metadata": {
"id": "xWHDUCpAvhPP",
"outputId": "c16c80a0-0364-4c18-f2aa-54f48bf8064a",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 290
}
},
"execution_count": 13,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Collecting bitsandbytes==0.44.1.dev0\n",
" Downloading https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_multi-backend-refactor/bitsandbytes-0.44.1.dev0-py3-none-manylinux_2_24_x86_64.whl (1.8 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m31.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hInstalling collected packages: bitsandbytes\n",
" Attempting uninstall: bitsandbytes\n",
" Found existing installation: bitsandbytes 0.45.2\n",
" Uninstalling bitsandbytes-0.45.2:\n",
" Successfully uninstalled bitsandbytes-0.45.2\n",
"Successfully installed bitsandbytes-0.44.1.dev0+9315692\n"
]
},
{
"output_type": "display_data",
"data": {
"application/vnd.colab-display-data+json": {
"pip_warning": {
"packages": [
"bitsandbytes"
]
},
"id": "578d82e53ffa4a4fbdffb187adb22712"
}
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"!pip install -q nnterp tiny-dashboard huggingface-hub datasets intel_extension_for_pytorch 'https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_multi-backend-refactor/bitsandbytes-0.44.1.dev0-py3-none-manylinux_2_24_x86_64.whl'"
],
"metadata": {
"id": "Wi7bvjIzCKhg",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "fed2e835-5060-4459-d0d0-4991a4aea257"
},
"execution_count": 11,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m105.4/105.4 MB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25h"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"## Architecture and test"
],
"metadata": {
"id": "ZxkpLbilUHjd"
}
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"id": "rnaERx7095-U"
},
"outputs": [],
"source": [
"import torch.nn as nn\n",
"import torch as th\n",
"from torch.nn.functional import relu\n",
"from warnings import warn\n",
"import einops\n",
"from huggingface_hub import PyTorchModelHubMixin\n",
"\n",
"\"\"\"\n",
"Code inspired by https://github.com/saprmarks/dictionary_learning\n",
"\"\"\"\n",
"\n",
"\n",
"class Encoder(nn.Module):\n",
" \"\"\"\n",
" A cross-coder encoder\n",
" \"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" activation_dim,\n",
" dict_size,\n",
" num_layers=None,\n",
" same_init_for_all_layers: bool = False,\n",
" norm_init_scale: float | None = None,\n",
" encoder_layers: list[int] | None = None,\n",
" ):\n",
" super().__init__()\n",
"\n",
" if encoder_layers is None:\n",
" if num_layers is None:\n",
" raise ValueError(\n",
" \"Either encoder_layers or num_layers must be specified\"\n",
" )\n",
" encoder_layers = list(range(num_layers))\n",
" else:\n",
" num_layers = len(encoder_layers)\n",
" self.encoder_layers = encoder_layers\n",
" self.activation_dim = activation_dim\n",
" self.dict_size = dict_size\n",
" self.num_layers = num_layers\n",
" if same_init_for_all_layers:\n",
" weight = nn.init.kaiming_uniform_(th.empty(activation_dim, dict_size))\n",
" weight = weight.repeat(num_layers, 1, 1)\n",
" else:\n",
" weight = nn.init.kaiming_uniform_(\n",
" th.empty(num_layers, activation_dim, dict_size)\n",
" )\n",
" if norm_init_scale is not None:\n",
" weight = weight / weight.norm(dim=1, keepdim=True) * norm_init_scale\n",
" self.weight = nn.Parameter(weight)\n",
" self.bias = nn.Parameter(th.zeros(dict_size))\n",
"\n",
" def forward(\n",
" self,\n",
" x: th.Tensor,\n",
" return_no_sum: bool = False,\n",
" select_features: list[int] | None = None,\n",
" ) -> th.Tensor: # (batch_size, activation_dim)\n",
" \"\"\"\n",
" Convert activations to features for each layer\n",
"\n",
" Args:\n",
" x: (batch_size, n_layers, activation_dim)\n",
" Returns:\n",
" f: (batch_size, dict_size)\n",
" \"\"\"\n",
" x = x[:, self.encoder_layers]\n",
" if select_features is not None:\n",
" w = self.weight[:, :, select_features]\n",
" bias = self.bias[select_features]\n",
" else:\n",
" w = self.weight\n",
" bias = self.bias\n",
" f = th.einsum(\"bld, ldf -> blf\", x, w)\n",
" return relu(f.sum(dim=1) + bias)\n",
"\n",
"\n",
"class CrossCoderDecoder(nn.Module):\n",
" \"\"\"\n",
" A crosscoder decoder\n",
" \"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" activation_dim,\n",
" dict_size,\n",
" num_layers,\n",
" same_init_for_all_layers: bool = True,\n",
" norm_init_scale: float | None = None,\n",
" init_with_weight: th.Tensor | None = None,\n",
" ):\n",
" super().__init__()\n",
" self.activation_dim = activation_dim\n",
" self.dict_size = dict_size\n",
" self.num_layers = num_layers\n",
" self.bias = nn.Parameter(th.zeros(num_layers, activation_dim))\n",
" if init_with_weight is not None:\n",
" self.weight = nn.Parameter(init_with_weight)\n",
" else:\n",
" if same_init_for_all_layers:\n",
" weight = nn.init.kaiming_uniform_(th.empty(dict_size, activation_dim))\n",
" weight = weight.repeat(num_layers, 1, 1)\n",
" else:\n",
" weight = nn.init.kaiming_uniform_(\n",
" th.empty(num_layers, dict_size, activation_dim)\n",
" )\n",
" if norm_init_scale is not None:\n",
" weight = weight / weight.norm(dim=2, keepdim=True) * norm_init_scale\n",
" self.weight = nn.Parameter(weight)\n",
"\n",
" def forward(\n",
" self,\n",
" f: th.Tensor,\n",
" select_features: list[int] | None = None,\n",
" add_bias: bool = True,\n",
" ) -> th.Tensor: # (batch_size, n_layers, activation_dim)\n",
" # f: (batch_size, n_layers, dict_size)\n",
" \"\"\"\n",
" Convert features to activations for each layer\n",
"\n",
" Args:\n",
" f: (batch_size, dict_size)\n",
" Returns:\n",
" x: (batch_size, n_layers, activation_dim)\n",
" \"\"\"\n",
" if select_features is not None:\n",
" w = self.weight[:, select_features]\n",
" else:\n",
" w = self.weight\n",
" x = th.einsum(\"bf, lfd -> bld\", f, w)\n",
" if add_bias:\n",
" x += self.bias\n",
" return x\n",
"\n",
"\n",
"class CrossCoder(PyTorchModelHubMixin, nn.Module):\n",
" \"\"\"\n",
" encoder: shape (num_layers, activation_dim, dict_size)\n",
" decoder: shape (num_layers, dict_size, activation_dim)\n",
" \"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" activation_dim,\n",
" dict_size,\n",
" num_layers,\n",
" same_init_for_all_layers=False,\n",
" norm_init_scale: float | None = None,\n",
" init_with_transpose=True,\n",
" encoder_layers: list[int] | None = None,\n",
" num_decoder_layers: int | None = None,\n",
" ):\n",
" \"\"\"\n",
" Args:\n",
" same_init_for_all_layers: if True, initialize all layers with the same vector\n",
" norm_init_scale: if not None, initialize the weights with a norm of this value\n",
" init_with_transpose: if True, initialize the decoder weights with the transpose of the encoder weights\n",
" encoder_layers: list of layers to use for the encoder. If None, num_layers must be specified.\n",
" num_decoder_layers: Number of decoder layers. If None, use num_layers.\n",
" \"\"\"\n",
" super().__init__()\n",
" if num_decoder_layers is None:\n",
" num_decoder_layers = num_layers\n",
"\n",
" self.activation_dim = activation_dim\n",
" self.dict_size = dict_size\n",
" self.num_layers = num_layers\n",
" self.encoder = Encoder(\n",
" activation_dim,\n",
" dict_size,\n",
" num_layers,\n",
" same_init_for_all_layers=same_init_for_all_layers,\n",
" norm_init_scale=norm_init_scale,\n",
" encoder_layers=encoder_layers,\n",
" )\n",
" if init_with_transpose:\n",
" decoder_weight = einops.rearrange(\n",
" self.encoder.weight.data.clone(),\n",
" \"num_layers activation_dim dict_size -> num_layers dict_size activation_dim\",\n",
" )\n",
" else:\n",
" decoder_weight = None\n",
" self.decoder = CrossCoderDecoder(\n",
" activation_dim,\n",
" dict_size,\n",
" num_decoder_layers,\n",
" same_init_for_all_layers=same_init_for_all_layers,\n",
" init_with_weight=decoder_weight,\n",
" norm_init_scale=norm_init_scale,\n",
" )\n",
"\n",
" def encode(\n",
" self, x: th.Tensor, **kwargs\n",
" ) -> th.Tensor: # (batch_size, n_layers, dict_size)\n",
" # x: (batch_size, n_layers, activation_dim)\n",
" return self.encoder(x, **kwargs)\n",
"\n",
" def get_activations(\n",
" self, x: th.Tensor, select_features: list[int] | None = None, **kwargs\n",
" ) -> th.Tensor:\n",
" f = self.encode(x, select_features=select_features, **kwargs)\n",
" if select_features is not None:\n",
" dw = self.decoder.weight[:, select_features]\n",
" else:\n",
" dw = self.decoder.weight\n",
" return f * dw.norm(dim=2).sum(dim=0, keepdim=True)\n",
"\n",
" def decode(\n",
" self, f: th.Tensor, **kwargs\n",
" ) -> th.Tensor: # (batch_size, n_layers, activation_dim)\n",
" # f: (batch_size, n_layers, dict_size)\n",
" return self.decoder(f, **kwargs)\n",
"\n",
" def forward(self, x: th.Tensor, output_features=False):\n",
" \"\"\"\n",
" Forward pass of the cross-coder.\n",
" x : activations to be encoded and decoded\n",
" output_features : if True, return the encoded features as well as the decoded x\n",
" \"\"\"\n",
" f = self.encode(x)\n",
" x_hat = self.decode(f)\n",
"\n",
" if output_features:\n",
" # Scale features by decoder column norms\n",
" f_scaled = f * self.decoder.weight.norm(dim=2).sum(\n",
" dim=0, keepdim=True\n",
" ) # Also sum across layers for the loss\n",
" return x_hat, f_scaled\n",
" else:\n",
" return x_hat\n",
"\n",
" @classmethod\n",
" def from_pretrained(\n",
" cls,\n",
" path: str,\n",
" dtype: th.dtype = th.float32,\n",
" device: th.device | None = None,\n",
" from_hub: bool = False,\n",
" **kwargs,\n",
" ):\n",
" \"\"\"\n",
" Load a pretrained cross-coder from a file.\n",
" \"\"\"\n",
" if from_hub:\n",
" return super().from_pretrained(path, device=device, dtype=dtype, **kwargs)\n",
"\n",
" state_dict = th.load(path, map_location=\"cpu\", weights_only=True)\n",
" if \"encoder.weight\" not in state_dict:\n",
" warn(\n",
" \"Cross-coder state dict was saved while torch.compiled was enabled. Fixing...\"\n",
" )\n",
" state_dict = {k.split(\"_orig_mod.\")[1]: v for k, v in state_dict.items()}\n",
" num_layers, activation_dim, dict_size = state_dict[\"encoder.weight\"].shape\n",
" cross_coder = cls(activation_dim, dict_size, num_layers)\n",
" cross_coder.load_state_dict(state_dict)\n",
"\n",
" if device is not None:\n",
" cross_coder = cross_coder.to(device)\n",
" return cross_coder.to(dtype=dtype)\n",
"\n",
" def resample_neurons(self, deads, activations):\n",
" # https://transformer-circuits.pub/2023/monosemantic-features/index.html#appendix-autoencoder-resampling\n",
" # compute loss for each activation\n",
" # impl from https://github.com/saprmarks/dictionary_learning\n",
" losses = (\n",
" (activations - self.forward(activations)).norm(dim=-1).mean(dim=-1).square()\n",
" )\n",
"\n",
" # sample input to create encoder/decoder weights from\n",
" n_resample = min([deads.sum(), losses.shape[0]])\n",
" print(\"Resampling\", n_resample, \"neurons\")\n",
" indices = th.multinomial(losses, num_samples=n_resample, replacement=False)\n",
" sampled_vecs = activations[indices] # (n_resample, num_layers, activation_dim)\n",
"\n",
" # get norm of the living neurons\n",
" # encoder.weight: (num_layers, activation_dim, dict_size)\n",
" # decoder.weight: (num_layers, dict_size, activation_dim)\n",
" alive_norm = self.encoder.weight[:, :, ~deads].norm(dim=-2)\n",
" alive_norm = alive_norm.mean(dim=-1) # (num_layers)\n",
" # convert to (num_layers, 1, 1)\n",
" alive_norm = einops.repeat(alive_norm, \"num_layers -> num_layers 1 1\")\n",
"\n",
" # resample first n_resample dead neurons\n",
" deads[deads.nonzero()[n_resample:]] = False\n",
" self.encoder.weight[:, :, deads] = (\n",
" sampled_vecs.permute(1, 2, 0) * alive_norm * 0.05\n",
" )\n",
" sampled_vecs = sampled_vecs.permute(1, 0, 2)\n",
" self.decoder.weight[:, deads, :] = th.nn.functional.normalize(\n",
" sampled_vecs, dim=-1\n",
" )\n",
" self.encoder.bias[deads] = 0.0"
]
},
{
"cell_type": "code",
"source": [
"# prompt: add tests of shapes\n",
"\n",
"import torch\n",
"import unittest\n",
"\n",
"class TestShapes(unittest.TestCase):\n",
"\n",
" def test_encoder_shape(self):\n",
" activation_dim = 64\n",
" dict_size = 128\n",
" num_layers = 4\n",
" encoder = Encoder(activation_dim, dict_size, num_layers)\n",
" x = torch.randn(10, num_layers, activation_dim) # Example input\n",
" f = encoder(x)\n",
" self.assertEqual(f.shape, (10, dict_size))\n",
"\n",
" def test_decoder_shape(self):\n",
" activation_dim = 64\n",
" dict_size = 128\n",
" num_layers = 4\n",
" decoder = CrossCoderDecoder(activation_dim, dict_size, num_layers)\n",
" f = torch.randn(10, dict_size) # Example input\n",
" x = decoder(f)\n",
" self.assertEqual(x.shape, (10, num_layers, activation_dim))\n",
"\n",
" def test_crosscoder_shape(self):\n",
" activation_dim = 64\n",
" dict_size = 128\n",
" num_layers = 4\n",
" crosscoder = CrossCoder(activation_dim, dict_size, num_layers)\n",
" x = torch.randn(10, num_layers, activation_dim) # Example input\n",
" x_hat = crosscoder(x)\n",
" self.assertEqual(x_hat.shape, (10, num_layers, activation_dim))\n",
"\n",
" def test_crosscoder_output_features_shape(self):\n",
" activation_dim = 64\n",
" dict_size = 128\n",
" num_layers = 4\n",
" crosscoder = CrossCoder(activation_dim, dict_size, num_layers)\n",
" x = torch.randn(10, num_layers, activation_dim)\n",
" x_hat, f_scaled = crosscoder(x, output_features=True)\n",
" self.assertEqual(x_hat.shape, (10, num_layers, activation_dim))\n",
" self.assertEqual(f_scaled.shape, (10, 128))\n",
"\n",
" def test_get_activations_shape(self):\n",
" activation_dim = 64\n",
" dict_size = 128\n",
" num_layers = 4\n",
" crosscoder = CrossCoder(activation_dim, dict_size, num_layers)\n",
" x = torch.randn(10, num_layers, activation_dim)\n",
" activations = crosscoder.get_activations(x)\n",
" self.assertEqual(activations.shape, (10, 128))\n",
"\n",
"\n",
"if __name__ == '__main__':\n",
" unittest.main(argv=['first-arg-is-ignored'], exit=False)\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "2VS-2YdiTo98",
"outputId": "8dfb9530-1324-4654-9724-b5d0bde669a6"
},
"execution_count": 8,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
".....\n",
"----------------------------------------------------------------------\n",
"Ran 5 tests in 0.017s\n",
"\n",
"OK\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"th.cuda.is_available()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "P6ew8ffOpQfd",
"outputId": "c17e6893-143f-4f94-9275-c60ba55ebd65"
},
"execution_count": 9,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"False"
]
},
"metadata": {},
"execution_count": 9
}
]
},
{
"cell_type": "code",
"source": [
"from nnsight import LanguageModel\n",
"import torch as th\n",
"dtype = th.bfloat16\n",
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"import torch\n",
"from transformers import BitsAndBytesConfig\n",
"\n",
"# Configure 4-bit quantization\n",
"quantization_config = BitsAndBytesConfig(\n",
" load_in_4bit=True,\n",
" bnb_4bit_compute_dtype=torch.float16,\n",
" bnb_4bit_quant_type=\"nf4\", # normalized float 4\n",
" bnb_4bit_use_double_quant=True\n",
")\n",
"\n",
"base_model = LanguageModel(\"Qwen/Qwen2.5-1.5B\", dispatch=True, torch_dtype=dtype, device_map=\"auto\", quantization_config=quantization_config)\n",
"chat_model = LanguageModel(\"Qwen/Qwen2.5-1.5B-Instruct\", dispatch=True, torch_dtype=dtype, device_map=\"auto\", quantization_config=quantization_config)\n"
],
"metadata": {
"id": "mSkUFTmOluLl",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 428
},
"outputId": "dc78c8d2-b9ba-43a2-d4c7-3a066700196d"
},
"execution_count": 10,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"CUDA is required but not available for bitsandbytes. Please consider installing the multi-platform enabled version of bitsandbytes, which is currently a work in progress. Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend\n"
]
},
{
"output_type": "error",
"ename": "RuntimeError",
"evalue": "CUDA is required but not available for bitsandbytes. Please consider installing the multi-platform enabled version of bitsandbytes, which is currently a work in progress. Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m<ipython-input-10-e7d8ab62cf34>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m 14\u001b[0m )\n\u001b[1;32m 15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m \u001b[0mbase_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLanguageModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Qwen/Qwen2.5-1.5B\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdispatch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch_dtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_map\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"auto\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquantization_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mquantization_config\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 17\u001b[0m \u001b[0mchat_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLanguageModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Qwen/Qwen2.5-1.5B-Instruct\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdispatch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch_dtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice_map\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"auto\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquantization_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mquantization_config\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/modeling/language.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, config, tokenizer, automodel, *args, **kwargs)\u001b[0m\n\u001b[1;32m 101\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrepo_id\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 102\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 105\u001b[0m self.generator: Envoy[InterventionProxyType, InterventionNodeType] = (\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/modeling/mixins/meta.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, dispatch, meta_buffers, rename, *args, **kwargs)\u001b[0m\n\u001b[1;32m 20\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatched\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 24\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/modeling/mixins/loadable.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, rename, *args, **kwargs)\u001b[0m\n\u001b[1;32m 12\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_load\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 16\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/modeling/language.py\u001b[0m in \u001b[0;36m_load\u001b[0;34m(self, repo_id, tokenizer_kwargs, patch_llama_scan, **kwargs)\u001b[0m\n\u001b[1;32m 173\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrope_scaling\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"rope_type\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"llama3\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 174\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 175\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautomodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_pretrained\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrepo_id\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 176\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 177\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/models/auto/auto_factory.py\u001b[0m in \u001b[0;36mfrom_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m 562\u001b[0m \u001b[0;32melif\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model_mapping\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 563\u001b[0m \u001b[0mmodel_class\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_get_model_class\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model_mapping\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 564\u001b[0;31m return model_class.from_pretrained(\n\u001b[0m\u001b[1;32m 565\u001b[0m \u001b[0mpretrained_model_name_or_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mmodel_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mhub_kwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 566\u001b[0m )\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/modeling_utils.py\u001b[0m in \u001b[0;36mfrom_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, weights_only, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m 3618\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3619\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mhf_quantizer\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3620\u001b[0;31m hf_quantizer.validate_environment(\n\u001b[0m\u001b[1;32m 3621\u001b[0m \u001b[0mtorch_dtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch_dtype\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3622\u001b[0m \u001b[0mfrom_tf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfrom_tf\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/quantizers/quantizer_bnb_4bit.py\u001b[0m in \u001b[0;36mvalidate_environment\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 81\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 82\u001b[0m \u001b[0mbnb_multibackend_is_enabled\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mis_bitsandbytes_multi_backend_available\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 83\u001b[0;31m \u001b[0mvalidate_bnb_backend_availability\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraise_exception\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 84\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 85\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"from_tf\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"from_flax\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/integrations/bitsandbytes.py\u001b[0m in \u001b[0;36mvalidate_bnb_backend_availability\u001b[0;34m(raise_exception)\u001b[0m\n\u001b[1;32m 557\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mis_bitsandbytes_multi_backend_available\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 558\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0m_validate_bnb_multi_backend_availability\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraise_exception\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 559\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0m_validate_bnb_cuda_backend_availability\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraise_exception\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/transformers/integrations/bitsandbytes.py\u001b[0m in \u001b[0;36m_validate_bnb_cuda_backend_availability\u001b[0;34m(raise_exception)\u001b[0m\n\u001b[1;32m 535\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mraise_exception\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 536\u001b[0m \u001b[0mlogger\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0merror\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlog_msg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 537\u001b[0;31m \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlog_msg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 538\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 539\u001b[0m \u001b[0mlogger\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwarning\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlog_msg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mRuntimeError\u001b[0m: CUDA is required but not available for bitsandbytes. Please consider installing the multi-platform enabled version of bitsandbytes, which is currently a work in progress. Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend"
]
}
]
},
{
"cell_type": "code",
"source": [
"hsize = chat_model._model.config.hidden_size\n",
"num_layers = chat_model._model.config.num_hidden_layers"
],
"metadata": {
"id": "psufXsdlsSCN"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# Interesting related models?\n",
"# math_model = LanguageModel(\"Qwen/Qwen2.5-Math-1.5B\", dispatch=True, torch_dtype=dtype)\n",
"# reasoning_model = LanguageModel(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\", dispatch=True, torch_dtype=dtype)"
],
"metadata": {
"id": "E54ssOk4meQL"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# Check for differences in tokenizer vocab and special tokens\n",
"def compare_tokenizers(tokenizer1, tokenizer2):\n",
" \"\"\"\n",
" Compares two tokenizers' vocabularies and special tokens.\n",
" \"\"\"\n",
" vocab_diff = set(tokenizer1.vocab.keys()) ^ set(tokenizer2.vocab.keys())\n",
" special_tokens_diff = set(tokenizer1.all_special_tokens) ^ set(tokenizer2.all_special_tokens)\n",
"\n",
" print(f\"Vocabulary differences: {vocab_diff}\")\n",
" print(f\"Special token differences: {special_tokens_diff}\")\n",
"\n",
" return vocab_diff, special_tokens_diff\n",
"\n",
"\n",
"# print(\"Comparing base_model and math_model:\")\n",
"# compare_tokenizers(base_model.tokenizer, math_model.tokenizer)\n",
"\n",
"# print(\"\\nComparing base_model and reasoning_model:\")\n",
"# compare_tokenizers(base_model.tokenizer, reasoning_model.tokenizer)\n",
"\n",
"print(\"\\nComparing base_model and chat_model:\")\n",
"compare_tokenizers(base_model.tokenizer, chat_model.tokenizer)\n"
],
"metadata": {
"id": "hk8dnWiRmJYz"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## Training (todo)"
],
"metadata": {
"id": "lNnHKQgdUKRm"
}
},
{
"cell_type": "markdown",
"source": [
"## Inference"
],
"metadata": {
"id": "b04Gu2uBlLlo"
}
},
{
"cell_type": "code",
"source": [
"from nnterp.nnsight_utils import get_layer, get_layer_output\n",
"def get_activations(prompts, base_model, chat_model, layer=14):\n",
" toks = chat_model.tokenizer.apply_chat_template(prompts, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024, return_dict=True)\n",
" attn_mask = toks.attention_mask.bool()\n",
" with chat_model.trace(toks):\n",
" chat_out = get_layer_output(chat_model, layer)[attn_mask].save()\n",
" get_layer(chat_model, layer).output.stop()\n",
" with base_model.trace(toks):\n",
" base_out = get_layer_output(base_model, layer)[attn_mask].save()\n",
" get_layer(base_model, layer).output.stop()\n",
" # activations: (batch_size, seq_len, d), with mask: (num_acts, d)\n",
" # return (num_acts,2, d)\n",
" return th.cat([base_out.unsqueeze(1), chat_out.unsqueeze(1)], dim=1)\n"
],
"metadata": {
"id": "I-euIOlMlP02"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"conversation = [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}, {\"role\": \"assistant\", \"content\": \"The capital of France is Paris.\"}, {\"role\": \"user\", \"content\": \"Who painted the Mona Lisa?\"}]\n",
"print(chat_model.tokenizer.apply_chat_template(conversation, tokenize=False))\n",
"import unittest\n",
"import torch as th\n",
"\n",
"class TestGetActivationsShape(unittest.TestCase):\n",
"\n",
" def test_get_activations_shape(self):\n",
" # Example prompts (replace with your actual prompts)\n",
" prompts = [conversation, conversation + [{\"role\": \"assistant\", \"content\": \"Mona Lisa was painted by Leonardo da Vinci.\"}]]\n",
"\n",
" toks = chat_model.tokenizer.apply_chat_template(prompts, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024, return_dict=True)\n",
" activations = get_activations(prompts, base_model, chat_model)\n",
"\n",
" # Check if the shape of the returned tensor is correct\n",
" self.assertEqual(activations.shape[0], sum(toks.attention_mask.sum(axis=1)))\n",
" self.assertEqual(activations.shape[1], 2) # 2 models (base and chat)\n",
" self.assertEqual(activations.shape[2], hsize) # hsize\n",
"\n",
"\n",
"if __name__ == '__main__':\n",
" unittest.main(argv=['first-arg-is-ignored'], exit=False)\n"
],
"metadata": {
"id": "dCEmlM7stRj1"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"from datasets import load_dataset\n",
"from torch.utils.data import DataLoader\n",
"\n",
"# Load the dataset (as you have in your code)\n",
"def collate_fn(batch):\n",
" return batch\n",
"dataset = load_dataset(\"HuggingFaceH4/ultrachat_200k\", streaming=True, split=\"test_sft\")\n",
"dict_size = 16_000\n",
"crosscoder = CrossCoder(hsize, dict_size, 2).to(\"cuda\" if th.cuda.is_available() else \"cpu\")\n",
"batch_size = 8\n",
"num_batches = 5\n",
"losses = []\n",
"dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)\n",
"with th.no_grad():\n",
" for batch in dataloader:\n",
" acts = get_activations(batch, base_model, chat_model).float()\n",
" reconstruction = crosscoder(acts)\n",
" loss = (reconstruction - acts).norm(dim=-1).mean()\n",
" losses.append(loss.item())\n",
" if len(losses) >= num_batches:\n",
" break\n",
"print(sum(losses) / len(losses))\n",
"\n"
],
"metadata": {
"id": "xBSDKQkRrvMS"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"from tiny_dashboard.dashboard_implementations import CrosscoderOnlineFeatureDashboard\n",
"\n",
"dashboard = CrosscoderOnlineFeatureDashboard(base_model, chat_model, crosscoder, 14)\n",
"dashboard.display()"
],
"metadata": {
"id": "hg9OVGRdnm6q"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [],
"metadata": {
"id": "14kBaEorpvjp"
},
"execution_count": null,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment