{ "cells": [ { "cell_type": "code", "execution_count": 45, "metadata": {}, "outputs": [], "source": [ "%reload_ext autoreload\n", "%autoreload 2" ] }, { "cell_type": "code", "execution_count": 46, "metadata": {}, "outputs": [], "source": [ "from datasets import load_dataset\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "from activation_store.collect import activation_store\n", "\n", "import torch" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load model" ] }, { "cell_type": "code", "execution_count": 47, "metadata": {}, "outputs": [], "source": [ "model_name = \"Qwen/Qwen2.5-0.5B-Instruct\"\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " model_name,\n", " torch_dtype=torch.bfloat16,\n", " device_map=\"auto\",\n", " attn_implementation=\"eager\", # flex_attention flash_attention_2 sdpa eager\n", ")\n", "tokenizer = AutoTokenizer.from_pretrained(model_name)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load data and tokenize" ] }, { "cell_type": "code", "execution_count": 48, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Dataset({\n", " features: ['attention_mask', 'input_ids'],\n", " num_rows: 20\n", "})" ] }, "execution_count": 48, "metadata": {}, "output_type": "execute_result" } ], "source": [ "N = 20\n", "max_length = 256\n", "\n", "imdb = load_dataset('wassname/imdb_dpo', split=f'test[:{N}]', keep_in_memory=False)\n", "\n", "\n", "def proc(row):\n", " messages = [\n", " {\"role\":\"user\", \"content\": row['prompt'] },\n", " {\"role\":\"assistant\", \"content\": row['chosen'] }\n", " ]\n", " return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_dict=True, max_length=max_length)\n", "\n", "ds2 = imdb.map(proc).with_format(\"torch\")\n", "new_cols = set(ds2.column_names) - set(imdb.column_names)\n", "ds2 = ds2.select_columns(new_cols)\n", "ds2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Data loader" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "from torch.utils.data import DataLoader\n", "def collate_fn(examples):\n", " # Pad the batch to max length within this batch\n", " return tokenizer.pad(\n", " examples,\n", " padding=True,\n", " return_tensors=\"pt\",\n", " max_length=max_length, \n", " truncation=True,\n", " )\n", "ds = DataLoader(ds2, batch_size=4, num_workers=0, collate_fn=collate_fn)\n", "print(ds)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Collect activations" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "['model.layers.0.mlp.down_proj',\n", " 'model.layers.1.mlp.down_proj',\n", " 'model.layers.2.mlp.down_proj',\n", " 'model.layers.3.mlp.down_proj',\n", " 'model.layers.4.mlp.down_proj',\n", " 'model.layers.5.mlp.down_proj',\n", " 'model.layers.6.mlp.down_proj',\n", " 'model.layers.7.mlp.down_proj',\n", " 'model.layers.8.mlp.down_proj',\n", " 'model.layers.9.mlp.down_proj',\n", " 'model.layers.10.mlp.down_proj',\n", " 'model.layers.11.mlp.down_proj',\n", " 'model.layers.12.mlp.down_proj',\n", " 'model.layers.13.mlp.down_proj',\n", " 'model.layers.14.mlp.down_proj',\n", " 'model.layers.15.mlp.down_proj',\n", " 'model.layers.16.mlp.down_proj',\n", " 'model.layers.17.mlp.down_proj',\n", " 'model.layers.18.mlp.down_proj',\n", " 'model.layers.19.mlp.down_proj',\n", " 'model.layers.20.mlp.down_proj',\n", " 'model.layers.21.mlp.down_proj',\n", " 'model.layers.22.mlp.down_proj',\n", " 'model.layers.23.mlp.down_proj']" ] }, "execution_count": 50, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# choose layers to cache\n", "layers = [k for k,v in model.named_modules() if k.endswith('mlp.down_proj')]\n", "layers" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2025-02-16 09:36:37.315\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m77\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__fac086acb713a85e.parquet\u001b[0m\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "8341bbff75634f0fb235e107abc2083d", "version_major": 2, "version_minor": 0 }, "text/plain": [ "collecting activations: 0%| | 0/5 [00:00 Float[Tensor, \"l b t h\"]:\n", " \"\"\"\n", " Novel experiment: Here we define a transform to isolate supressed activations, where we hypothesis that style/concepts/scratchpads and other internal only representations must be stored.\n", "\n", " See the following references for more information:\n", "\n", " - https://arxiv.org/pdf/2401.12181\n", " - > Suppression neurons that are similar, except decrease the probability of a group of related tokens\n", "\n", " - https://arxiv.org/html/2406.19384\n", " - > Previous work suggests that networks contain ensembles of “prediction\" neurons, which act as probability promoters [66, 24, 32] and work in tandem with suppression neurons (Section 5.4).\n", "\n", " - https://arxiv.org/pdf/2401.12181\n", " > We find a striking pattern which is remarkably consistent across the different seeds: after about the halfway point in the model, prediction neurons become increasingly prevalent until the very end of the network where there is a sudden shift towards a much larger number of suppression neurons.\n", " \"\"\"\n", " with torch.no_grad():\n", " # here we pass the hs through the last layer, take a diff, and then project it back to find which activation changes lead to supressed\n", " hs2 = rearrange(hs[:, :, -1:], \"l b t h -> (l b t) h\")\n", " hs_out2 = torch.nn.functional.linear(hs2, w_out)\n", " hs_out = rearrange(\n", " hs_out2, \"(l b t) h -> l b t h\", l=hs.shape[0], b=hs.shape[1], t=1\n", " )\n", " diffs = hs_out[:, :, :].diff(dim=0)\n", " diffs2 = rearrange(diffs, \"l b t h -> (l b t) h\")\n", " # W_inv = get_cache_inv(w_out)\n", "\n", " diffs_inv2 = torch.nn.functional.linear(diffs2.to(dtype=w_inv.dtype), w_inv)\n", " diffs_inv = rearrange(\n", " diffs_inv2, \"(l b t) h -> l b t h\", l=hs.shape[0] - 1, b=hs.shape[1], t=1\n", " ).to(w_out.dtype)\n", " # TODO just return this?\n", " eps = 1.0e-1\n", " supressed_mask = (diffs_inv < -eps).to(hs.dtype)\n", " # supressed_mask = repeat(supressed_mask, 'l b 1 h -> l b t h', t=hs.shape[2])\n", " supressed_act = hs[1:] * supressed_mask\n", " return supressed_act" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from activation_store.collect import default_postprocess_result\n", "\n", "Wo = model.get_output_embeddings().weight.detach().clone().cpu()\n", "Wo_inv = torch.pinverse(Wo.clone().float())\n", "\n", "@torch.no_grad()\n", "def sup_postproc(input, trace, output, model):\n", "\n", " \n", " o = default_postprocess_result(input, trace, output, model)\n", " \n", " hs = o.pop('hidden_states')\n", " hs = rearrange(hs, \"b l t h -> l b t h\")\n", " hs_s = get_supressed_activations(hs, Wo.to(hs.dtype), Wo_inv.to(hs.dtype))\n", " hs_s = rearrange(hs_s, \"l b t h -> b l t h\")\n", " o['hidden_states_supressed'] = hs_s.half()\n", " \n", " return o\n" ] }, { "cell_type": "code", "execution_count": 86, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2025-02-16 09:52:12.917\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m78\u001b[0m - \u001b[1mcreating dataset /media/wassname/SGIronWolf/projects5/elk/cache_transformer_acts/outputs/.ds/ds__115ab10dde7bd7a3.parquet\u001b[0m\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "c0c38f37f9934a0dbe7086b695624548", "version_major": 2, "version_minor": 0 }, "text/plain": [ "collecting activations: 0%| | 0/5 [00:00