This commit is contained in:
wassname
2025-05-04 14:56:08 +08:00
parent bd1229636c
commit 1f6c6d68a4
+34 -509
View File
@@ -13,7 +13,7 @@
},
{
"cell_type": "code",
"execution_count": 67,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -23,7 +23,7 @@
},
{
"cell_type": "code",
"execution_count": 68,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -35,7 +35,7 @@
},
{
"cell_type": "code",
"execution_count": 69,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -74,7 +74,7 @@
},
{
"cell_type": "code",
"execution_count": 70,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -101,7 +101,7 @@
},
{
"cell_type": "code",
"execution_count": 71,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -130,24 +130,9 @@
},
{
"cell_type": "code",
"execution_count": 72,
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "68b61ba11b344abc87c84919121f02cb",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading checkpoint shards: 0%| | 0/2 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"outputs": [],
"source": [
"\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
@@ -216,18 +201,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|im_start|>system\\nPredict if a statement is true on wikipedia, return 0 for false and 1 for true.\\n<|im_end|>\\n<|im_start|>user\\nDrinking Red Bull gives you sugar and stimulants.<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\nThe answer is '"
]
},
"execution_count": 49,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"tokenizer.batch_decode(ds2['input_ids'])[0]"
]
@@ -243,15 +217,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<torch.utils.data.dataloader.DataLoader object at 0x76585e484e80>\n"
]
}
],
"outputs": [],
"source": [
"collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
"ds = DataLoader(ds2, batch_size=batch_size, collate_fn=collate_fn)\n",
@@ -287,71 +253,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'mlp.down_proj': ['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',\n",
" 'model.layers.24.mlp.down_proj',\n",
" 'model.layers.25.mlp.down_proj'],\n",
" 'self_attn': ['model.layers.8.self_attn',\n",
" 'model.layers.9.self_attn',\n",
" 'model.layers.10.self_attn',\n",
" 'model.layers.11.self_attn',\n",
" 'model.layers.12.self_attn',\n",
" 'model.layers.13.self_attn',\n",
" 'model.layers.14.self_attn',\n",
" 'model.layers.15.self_attn',\n",
" 'model.layers.16.self_attn',\n",
" 'model.layers.17.self_attn',\n",
" 'model.layers.18.self_attn',\n",
" 'model.layers.19.self_attn',\n",
" 'model.layers.20.self_attn',\n",
" 'model.layers.21.self_attn',\n",
" 'model.layers.22.self_attn',\n",
" 'model.layers.23.self_attn',\n",
" 'model.layers.24.self_attn',\n",
" 'model.layers.25.self_attn'],\n",
" 'mlp.up_proj': ['model.layers.8.mlp.up_proj',\n",
" 'model.layers.9.mlp.up_proj',\n",
" 'model.layers.10.mlp.up_proj',\n",
" 'model.layers.11.mlp.up_proj',\n",
" 'model.layers.12.mlp.up_proj',\n",
" 'model.layers.13.mlp.up_proj',\n",
" 'model.layers.14.mlp.up_proj',\n",
" 'model.layers.15.mlp.up_proj',\n",
" 'model.layers.16.mlp.up_proj',\n",
" 'model.layers.17.mlp.up_proj',\n",
" 'model.layers.18.mlp.up_proj',\n",
" 'model.layers.19.mlp.up_proj',\n",
" 'model.layers.20.mlp.up_proj',\n",
" 'model.layers.21.mlp.up_proj',\n",
" 'model.layers.22.mlp.up_proj',\n",
" 'model.layers.23.mlp.up_proj',\n",
" 'model.layers.24.mlp.up_proj',\n",
" 'model.layers.25.mlp.up_proj']}"
]
},
"execution_count": 65,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"# choose layers to cache\n",
"n_layers = model.config.num_hidden_layers\n",
@@ -394,18 +296,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-110.parquet')"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"\n",
"acts_outfile = Path(f'/tmp/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}_v2.parquet')\n",
@@ -417,29 +308,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'mlp.down_proj': ['model.layers.14.mlp.down_proj',\n",
" 'model.layers.17.mlp.down_proj',\n",
" 'model.layers.20.mlp.down_proj',\n",
" 'model.layers.23.mlp.down_proj'],\n",
" 'self_attn': ['model.layers.14.self_attn',\n",
" 'model.layers.17.self_attn',\n",
" 'model.layers.20.self_attn',\n",
" 'model.layers.23.self_attn'],\n",
" 'mlp.up_proj': ['model.layers.14.mlp.up_proj',\n",
" 'model.layers.17.mlp.up_proj',\n",
" 'model.layers.20.mlp.up_proj',\n",
" 'model.layers.23.mlp.up_proj']}"
]
},
"execution_count": 53,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"def collect_all_tokens(*args, **kwargs):\n",
" return default_postprocess_result(*args, **kwargs, last_token=False)\n",
@@ -455,35 +324,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "d4d0acec794f4102a361592f358ed92b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading dataset shards: 0%| | 0/53 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"Dataset({\n",
" features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label'],\n",
" num_rows: 316\n",
"})"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"# TODO which is better for mem, this or below?\n",
"ds_a = load_dataset(\"parquet\", split='train', data_files=str(f), keep_in_memory=False).with_format(\"torch\")\n",
@@ -494,18 +335,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
]
},
"execution_count": 55,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"# # TODO which is better for mem, this or above?\n",
"# ds_a = Dataset.from_parquet(str(f), split=split, keep_in_memory=False).with_format(\"torch\")\n",
@@ -516,25 +346,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2025-05-04 14:49:53.530\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m174\u001b[0m - \u001b[33m\u001b[1mfile /tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet already exists, skipping\u001b[0m\n"
]
},
{
"data": {
"text/plain": [
"PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
]
},
"execution_count": 56,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"act_groups = [c for c in ds_a.column_names if c.startswith('acts-')]\n",
"act_groups"
@@ -544,49 +356,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c331b2a7429145d6a9618e9c0a76816b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading dataset shards: 0%| | 0/27 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c331b2a7429145d6a9618e9c0a76816b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading dataset shards: 0%| | 0/27 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"Dataset({\n",
" features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label'],\n",
" num_rows: 316\n",
"})"
]
},
"execution_count": 57,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"ds_a"
]
@@ -595,22 +365,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"acts-mlp.down_proj torch.Size([18, 110, 2048])\n",
"acts-self_attn torch.Size([18, 110, 2048])\n",
"acts-mlp.up_proj torch.Size([18, 110, 6144])\n",
"loss torch.Size([])\n",
"logits torch.Size([110, 151936])\n",
"hidden_states torch.Size([29, 110, 2048])\n",
"attention_mask torch.Size([110])\n",
"label torch.Size([])\n"
]
}
],
"outputs": [],
"source": [
"for k,v in ds_a[0].items():\n",
" if hasattr(v, 'shape'):\n",
@@ -623,18 +378,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']"
]
},
"execution_count": 59,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"# sanity test generate\n",
"b = next(iter(ds))\n",
@@ -662,22 +406,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"acts-mlp.down_proj torch.Size([4, 90, 2048])\n",
"acts-self_attn torch.Size([4, 90, 2048])\n",
"acts-mlp.up_proj torch.Size([4, 90, 6144])\n",
"loss torch.Size([])\n",
"logits torch.Size([90, 151936])\n",
"hidden_states torch.Size([29, 90, 2048])\n",
"attention_mask torch.Size([90])\n",
"label torch.Size([])\n"
]
}
],
"outputs": [],
"source": [
"@torch.no_grad()\n",
"def get_supressed_activations(\n",
@@ -726,19 +455,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"before ['0', '0 ', '0\\n', 'false', 'False ']\n",
"after ['false', 'False', '0']\n",
"before ['1', '1 ', '1\\n', 'true', 'True ']\n",
"after ['1', 'True', 'true']\n",
"QC: manually check that these are equivilent (no <end_of_text> or newline)\n"
]
}
],
"outputs": [],
"source": [
"def get_uniq_token_ids(tokens):\n",
" token_ids = tokenizer(\n",
@@ -766,16 +483,17 @@
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "af953342d91b4ed38e5dbaae9e9a8c64",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Dataset({\n",
" features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts'],\n",
" num_rows: 316\n",
"})"
"Map: 0%| | 0/316 [00:00<?, ? examples/s]"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
"output_type": "display_data"
}
],
"source": [
@@ -821,19 +539,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"before ['0', '0 ', '0\\n', 'false', 'False ']\n",
"after ['false', 'False', '0']\n",
"before ['1', '1 ', '1\\n', 'true', 'True ']\n",
"after ['1', 'True', 'true']\n",
"QC: manually check that these are equivilent (no <end_of_text> or newline)\n"
]
}
],
"outputs": [],
"source": [
"model = Wo = Wo_inv = tokenizer = None\n",
"clear_mem()"
@@ -843,71 +549,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "17b983edf4d64ca995dc90e063b4eb8a",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Map: 0%| | 0/316 [00:00<?, ? examples/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# # FIXME proboly not need to store all these\n",
"# def calc_hs_sup(o, threshholds = [-10, -5, -1, -0.5, -0.1, -0.01, -0, 0, 0.01, 0.1, 0.5, 1, 10]):\n",
"# \"\"\"\n",
"# Calc supressed activations for a certain threshold\n",
"# \"\"\"\n",
"# supr_amounts = o[\"supr_amounts\"]\n",
"# hs = o[\"hidden_states\"] # [b l h]\n",
"# out = {}\n",
"# for thresh in threshholds:\n",
"# if thresh > 0:\n",
"# supressed_mask = (supr_amounts > thresh).to(hs.dtype)# [b l h]\n",
"# else:\n",
"# supressed_mask = (supr_amounts < thresh).to(hs.dtype)\n",
"# out[f'supressed_hs_{thresh}'] = hs * supressed_mask\n",
"# return out\n",
"\n",
"# ds_a2 = ds_a2.map(calc_hs_sup, writer_batch_size=1)\n",
"\n",
"# # for eps in [-10, -5, -1, -0.5, -0.1, -0.01, -0, 0, 0.01, 0.1, 0.5, 1, 10]:\n",
"# # ds_a2 = ds_a2.map(calc_hs_sup, fn_kwargs={'eps': eps}, writer_batch_size=1, num_proc=None)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'acts-mlp.down_proj': torch.Size([18, 110, 2048]),\n",
" 'acts-self_attn': torch.Size([18, 110, 2048]),\n",
" 'acts-mlp.up_proj': torch.Size([18, 110, 6144]),\n",
" 'loss': torch.Size([]),\n",
" 'logits': torch.Size([110, 151936]),\n",
" 'hidden_states': torch.Size([13, 110, 2048]),\n",
" 'attention_mask': torch.Size([110]),\n",
" 'label': torch.Size([]),\n",
" 'llm_ans': torch.Size([2]),\n",
" 'llm_log_prob_true': torch.Size([]),\n",
" 'supr_amounts': torch.Size([13, 1, 2048])}"
]
},
"execution_count": 24,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"{k: v.shape for k,v in ds_a2[0].items() if isinstance(v, torch.Tensor)}\n"
]
@@ -916,21 +558,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Dataset({\n",
" features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts'],\n",
" num_rows: 316\n",
"})"
]
},
"execution_count": 25,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"ds_a2"
]
@@ -946,18 +574,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"88"
]
},
"execution_count": 26,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"test_fraction = 0.2\n",
"TRAIN_TEST_SPLIT = int(max_length * (1- test_fraction))\n",
@@ -1282,39 +899,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n",
"threshold: 0.02307834103703499\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n",
"threshold: 0.02307834103703499\n"
]
},
{
"data": {
"text/plain": [
"(torch.Size([316, 13, 110, 2048]),\n",
" torch.Size([316, 13, 110, 2048]),\n",
" tensor(4856864.5000),\n",
" tensor(4856636.5000),\n",
" torch.Size([316, 13, 110, 1]),\n",
" tensor(0.9752))"
]
},
"execution_count": 81,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"# X = ds_a2['hidden_states']\n",
"# X2, mask = entropy_guided_filter(X)\n",
@@ -1382,19 +967,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[('supressed_hs_-0.5', 'magnitude', 'first'),\n",
" ('supressed_hs_-0.5', 'entropy', 'sum')]"
]
},
"execution_count": 49,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"X_cols = [\"supr_amounts\", \"hidden_states\",] + act_groups\n",
"\n",
@@ -1543,54 +1116,6 @@
" clear_mem()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"X.view(len(X), -1)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"# results = []\n",
"\n",
"\n",
"\n",
"# # First loop: Use only the token-level reductions.\n",
"# for ds_col in tqdm(X_cols):\n",
"# for token_key, token_func in token_level_funcs.items():\n",
"# reduction_name = f\"{ds_col} {token_key}\"\n",
"# try:\n",
"# X = torch.stack([token_func(x.float()) for x in ds_a2[ds_col]])\n",
"# score = train_linear_prob_on_dataset(X, reduction_name)\n",
"# results.append((reduction_name, score))\n",
"# except Exception as e:\n",
"# logger.error(f\"error with {reduction_name} {e}\")\n",
"# continue\n",
"# clear_mem()\n",
"\n",
"# # Second loop: Apply token-level reduction then aggregate the token outputs.\n",
"# for ds_col in tqdm(X_cols):\n",
"# for token_key, token_func in token_level_funcs.items():\n",
"# for final_key, final_func in token_aggs.items():\n",
"# reduction_name = f\"{ds_col} {token_key}_{final_key}\"\n",
"# try:\n",
"# X = torch.stack([final_func(token_func(x.float())) for x in ds_a2[ds_col]])\n",
"# score = train_linear_prob_on_dataset(X, reduction_name)\n",
"# results.append((reduction_name, score))\n",
"# except Exception as e:\n",
"# logger.error(f\"error with {reduction_name} {e}\")\n",
"# continue\n",
"# clear_mem()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -1638,7 +1163,7 @@
"metadata": {},
"outputs": [],
"source": [
"X = torch.sigmoid(ds_a2['llm_log_prob_true']/10)\n",
"X = torch.sigmoid(ds_a2['llm_log_prob_true']/10) # would be better to calibrate or logreg\n",
"y = ds_a2['label']\n",
"\n",
"X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_fraction, shuffle=False)\n",