Last active
February 17, 2025 11:50
-
-
Save Butanium/805d5965af0c9a9ce903e330efb65118 to your computer and use it in GitHub Desktop.
crosscoder.ipynb
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| "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