wassname/LoRA_are_lie_detectors
1
0
Fork 0
mirror of https://github.com/wassname/LoRA_are_lie_detectors.git synced 2026-09-10 11:40:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

out of gpu mem

Browse Source
This commit is contained in:
wassname
2023-12-18 10:29:32 +08:00
parent a53f328d25
commit bd2fabb1bf
15 changed files with 1161 additions and 1384 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -1,3 +1,5 @@
notebooks/phi_azure_deleteme/
lightning_logs/
.pkl_cache/
.ds/
@@ -164,4 +166,3 @@ cython_debug/
*.bin
.DS_Store
notebooks/01_mjc.ipynb
+710 -15
View File
@@ -2,9 +2,38 @@
"cells": [
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"# os.environ[\"CUDA_VISIBLE_DEVICES\"]= \"1\"\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
},
{
"data": {
"text/plain": [
"1"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import os\n",
"import numpy as np\n",
@@ -26,7 +55,9 @@
"\n",
"from pathlib import Path\n",
"import transformers\n",
"\n",
"from peft import get_peft_config, get_peft_model, LoraConfig, TaskType, LoftQConfig\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
"from datasets import Dataset\n",
"\n",
"from loguru import logger\n",
"\n",
@@ -35,7 +66,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
@@ -43,27 +74,157 @@
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"from src.extraction.config import ExtractConfig\n",
"from src.config import ExtractConfig\n",
"from src.prompts.prompt_loading import load_preproc_dataset\n",
"from src.models.load import load_model\n",
"# from src.prompts.prompt_loading import load_prompt_structure\n",
"# from src.prompts.prompt_loading import load_prompt_structure\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# params\n",
"max_epochs = 100\n",
"device = \"cuda:0\"\n",
"\n",
"# quiet please\n",
"torch.set_float32_matmul_precision(\"medium\")\n",
"import warnings\n",
"\n",
"warnings.filterwarnings(\"ignore\", \".*does not have many workers.*\")\n",
"warnings.filterwarnings(\n",
" \"ignore\", \".*sampler has shuffling enabled, it is strongly recommended that.*\"\n",
")\n",
"warnings.filterwarnings(\"ignore\", \".*has been removed as a dependency of.*\")\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2023-12-18 10:28:32.914\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m16\u001b[0m - \u001b[1mtokenizer does not have use_cache\u001b[0m\n",
"2023-12-18T10:28:32.914981+0800 INFO tokenizer does not have use_cache\n",
"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
"\u001b[32m2023-12-18 10:28:33.244\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m21\u001b[0m - \u001b[1mchanging pad_token_id from None to 0\u001b[0m\n",
"2023-12-18T10:28:33.244305+0800 INFO changing pad_token_id from None to 0\n",
"\u001b[32m2023-12-18 10:28:33.244\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m21\u001b[0m - \u001b[1mchanging padding_side from right to left\u001b[0m\n",
"2023-12-18T10:28:33.244820+0800 INFO changing padding_side from right to left\n",
"\u001b[32m2023-12-18 10:28:33.245\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.models.load\u001b[0m:\u001b[36mverbose_change_param\u001b[0m:\u001b[36m21\u001b[0m - \u001b[1mchanging truncation_side from right to left\u001b[0m\n",
"2023-12-18T10:28:33.245209+0800 INFO changing truncation_side from right to left\n",
"Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00, 1.51it/s]\n"
]
}
],
"source": [
"# params\n",
"cfg = ExtractConfig(\n",
" batch_size=2,\n",
" max_examples=(400, 400),\n",
" intervention_fit_examples=160,\n",
")\n",
"model, tokenizer = load_model(\n",
" cfg.model, disable_exllama=False, device=device,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"# model.to(device)\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"# TODO I would like to only have biases, but for now lets just try a very small intervention on the last parts of a layer...\n",
"peft_config = LoraConfig(\n",
" target_modules=['out_proj', 'mlp.fc2',], # only the layers that go directly to the residual\n",
" bias='lora_only',\n",
" task_type=TaskType.CAUSAL_LM,\n",
" inference_mode=False,\n",
" r=1,\n",
" lora_alpha=1,\n",
" lora_dropout=0.,\n",
" # layers_pattern='dsf4gg',\n",
")\n",
"# model = get_peft_model(model, peft_config)\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"# params\n",
"cfg = ExtractConfig(\n",
" max_examples=(400, 400),\n",
" intervention_fit_examples=160,\n",
")\n",
"model, tokenizer = load_model(\n",
" cfg.model, pad_token_id=cfg.pad_token_id, disable_exllama=False\n",
")\n",
"# model.to(device)\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"model.add_adapter(peft_config)\n",
"# model.print_trainable_parameters()\n"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"device(type='cuda', index=0)"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model.device\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2023-12-18 10:28:36.441\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m364\u001b[0m - \u001b[1mmedian token length: 433.0 for amazon_polarity. max_length=1000\u001b[0m\n",
"2023-12-18T10:28:36.441848+0800 INFO median token length: 433.0 for amazon_polarity. max_length=1000\n",
"\u001b[32m2023-12-18 10:28:36.443\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m368\u001b[0m - \u001b[1mtruncation rate: 0.00% on amazon_polarity\u001b[0m\n",
"2023-12-18T10:28:36.443448+0800 INFO truncation rate: 0.00% on amazon_polarity\n",
"Filter: 100%|██████████| 2402/2402 [00:01<00:00, 2114.02 examples/s]\n",
"\u001b[32m2023-12-18 10:28:37.594\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msrc.prompts.prompt_loading\u001b[0m:\u001b[36mload_preproc_dataset\u001b[0m:\u001b[36m377\u001b[0m - \u001b[1mnum_rows (after filtering out truncated rows) 2402=>2402\u001b[0m\n",
"2023-12-18T10:28:37.594298+0800 INFO num_rows (after filtering out truncated rows) 2402=>2402\n"
]
}
],
"source": [
"N = sum(cfg.max_examples)\n",
"ds_name = \"amazon_polarity\"\n",
"ds_tokens = load_preproc_dataset(\n",
" ds_name,\n",
@@ -75,11 +236,545 @@
" prompt_format=cfg.prompt_format,\n",
")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Lora train"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"# from https://github.com/jonkrohn/NLP-with-LLMs/blob/main/code/Finetune-T5-on-GPU.ipynb\n",
"from pytorch_optimizer import Ranger21\n",
"import lightning.pytorch as pl\n",
"from torchmetrics import Metric, MetricCollection, Accuracy, AUROC\n",
"from torchmetrics.functional import accuracy\n"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"# to_tensor = lambda x: x # torch.from_numpy(x).float()\n",
"# to_ds = lambda hs0, hs1, y: TensorDataset(to_tensor(hs0), to_tensor(hs1), to_tensor(y))\n",
"\n",
"class DeceptionDataModule(pl.LightningDataModule):\n",
"\n",
" def __init__(self,\n",
" ds: Dataset,\n",
" batch_size: int=32,\n",
" # x_cols = ['input_ids', 'attention_mask', 'label_true', 'label_instructed', 'choice_ids'],\n",
" ):\n",
" super().__init__()\n",
" self.save_hyperparameters(ignore=[\"ds\"])\n",
" self.ds = ds.with_format('torch')\n",
" # self.x_cols = x_cols\n",
" self.setup('train')\n",
"\n",
" def setup(self, stage: str):\n",
" h = self.hparams\n",
"\n",
" n = len(self.ds)\n",
" self.splits = {\n",
" 'train': (0, int(n * 0.5)),\n",
" 'val': (int(n * 0.5), int(n * 0.75)),\n",
" 'test': (int(n * 0.75), n),\n",
" }\n",
"\n",
" self.datasets = {key: self.ds.select(range(start, end)) for key, (start, end) in self.splits.items()}\n",
"\n",
" def create_dataloader(self, ds, shuffle=False):\n",
" return DataLoader(ds, batch_size=self.hparams.batch_size, drop_last=False, shuffle=shuffle)\n",
"\n",
" def train_dataloader(self):\n",
" return self.create_dataloader(self.datasets['train'], shuffle=True)\n",
"\n",
" def val_dataloader(self):\n",
" return self.create_dataloader(self.datasets['val'])\n",
"\n",
" def test_dataloader(self):\n",
" return self.create_dataloader(self.datasets['test'])\n",
"\n",
"# https://huggingface.co/docs/datasets/use_with_pytorch#data-loading\n"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"\n",
"\n",
"\n",
"class LoraFinetuner(pl.LightningModule):\n",
" def __init__(\n",
" self, model: AutoModelForCausalLM, tokenizerm: AutoTokenizer, total_steps: int, lr=4e-3, weight_decay=1e-9\n",
" ):\n",
" super().__init__()\n",
" self.model = model\n",
" self.tokenizer = tokenizer\n",
" self.save_hyperparameters(\n",
" ignore=[\"model\", 'tokenizer'],\n",
" )\n",
"\n",
" def forward(self, batch):\n",
" b_in = dict(input_ids=batch['input_ids'],\n",
" attention_mask=batch['attention_mask'],)\n",
" # b_in = {k: v.to(self.model.device) for k, v in b_in.items()}\n",
" \n",
" return self.model(\n",
" **b_in,\n",
" use_cache=False,\n",
" output_hidden_states=True,\n",
" return_dict=True\n",
" )\n",
" # odict_keys(['logits', 'hidden_states', 'attentions'])\n",
"\n",
" def _step(self, batch, batch_idx=0, stage=\"train\"):\n",
" self.model.disable_adapters()\n",
" with torch.no_grad():\n",
" out = self(batch)\n",
" log_probs = torch.log_softmax(out['logits'][:, -1,], -1)\n",
" del out\n",
" \n",
" self.model.enable_adapters()\n",
" out2 = self(batch)\n",
" log_probs2 = torch.log_softmax(out2['logits'][:, -1,], -1)\n",
"\n",
" if stage == \"pred\":\n",
" return log_probs2.exp()\n",
" \n",
" \n",
" # get loss, so that our adapter returns switched probs for our choices (e.g. Yes <> No)\n",
" id_neg = batch['choice_ids'][:, 0]\n",
" id_pos = batch['choice_ids'][:, 1]\n",
"\n",
" opposite_log_probs = log_probs.clone()\n",
" for i in range(id_neg.shape[1]):\n",
" opposite_log_probs[:, id_neg[:, i]] = log_probs[:, id_pos[:, i]]\n",
" loss = F.kl_div(log_probs2, opposite_log_probs, reduction='batchmean', log_target=True)\n",
"\n",
" self.log(f\"{stage}/loss\", loss, on_epoch=True, on_step=False, prog_bar=True)\n",
" self.log(\n",
" f\"{stage}/n\", len(id_neg), on_epoch=True, on_step=False, reduce_fx=torch.sum\n",
" )\n",
" return loss\n",
"\n",
" def training_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
" return self._step(batch, batch_idx)\n",
"\n",
" def validation_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
" return self._step(batch, batch_idx, stage=\"val\")\n",
"\n",
" def predict_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
" return self._step(batch, batch_idx, stage=\"pred\").cpu().detach()\n",
"\n",
" def test_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
" return self._step(batch, batch_idx, stage=\"test\")\n",
"\n",
" def configure_optimizers(self):\n",
" \"\"\"use ranger21 from https://github.com/kozistr/pytorch_optimizer\"\"\"\n",
" optimizer = Ranger21(\n",
" self.parameters(),\n",
" lr=self.hparams.lr,\n",
" weight_decay=self.hparams.weight_decay,\n",
" num_iterations=self.hparams.total_steps,\n",
" )\n",
" return optimizer\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<__main__.DeceptionDataModule at 0x7f8f58429390>"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"dm = DeceptionDataModule(ds_tokens, batch_size=cfg.batch_size)\n",
"dm\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"dict_keys(['ds_string', 'example_i', 'answer', 'messages', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'question', 'input_ids', 'attention_mask', 'truncated', 'length', 'prompt_truncated', 'choice_ids']) torch.Size([2, 1000])\n"
]
},
{
"data": {
"text/plain": [
"1000"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"dl_train = dm.train_dataloader()\n",
"dl_val = dm.val_dataloader()\n",
"b = next(iter(dl_train))\n",
"print(b.keys(), b['input_ids'].shape)\n",
"c_in = b['input_ids'].shape[1]\n",
"c_in\n"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [],
"source": [
"# from accelerate import Accelerator\n",
"\n",
"# accelerator = Accelerator(device_placement=False)\n",
"# model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n",
"# model, None, dl_train, dl_val, device_placement=[False, False, False, False]\n",
"# )\n"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"dict_keys(['ds_string', 'example_i', 'answer', 'messages', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'question', 'input_ids', 'attention_mask', 'truncated', 'length', 'prompt_truncated', 'choice_ids']) torch.Size([2, 1000])\n"
]
},
{
"data": {
"text/plain": [
"1000"
]
},
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"b = next(iter(dl_train))\n",
"print(b.keys(), b['input_ids'].shape)\n",
"c_in = b['input_ids'].shape[1]\n",
"c_in\n"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1000\n"
]
}
],
"source": [
"net = LoraFinetuner(model, tokenizer, lr=3e-4, weight_decay=1e-9, total_steps=len(dl_train)*max_epochs)\n",
"\n",
"print(c_in)\n",
"# net.model.enable_adapters()\n",
"\n",
"\n",
"# net = accelerator.prepare(\n",
"# net, device_placement=[False]\n",
"# )\n"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"GPU available: True (cuda), used: True\n",
"TPU available: False, using: 0 TPU cores\n",
"IPU available: False, using: 0 IPUs\n",
"HPU available: False, using: 0 HPUs\n"
]
}
],
"source": [
"# we want to init lightning early, so it inits accelerate\n",
"trainer1 = pl.Trainer(\n",
" precision=\"16-true\",\n",
" # precision=\"16-mixed\",\n",
" # precision=\"b16-mixed\",\n",
" # precision=\"b16-mixed\",\n",
" # gradient_clip_val=20,\n",
"\n",
" # accelerator=\"auto\",\n",
" # devices=\"1\",\n",
"\n",
" accelerator=\"gpu\",\n",
" devices=[0],\n",
"\n",
" max_epochs=max_epochs,\n",
" log_every_n_steps=3,\n",
" # enable_progress_bar=False, \n",
" enable_model_summary=False\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"device(type='cuda', index=0)"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"net.model.device\n"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"# b_in = dict(input_ids=b['input_ids'],\n",
"# attention_mask=b['attention_mask'],)\n",
"# b_in = {k: v.to(net.model.device) for k, v in b_in.items()}\n",
"\n",
"# with torch.no_grad():\n",
"# y = net(b_in, )\n",
"# y.keys()\n"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [],
"source": [
"# from torchinfo import summary\n",
"# net.model.disable_adapters()\n",
"# summary(net, input_data=(dict(input_ids=b['input_ids'], attention_mask=b['attention_mask'],),), depth=4)\n"
]
},
{
"cell_type": "code",
"execution_count": 24,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sanity Checking DataLoader 0: 50%|█████ | 1/2 [00:00<00:00, 1.08it/s]"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/utilities/data.py:77: Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is 2. To avoid any miscalculations, use `self.log(..., batch_size=batch_size)`.\n",
"/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: You called `self.log('val/n', ...)` in your `validation_step` but the value needs to be floating to be reduced. Converting it to torch.float32. You can silence this warning by converting the value to floating point yourself. If you don't intend to reduce the value (for instance when logging the global step or epoch) then you can use `self.logger.log_metrics({'val/n': ...})` instead.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: 0%| | 0/200 [00:00<?, ?it/s] "
]
},
{
"ename": "OutOfMemoryError",
"evalue": "CUDA out of memory. Tried to allocate 50.00 MiB. GPU 0 has a total capacty of 23.67 GiB of which 1.13 GiB is free. Including non-PyTorch memory, this process has 21.46 GiB memory in use. Of the allocated memory 20.92 GiB is allocated by PyTorch, and 235.11 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mOutOfMemoryError\u001b[0m Traceback (most recent call last)",
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb Cell 26\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m trainer1\u001b[39m.\u001b[39;49mfit(model\u001b[39m=\u001b[39;49mnet, train_dataloaders\u001b[39m=\u001b[39;49mdl_train, val_dataloaders\u001b[39m=\u001b[39;49mdl_val);\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:544\u001b[0m, in \u001b[0;36mTrainer.fit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m 542\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mstatus \u001b[39m=\u001b[39m TrainerStatus\u001b[39m.\u001b[39mRUNNING\n\u001b[1;32m 543\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining \u001b[39m=\u001b[39m \u001b[39mTrue\u001b[39;00m\n\u001b[0;32m--> 544\u001b[0m call\u001b[39m.\u001b[39;49m_call_and_handle_interrupt(\n\u001b[1;32m 545\u001b[0m \u001b[39mself\u001b[39;49m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path\n\u001b[1;32m 546\u001b[0m )\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py:44\u001b[0m, in \u001b[0;36m_call_and_handle_interrupt\u001b[0;34m(trainer, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m 42\u001b[0m \u001b[39mif\u001b[39;00m trainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mlauncher \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 43\u001b[0m \u001b[39mreturn\u001b[39;00m trainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mlauncher\u001b[39m.\u001b[39mlaunch(trainer_fn, \u001b[39m*\u001b[39margs, trainer\u001b[39m=\u001b[39mtrainer, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m---> 44\u001b[0m \u001b[39mreturn\u001b[39;00m trainer_fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 46\u001b[0m \u001b[39mexcept\u001b[39;00m _TunerExitException:\n\u001b[1;32m 47\u001b[0m _call_teardown_hook(trainer)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:580\u001b[0m, in \u001b[0;36mTrainer._fit_impl\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m 573\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mfn \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m 574\u001b[0m ckpt_path \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_checkpoint_connector\u001b[39m.\u001b[39m_select_ckpt_path(\n\u001b[1;32m 575\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mfn,\n\u001b[1;32m 576\u001b[0m ckpt_path,\n\u001b[1;32m 577\u001b[0m model_provided\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m,\n\u001b[1;32m 578\u001b[0m model_connected\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlightning_module \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m 579\u001b[0m )\n\u001b[0;32m--> 580\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_run(model, ckpt_path\u001b[39m=\u001b[39;49mckpt_path)\n\u001b[1;32m 582\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m.\u001b[39mstopped\n\u001b[1;32m 583\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:989\u001b[0m, in \u001b[0;36mTrainer._run\u001b[0;34m(self, model, ckpt_path)\u001b[0m\n\u001b[1;32m 984\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_signal_connector\u001b[39m.\u001b[39mregister_signal_handlers()\n\u001b[1;32m 986\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m 987\u001b[0m \u001b[39m# RUN THE TRAINER\u001b[39;00m\n\u001b[1;32m 988\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[0;32m--> 989\u001b[0m results \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_run_stage()\n\u001b[1;32m 991\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m 992\u001b[0m \u001b[39m# POST-Training CLEAN UP\u001b[39;00m\n\u001b[1;32m 993\u001b[0m \u001b[39m# ----------------------------\u001b[39;00m\n\u001b[1;32m 994\u001b[0m log\u001b[39m.\u001b[39mdebug(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m: trainer tearing down\u001b[39m\u001b[39m\"\u001b[39m)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:1035\u001b[0m, in \u001b[0;36mTrainer._run_stage\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 1033\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_run_sanity_check()\n\u001b[1;32m 1034\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mautograd\u001b[39m.\u001b[39mset_detect_anomaly(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_detect_anomaly):\n\u001b[0;32m-> 1035\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mfit_loop\u001b[39m.\u001b[39;49mrun()\n\u001b[1;32m 1036\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m 1037\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mUnexpected state \u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstate\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py:202\u001b[0m, in \u001b[0;36m_FitLoop.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 200\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 201\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mon_advance_start()\n\u001b[0;32m--> 202\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49madvance()\n\u001b[1;32m 203\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mon_advance_end()\n\u001b[1;32m 204\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_restarting \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py:359\u001b[0m, in \u001b[0;36m_FitLoop.advance\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 357\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtrainer\u001b[39m.\u001b[39mprofiler\u001b[39m.\u001b[39mprofile(\u001b[39m\"\u001b[39m\u001b[39mrun_training_epoch\u001b[39m\u001b[39m\"\u001b[39m):\n\u001b[1;32m 358\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_data_fetcher \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[0;32m--> 359\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mepoch_loop\u001b[39m.\u001b[39;49mrun(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_data_fetcher)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py:136\u001b[0m, in \u001b[0;36m_TrainingEpochLoop.run\u001b[0;34m(self, data_fetcher)\u001b[0m\n\u001b[1;32m 134\u001b[0m \u001b[39mwhile\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdone:\n\u001b[1;32m 135\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 136\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49madvance(data_fetcher)\n\u001b[1;32m 137\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mon_advance_end(data_fetcher)\n\u001b[1;32m 138\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_restarting \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py:240\u001b[0m, in \u001b[0;36m_TrainingEpochLoop.advance\u001b[0;34m(self, data_fetcher)\u001b[0m\n\u001b[1;32m 237\u001b[0m \u001b[39mwith\u001b[39;00m trainer\u001b[39m.\u001b[39mprofiler\u001b[39m.\u001b[39mprofile(\u001b[39m\"\u001b[39m\u001b[39mrun_training_batch\u001b[39m\u001b[39m\"\u001b[39m):\n\u001b[1;32m 238\u001b[0m \u001b[39mif\u001b[39;00m trainer\u001b[39m.\u001b[39mlightning_module\u001b[39m.\u001b[39mautomatic_optimization:\n\u001b[1;32m 239\u001b[0m \u001b[39m# in automatic optimization, there can only be one optimizer\u001b[39;00m\n\u001b[0;32m--> 240\u001b[0m batch_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mautomatic_optimization\u001b[39m.\u001b[39;49mrun(trainer\u001b[39m.\u001b[39;49moptimizers[\u001b[39m0\u001b[39;49m], batch_idx, kwargs)\n\u001b[1;32m 241\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 242\u001b[0m batch_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmanual_optimization\u001b[39m.\u001b[39mrun(kwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py:187\u001b[0m, in \u001b[0;36m_AutomaticOptimization.run\u001b[0;34m(self, optimizer, batch_idx, kwargs)\u001b[0m\n\u001b[1;32m 180\u001b[0m closure()\n\u001b[1;32m 182\u001b[0m \u001b[39m# ------------------------------\u001b[39;00m\n\u001b[1;32m 183\u001b[0m \u001b[39m# BACKWARD PASS\u001b[39;00m\n\u001b[1;32m 184\u001b[0m \u001b[39m# ------------------------------\u001b[39;00m\n\u001b[1;32m 185\u001b[0m \u001b[39m# gradient update with accumulated gradients\u001b[39;00m\n\u001b[1;32m 186\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 187\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_optimizer_step(batch_idx, closure)\n\u001b[1;32m 189\u001b[0m result \u001b[39m=\u001b[39m closure\u001b[39m.\u001b[39mconsume_result()\n\u001b[1;32m 190\u001b[0m \u001b[39mif\u001b[39;00m result\u001b[39m.\u001b[39mloss \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py:265\u001b[0m, in \u001b[0;36m_AutomaticOptimization._optimizer_step\u001b[0;34m(self, batch_idx, train_step_and_backward_closure)\u001b[0m\n\u001b[1;32m 262\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moptim_progress\u001b[39m.\u001b[39moptimizer\u001b[39m.\u001b[39mstep\u001b[39m.\u001b[39mincrement_ready()\n\u001b[1;32m 264\u001b[0m \u001b[39m# model hook\u001b[39;00m\n\u001b[0;32m--> 265\u001b[0m call\u001b[39m.\u001b[39;49m_call_lightning_module_hook(\n\u001b[1;32m 266\u001b[0m trainer,\n\u001b[1;32m 267\u001b[0m \u001b[39m\"\u001b[39;49m\u001b[39moptimizer_step\u001b[39;49m\u001b[39m\"\u001b[39;49m,\n\u001b[1;32m 268\u001b[0m trainer\u001b[39m.\u001b[39;49mcurrent_epoch,\n\u001b[1;32m 269\u001b[0m batch_idx,\n\u001b[1;32m 270\u001b[0m optimizer,\n\u001b[1;32m 271\u001b[0m train_step_and_backward_closure,\n\u001b[1;32m 272\u001b[0m )\n\u001b[1;32m 274\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m should_accumulate:\n\u001b[1;32m 275\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moptim_progress\u001b[39m.\u001b[39moptimizer\u001b[39m.\u001b[39mstep\u001b[39m.\u001b[39mincrement_completed()\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py:157\u001b[0m, in \u001b[0;36m_call_lightning_module_hook\u001b[0;34m(trainer, hook_name, pl_module, *args, **kwargs)\u001b[0m\n\u001b[1;32m 154\u001b[0m pl_module\u001b[39m.\u001b[39m_current_fx_name \u001b[39m=\u001b[39m hook_name\n\u001b[1;32m 156\u001b[0m \u001b[39mwith\u001b[39;00m trainer\u001b[39m.\u001b[39mprofiler\u001b[39m.\u001b[39mprofile(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m[LightningModule]\u001b[39m\u001b[39m{\u001b[39;00mpl_module\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m{\u001b[39;00mhook_name\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m):\n\u001b[0;32m--> 157\u001b[0m output \u001b[39m=\u001b[39m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 159\u001b[0m \u001b[39m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m 160\u001b[0m pl_module\u001b[39m.\u001b[39m_current_fx_name \u001b[39m=\u001b[39m prev_fx_name\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/core/module.py:1282\u001b[0m, in \u001b[0;36mLightningModule.optimizer_step\u001b[0;34m(self, epoch, batch_idx, optimizer, optimizer_closure)\u001b[0m\n\u001b[1;32m 1243\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39moptimizer_step\u001b[39m(\n\u001b[1;32m 1244\u001b[0m \u001b[39mself\u001b[39m,\n\u001b[1;32m 1245\u001b[0m epoch: \u001b[39mint\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1248\u001b[0m optimizer_closure: Optional[Callable[[], Any]] \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m 1249\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 1250\u001b[0m \u001b[39m \u001b[39m\u001b[39mr\u001b[39m\u001b[39m\"\"\"Override this method to adjust the default way the :class:`~lightning.pytorch.trainer.trainer.Trainer` calls\u001b[39;00m\n\u001b[1;32m 1251\u001b[0m \u001b[39m the optimizer.\u001b[39;00m\n\u001b[1;32m 1252\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1280\u001b[0m \n\u001b[1;32m 1281\u001b[0m \u001b[39m \"\"\"\u001b[39;00m\n\u001b[0;32m-> 1282\u001b[0m optimizer\u001b[39m.\u001b[39;49mstep(closure\u001b[39m=\u001b[39;49moptimizer_closure)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/core/optimizer.py:151\u001b[0m, in \u001b[0;36mLightningOptimizer.step\u001b[0;34m(self, closure, **kwargs)\u001b[0m\n\u001b[1;32m 148\u001b[0m \u001b[39mraise\u001b[39;00m MisconfigurationException(\u001b[39m\"\u001b[39m\u001b[39mWhen `optimizer.step(closure)` is called, the closure should be callable\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m 150\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_strategy \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[0;32m--> 151\u001b[0m step_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_strategy\u001b[39m.\u001b[39;49moptimizer_step(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_optimizer, closure, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 153\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_on_after_step()\n\u001b[1;32m 155\u001b[0m \u001b[39mreturn\u001b[39;00m step_output\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py:230\u001b[0m, in \u001b[0;36mStrategy.optimizer_step\u001b[0;34m(self, optimizer, closure, model, **kwargs)\u001b[0m\n\u001b[1;32m 228\u001b[0m \u001b[39m# TODO(fabric): remove assertion once strategy's optimizer_step typing is fixed\u001b[39;00m\n\u001b[1;32m 229\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39misinstance\u001b[39m(model, pl\u001b[39m.\u001b[39mLightningModule)\n\u001b[0;32m--> 230\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mprecision_plugin\u001b[39m.\u001b[39;49moptimizer_step(optimizer, model\u001b[39m=\u001b[39;49mmodel, closure\u001b[39m=\u001b[39;49mclosure, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py:117\u001b[0m, in \u001b[0;36mPrecision.optimizer_step\u001b[0;34m(self, optimizer, model, closure, **kwargs)\u001b[0m\n\u001b[1;32m 115\u001b[0m \u001b[39m\u001b[39m\u001b[39m\"\"\"Hook to run the optimizer step.\"\"\"\u001b[39;00m\n\u001b[1;32m 116\u001b[0m closure \u001b[39m=\u001b[39m partial(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_wrap_closure, model, optimizer, closure)\n\u001b[0;32m--> 117\u001b[0m \u001b[39mreturn\u001b[39;00m optimizer\u001b[39m.\u001b[39;49mstep(closure\u001b[39m=\u001b[39;49mclosure, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/optim/optimizer.py:373\u001b[0m, in \u001b[0;36mOptimizer.profile_hook_step.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 368\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 369\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[1;32m 370\u001b[0m \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00mfunc\u001b[39m}\u001b[39;00m\u001b[39m must return None or a tuple of (new_args, new_kwargs), but got \u001b[39m\u001b[39m{\u001b[39;00mresult\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 371\u001b[0m )\n\u001b[0;32m--> 373\u001b[0m out \u001b[39m=\u001b[39m func(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 374\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_optimizer_step_code()\n\u001b[1;32m 376\u001b[0m \u001b[39m# call optimizer step post hooks\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 112\u001b[0m \u001b[39m@functools\u001b[39m\u001b[39m.\u001b[39mwraps(func)\n\u001b[1;32m 113\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdecorate_context\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m 114\u001b[0m \u001b[39mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/pytorch_optimizer/optimizer/ranger21.py:181\u001b[0m, in \u001b[0;36mRanger21.step\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m 179\u001b[0m \u001b[39mif\u001b[39;00m closure \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 180\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39menable_grad():\n\u001b[0;32m--> 181\u001b[0m loss \u001b[39m=\u001b[39m closure()\n\u001b[1;32m 183\u001b[0m param_size: \u001b[39mint\u001b[39m \u001b[39m=\u001b[39m \u001b[39m0\u001b[39m\n\u001b[1;32m 184\u001b[0m variance_ma_sum: \u001b[39mfloat\u001b[39m \u001b[39m=\u001b[39m \u001b[39m1.0\u001b[39m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py:104\u001b[0m, in \u001b[0;36mPrecision._wrap_closure\u001b[0;34m(self, model, optimizer, closure)\u001b[0m\n\u001b[1;32m 91\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_wrap_closure\u001b[39m(\n\u001b[1;32m 92\u001b[0m \u001b[39mself\u001b[39m,\n\u001b[1;32m 93\u001b[0m model: \u001b[39m\"\u001b[39m\u001b[39mpl.LightningModule\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m 94\u001b[0m optimizer: Optimizer,\n\u001b[1;32m 95\u001b[0m closure: Callable[[], Any],\n\u001b[1;32m 96\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Any:\n\u001b[1;32m 97\u001b[0m \u001b[39m \u001b[39m\u001b[39m\"\"\"This double-closure allows makes sure the ``closure`` is executed before the ``on_before_optimizer_step``\u001b[39;00m\n\u001b[1;32m 98\u001b[0m \u001b[39m hook is called.\u001b[39;00m\n\u001b[1;32m 99\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 102\u001b[0m \n\u001b[1;32m 103\u001b[0m \u001b[39m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 104\u001b[0m closure_result \u001b[39m=\u001b[39m closure()\n\u001b[1;32m 105\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_after_closure(model, optimizer)\n\u001b[1;32m 106\u001b[0m \u001b[39mreturn\u001b[39;00m closure_result\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py:140\u001b[0m, in \u001b[0;36mClosure.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 139\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__call__\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39m*\u001b[39margs: Any, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs: Any) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Optional[Tensor]:\n\u001b[0;32m--> 140\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mclosure(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 141\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_result\u001b[39m.\u001b[39mloss\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 112\u001b[0m \u001b[39m@functools\u001b[39m\u001b[39m.\u001b[39mwraps(func)\n\u001b[1;32m 113\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdecorate_context\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m 114\u001b[0m \u001b[39mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py:126\u001b[0m, in \u001b[0;36mClosure.closure\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 124\u001b[0m \u001b[39m@torch\u001b[39m\u001b[39m.\u001b[39menable_grad()\n\u001b[1;32m 125\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mclosure\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39m*\u001b[39margs: Any, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs: Any) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m ClosureResult:\n\u001b[0;32m--> 126\u001b[0m step_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_step_fn()\n\u001b[1;32m 128\u001b[0m \u001b[39mif\u001b[39;00m step_output\u001b[39m.\u001b[39mclosure_loss \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 129\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mwarning_cache\u001b[39m.\u001b[39mwarn(\u001b[39m\"\u001b[39m\u001b[39m`training_step` returned `None`. If this was on purpose, ignore this warning...\u001b[39m\u001b[39m\"\u001b[39m)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py:315\u001b[0m, in \u001b[0;36m_AutomaticOptimization._training_step\u001b[0;34m(self, kwargs)\u001b[0m\n\u001b[1;32m 312\u001b[0m trainer \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtrainer\n\u001b[1;32m 314\u001b[0m \u001b[39m# manually capture logged metrics\u001b[39;00m\n\u001b[0;32m--> 315\u001b[0m training_step_output \u001b[39m=\u001b[39m call\u001b[39m.\u001b[39;49m_call_strategy_hook(trainer, \u001b[39m\"\u001b[39;49m\u001b[39mtraining_step\u001b[39;49m\u001b[39m\"\u001b[39;49m, \u001b[39m*\u001b[39;49mkwargs\u001b[39m.\u001b[39;49mvalues())\n\u001b[1;32m 316\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtrainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mpost_training_step() \u001b[39m# unused hook - call anyway for backward compatibility\u001b[39;00m\n\u001b[1;32m 318\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39moutput_result_cls\u001b[39m.\u001b[39mfrom_training_step_output(training_step_output, trainer\u001b[39m.\u001b[39maccumulate_grad_batches)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py:309\u001b[0m, in \u001b[0;36m_call_strategy_hook\u001b[0;34m(trainer, hook_name, *args, **kwargs)\u001b[0m\n\u001b[1;32m 306\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m 308\u001b[0m \u001b[39mwith\u001b[39;00m trainer\u001b[39m.\u001b[39mprofiler\u001b[39m.\u001b[39mprofile(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m[Strategy]\u001b[39m\u001b[39m{\u001b[39;00mtrainer\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39m\u001b[39m__class__\u001b[39m\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m{\u001b[39;00mhook_name\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m):\n\u001b[0;32m--> 309\u001b[0m output \u001b[39m=\u001b[39m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 311\u001b[0m \u001b[39m# restore current_fx when nested context\u001b[39;00m\n\u001b[1;32m 312\u001b[0m pl_module\u001b[39m.\u001b[39m_current_fx_name \u001b[39m=\u001b[39m prev_fx_name\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py:382\u001b[0m, in \u001b[0;36mStrategy.training_step\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 380\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel \u001b[39m!=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlightning_module:\n\u001b[1;32m 381\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_redirection(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlightning_module, \u001b[39m\"\u001b[39m\u001b[39mtraining_step\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m--> 382\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mlightning_module\u001b[39m.\u001b[39;49mtraining_step(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb Cell 26\u001b[0m line \u001b[0;36m5\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=54'>55</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mtraining_step\u001b[39m(\u001b[39mself\u001b[39m, batch, batch_idx\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m, dataloader_idx\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m):\n\u001b[0;32m---> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=55'>56</a>\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_step(batch, batch_idx)\n",
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb Cell 26\u001b[0m line \u001b[0;36m3\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=29'>30</a>\u001b[0m \u001b[39mdel\u001b[39;00m out\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=31'>32</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel\u001b[39m.\u001b[39menable_adapters()\n\u001b[0;32m---> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=32'>33</a>\u001b[0m out2 \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m(batch)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=33'>34</a>\u001b[0m log_probs2 \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mlog_softmax(out2[\u001b[39m'\u001b[39m\u001b[39mlogits\u001b[39m\u001b[39m'\u001b[39m][:, \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m,], \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=35'>36</a>\u001b[0m \u001b[39mif\u001b[39;00m stage \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mpred\u001b[39m\u001b[39m\"\u001b[39m:\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb Cell 26\u001b[0m line \u001b[0;36m1\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=12'>13</a>\u001b[0m b_in \u001b[39m=\u001b[39m \u001b[39mdict\u001b[39m(input_ids\u001b[39m=\u001b[39mbatch[\u001b[39m'\u001b[39m\u001b[39minput_ids\u001b[39m\u001b[39m'\u001b[39m],\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=13'>14</a>\u001b[0m attention_mask\u001b[39m=\u001b[39mbatch[\u001b[39m'\u001b[39m\u001b[39mattention_mask\u001b[39m\u001b[39m'\u001b[39m],)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=14'>15</a>\u001b[0m \u001b[39m# b_in = {k: v.to(self.model.device) for k, v in b_in.items()}\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=16'>17</a>\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel(\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=17'>18</a>\u001b[0m \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mb_in,\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=18'>19</a>\u001b[0m use_cache\u001b[39m=\u001b[39;49m\u001b[39mFalse\u001b[39;49;00m,\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=19'>20</a>\u001b[0m output_hidden_states\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=20'>21</a>\u001b[0m return_dict\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/01_mjc.ipynb#X24sZmlsZQ%3D%3D?line=21'>22</a>\u001b[0m )\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n",
"File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-w_hidden_states/e86b8e062f3e002d3c03182e0122d6fbe5e37039/modeling_phi.py:957\u001b[0m, in \u001b[0;36mPhiForCausalLM.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, labels, **kwargs)\u001b[0m\n\u001b[1;32m 949\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\n\u001b[1;32m 950\u001b[0m \u001b[39mself\u001b[39m,\n\u001b[1;32m 951\u001b[0m input_ids: torch\u001b[39m.\u001b[39mLongTensor,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 955\u001b[0m \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs,\n\u001b[1;32m 956\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m CausalLMOutputWithPast:\n\u001b[0;32m--> 957\u001b[0m outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtransformer(input_ids, past_key_values\u001b[39m=\u001b[39;49mpast_key_values, attention_mask\u001b[39m=\u001b[39;49mattention_mask)\n\u001b[1;32m 958\u001b[0m lm_logits \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlm_head(outputs\u001b[39m.\u001b[39mlast_hidden_state)\n\u001b[1;32m 960\u001b[0m loss \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n",
"File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-w_hidden_states/e86b8e062f3e002d3c03182e0122d6fbe5e37039/modeling_phi.py:917\u001b[0m, in \u001b[0;36mPhiModel.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask)\u001b[0m\n\u001b[1;32m 915\u001b[0m all_hidden_states \u001b[39m=\u001b[39m [hidden_states]\n\u001b[1;32m 916\u001b[0m \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mh:\n\u001b[0;32m--> 917\u001b[0m hidden_states, attn_outputs \u001b[39m=\u001b[39m layer(\n\u001b[1;32m 918\u001b[0m hidden_states,\n\u001b[1;32m 919\u001b[0m past_key_values\u001b[39m=\u001b[39;49mpast_key_values,\n\u001b[1;32m 920\u001b[0m attention_mask\u001b[39m=\u001b[39;49mattention_mask,\n\u001b[1;32m 921\u001b[0m )\n\u001b[1;32m 922\u001b[0m all_hidden_states\u001b[39m.\u001b[39mappend(hidden_states)\n\u001b[1;32m 923\u001b[0m all_self_attns\u001b[39m.\u001b[39mappend(attn_outputs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n",
"File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-w_hidden_states/e86b8e062f3e002d3c03182e0122d6fbe5e37039/modeling_phi.py:779\u001b[0m, in \u001b[0;36mParallelBlock.forward\u001b[0;34m(self, hidden_states, past_key_values, attention_mask, **kwargs)\u001b[0m\n\u001b[1;32m 776\u001b[0m attn_outputs \u001b[39m=\u001b[39m attn_outputs[\u001b[39m0\u001b[39m]\n\u001b[1;32m 778\u001b[0m attn_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresid_dropout(attn_outputs)\n\u001b[0;32m--> 779\u001b[0m feed_forward_hidden_states \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresid_dropout(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmlp(hidden_states))\n\u001b[1;32m 781\u001b[0m hidden_states \u001b[39m=\u001b[39m attn_outputs \u001b[39m+\u001b[39m feed_forward_hidden_states \u001b[39m+\u001b[39m residual\n\u001b[1;32m 783\u001b[0m \u001b[39mreturn\u001b[39;00m hidden_states, attn_outputs\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n",
"File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/wassname/phi-2-w_hidden_states/e86b8e062f3e002d3c03182e0122d6fbe5e37039/modeling_phi.py:322\u001b[0m, in \u001b[0;36mMLP.forward\u001b[0;34m(self, hidden_states)\u001b[0m\n\u001b[1;32m 320\u001b[0m hidden_states \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfc1(hidden_states)\n\u001b[1;32m 321\u001b[0m hidden_states \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mact(hidden_states)\n\u001b[0;32m--> 322\u001b[0m hidden_states \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mfc2(hidden_states)\n\u001b[1;32m 324\u001b[0m \u001b[39mreturn\u001b[39;00m hidden_states\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/peft/tuners/lora/bnb.py:288\u001b[0m, in \u001b[0;36mLinear4bit.forward\u001b[0;34m(self, x, *args, **kwargs)\u001b[0m\n\u001b[1;32m 286\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbase_layer(x, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 287\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 288\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbase_layer(x, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 289\u001b[0m \u001b[39m# As per Tim Dettmers, for 4bit, we need to defensively clone here.\u001b[39;00m\n\u001b[1;32m 290\u001b[0m \u001b[39m# The reason is that in some cases, an error can occur that backprop\u001b[39;00m\n\u001b[1;32m 291\u001b[0m \u001b[39m# does not work on a manipulated view. This issue may be solved with\u001b[39;00m\n\u001b[1;32m 292\u001b[0m \u001b[39m# newer PyTorch versions but this would need extensive testing to be\u001b[39;00m\n\u001b[1;32m 293\u001b[0m \u001b[39m# sure.\u001b[39;00m\n\u001b[1;32m 294\u001b[0m result \u001b[39m=\u001b[39m result\u001b[39m.\u001b[39mclone()\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1525\u001b[0m \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1526\u001b[0m \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 1530\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 163\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[1;32m 164\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m output \u001b[39m=\u001b[39m old_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m 166\u001b[0m \u001b[39mreturn\u001b[39;00m module\u001b[39m.\u001b[39m_hf_hook\u001b[39m.\u001b[39mpost_forward(module, output)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/bitsandbytes/nn/modules.py:256\u001b[0m, in \u001b[0;36mLinear4bit.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 253\u001b[0m x \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mto(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcompute_dtype)\n\u001b[1;32m 255\u001b[0m bias \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbias \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mbias\u001b[39m.\u001b[39mto(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcompute_dtype)\n\u001b[0;32m--> 256\u001b[0m out \u001b[39m=\u001b[39m bnb\u001b[39m.\u001b[39;49mmatmul_4bit(x, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight\u001b[39m.\u001b[39;49mt(), bias\u001b[39m=\u001b[39;49mbias, quant_state\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight\u001b[39m.\u001b[39;49mquant_state)\n\u001b[1;32m 258\u001b[0m out \u001b[39m=\u001b[39m out\u001b[39m.\u001b[39mto(inp_dtype)\n\u001b[1;32m 260\u001b[0m \u001b[39mreturn\u001b[39;00m out\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:577\u001b[0m, in \u001b[0;36mmatmul_4bit\u001b[0;34m(A, B, quant_state, out, bias)\u001b[0m\n\u001b[1;32m 575\u001b[0m \u001b[39mreturn\u001b[39;00m out\n\u001b[1;32m 576\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 577\u001b[0m \u001b[39mreturn\u001b[39;00m MatMul4Bit\u001b[39m.\u001b[39;49mapply(A, B, out, bias, quant_state)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/torch/autograd/function.py:539\u001b[0m, in \u001b[0;36mFunction.apply\u001b[0;34m(cls, *args, **kwargs)\u001b[0m\n\u001b[1;32m 536\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m torch\u001b[39m.\u001b[39m_C\u001b[39m.\u001b[39m_are_functorch_transforms_active():\n\u001b[1;32m 537\u001b[0m \u001b[39m# See NOTE: [functorch vjp and autograd interaction]\u001b[39;00m\n\u001b[1;32m 538\u001b[0m args \u001b[39m=\u001b[39m _functorch\u001b[39m.\u001b[39mutils\u001b[39m.\u001b[39munwrap_dead_wrappers(args)\n\u001b[0;32m--> 539\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39msuper\u001b[39;49m()\u001b[39m.\u001b[39;49mapply(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs) \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 541\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mcls\u001b[39m\u001b[39m.\u001b[39msetup_context \u001b[39m==\u001b[39m _SingleLevelFunction\u001b[39m.\u001b[39msetup_context:\n\u001b[1;32m 542\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[1;32m 543\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mIn order to use an autograd.Function with functorch transforms \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 544\u001b[0m \u001b[39m\"\u001b[39m\u001b[39m(vmap, grad, jvp, jacrev, ...), it must override the setup_context \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 545\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mstaticmethod. For more details, please see \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 546\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mhttps://pytorch.org/docs/master/notes/extending.func.html\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 547\u001b[0m )\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/.venv/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:516\u001b[0m, in \u001b[0;36mMatMul4Bit.forward\u001b[0;34m(ctx, A, B, out, bias, quant_state)\u001b[0m\n\u001b[1;32m 511\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39mempty(A\u001b[39m.\u001b[39mshape[:\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m] \u001b[39m+\u001b[39m B_shape[:\u001b[39m1\u001b[39m], dtype\u001b[39m=\u001b[39mA\u001b[39m.\u001b[39mdtype, device\u001b[39m=\u001b[39mA\u001b[39m.\u001b[39mdevice)\n\u001b[1;32m 514\u001b[0m \u001b[39m# 1. Dequantize\u001b[39;00m\n\u001b[1;32m 515\u001b[0m \u001b[39m# 2. MatmulnN\u001b[39;00m\n\u001b[0;32m--> 516\u001b[0m output \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mnn\u001b[39m.\u001b[39mfunctional\u001b[39m.\u001b[39mlinear(A, F\u001b[39m.\u001b[39;49mdequantize_4bit(B, quant_state)\u001b[39m.\u001b[39;49mto(A\u001b[39m.\u001b[39;49mdtype)\u001b[39m.\u001b[39mt(), bias)\n\u001b[1;32m 518\u001b[0m \u001b[39m# 3. Save state\u001b[39;00m\n\u001b[1;32m 519\u001b[0m ctx\u001b[39m.\u001b[39mstate \u001b[39m=\u001b[39m quant_state\n",
"\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 50.00 MiB. GPU 0 has a total capacty of 23.67 GiB of which 1.13 GiB is free. Including non-PyTorch memory, this process has 21.46 GiB memory in use. Of the allocated memory 20.92 GiB is allocated by PyTorch, and 235.11 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
]
}
],
"source": [
"\n",
"trainer1.fit(model=net, train_dataloaders=dl_train, val_dataloaders=dl_val);\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"df_hist = read_metrics_csv(trainer1.logger.experiment.metrics_file_path).ffill().bfill()\n",
"for key in [\"loss_rec\"]:\n",
" df_hist[[c for c in df_hist.columns if key in c]].plot(logy=True)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"# predict\n",
"dl_test = dm.test_dataloader()\n",
"# print(f\"training with x_feats={x_feats} with c={c}\")\n",
"rs = trainer2.test(net, dataloaders=[dl_train, dl_val, dl_test, dl_oos])\n",
"\n",
"testval_metrics = calc_metrics(dm, trainer2, net, use_val=True)\n",
"rs = rename(rs, [\"train\", \"val\", \"test\", \"oos\"])\n",
"# rs['test'] = {**rs['test'], **test_metrics}\n",
"rs[\"test\"][\"acc_lie_lie\"] = testval_metrics[\"acc_lie_lie\"]\n",
"rs[\"testval_metrics\"] = rs[\"test\"]\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python"
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.0rc1"
}
},
"nbformat": 4,
poetry.lock
Generated
+145 -1350
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+4 -5
View File
@@ -8,13 +8,12 @@ readme = "README.md"
[tool.poetry.dependencies]
python = ">=3.10,<3.13"
accelerate = "^0.23.0"
accelerate = "0.23.0"
torch = {version = "^2.1.0+cu118", source = "pytorch"}
simple-parsing = "^0.1.4"
tqdm = "^4.66.1"
datasets = "^2.14.5"
transformers = "^4.36.1"
auto-gptq = "^0.6.0"
transformers = "4.34.0"
optimum = "^1.13.2"
numpy = "^1.26.1"
pandas = "^2.1.1"
@@ -22,15 +21,15 @@ lightning = "^2.1.0"
matplotlib = "^3.8.0"
loguru = "^0.7.2"
einops = "^0.7.0"
baukit = {git = "https://github.com/davidbau/baukit.git"}
eleuther-elk = "0.1.1"
scikit-learn = "^1.3.1"
pytorch-optimizer = "^2.12.0"
pathvalidate = "^3.2.0"
torchinfo = "^1.8.0"
jaxtyping = "^0.2.24"
autoawq = "^0.1.7"
bitsandbytes = "^0.41.3.post2"
packaging = "^23.2"
peft = "^0.7.1"
[[tool.poetry.source]]
name = "pytorch"
research_log.md
+53
View File
@@ -0,0 +1,53 @@
# 2023-12-17 20:04:28
The idea is to use Lora as a probe. Only activated on the biases of the last layers or something like that.
This way we easily use backprop to chose the probe.
Then we test it for lie detection on a new dataset to check acc and generalizaion.
# 2023-12-17 20:05:09
Examples of peft
- jdp https://github.com/JD-P/minihf/blob/b54075c34ef88d9550e37fdf709e78e5a68787c4/lora_tune.py#L50
The idea is that we do something like this
```py
with torch.no_rgrad():
probs = model(input)
choice_probs = probs[choices]
opposite_probs = choice_probs[::-1].detach()
with adapter('intervention'):
probs2 = model(input)
choice_probs2 = probs2[choices]
# or do we keep everything else the same?
loss = kl_div(choice_probs2, opposite_probs)
loss.backward()
optimizer.step()
```
# 2023-12-18 07:30:23
I've having device problems. I probobly need to load a non quant version.
can;t install flash mem
can't train in 8bit mode?
bitsandbytes = "^0.41.3.post2" ?
and if I use 16 bit and my 25g gpu it's out of mem. I can't use a batch of one or lightnign fauils
- ... so don't use liughtning
- or use Loraq https://github.com/huggingface/peft/blob/46a84bd395f1b486b7b076acfa8f6df3dfad26b8/examples/loftq_finetuning/README.md?plain=1#L2
- or work out how to use butsanybytes during train
- or gradient accum?
src/config.py
+7 -1
View File
@@ -14,7 +14,7 @@ class ExtractConfig(Serializable):
datasets: tuple[str, ...] = ("amazon_polarity", "super_glue:boolq", "glue:qnli", "imdb")
"""Names of HF datasets to use, e.g. `"super_glue:boolq"` or `"imdb"` `"glue:qnli"""
model: str = "wassname/phi-2-GPTQ_w_hidden_states"
model: str = "wassname/phi-2-w_hidden_states"
batch_size: int = 5
@@ -31,3 +31,9 @@ class ExtractConfig(Serializable):
intervention_fit_examples: int = 200
"""how many example to use for intervention calibration"""
max_examples: tuple[int, int] = (1000, 200)
"""Maximum number of examples to use from each split of the dataset."""
seed: int = 42
"""Random seed."""
src/datasets/__init__.py
View File
src/datasets/scores.py
+100
View File
@@ -0,0 +1,100 @@
from typing import Optional, List, Tuple, Dict
import numpy as np
import torch
import torch.nn.functional as F
import functools
import itertools
from transformers import (
PreTrainedTokenizer,
PreTrainedModel
)
default_class2choices = [['No', 'Negative', 'negative', 'no', 'false', 'wrong', 'False', '0'], ['Yes', 'Positive', 'positive', 'yes', 'true', 'correct', 'right', 'True', '1']]
def scores2choice_probs(row, class2_ids: List[List[int]], keys=["scores0"], prefix=""):
""" Given next_token scores (logits) we take only the subset the corresponds to our
- negative tokens (e.g. False, no, ...)
- and positive tokens (e.g. Yes, yes, affirmative, ...).
example output:
{'choice_probs1': array([0.39, 0.31 ], dtype=float32),
'ans1': 0.44,
'choice_probs2': array([0.44, 0.45], dtype=float32),
'ans2': 0.502,}
"""
eps = 1e-5
out = {}
for key in keys:
scores = torch.tensor(row[key])
probs = F.softmax(scores, 0).numpy() # shape [tokens, inferences)
probs_c = [np.sum([probs[cc] for cc in c], 0) for c in class2_ids] # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']]
probs_c = np.stack(probs_c, 0) # shape [choices, inferences]
# balance of probs
out[prefix+key.replace("scores", "choice_probs")] = probs_c
out[prefix+key.replace("scores", "ans")] = probs_c[1] / (np.sum(probs_c, 0) + eps) # shape is [inferences]
# # balance of logits (much more exaggerated)
# scores_c = [scores[class2_ids[c]].sum() for c in class2_ids]
# out[key.replace("scores", "ansb")] = torch.tensor(scores_c).softmax(-1)[1].item()
return out
@functools.lru_cache()
def choice2id(tokenizer, c: str, whitespace_first=False) -> List[int]:
"""convert a choice to a single token"""
# HACK: this whole function is messy, and specific to the llama tokenizer :(. I don't want it to fail silently, so I'm adding a few asserts. It's better to find out before 4 hours of data collection
# Note some tokenizers differentiate between "yes", "\nyes" and " yes", and ideally we want all!
ids2 = []
ids2 += tokenizer(f' {c}', add_special_tokens=False)["input_ids"]
ids2 += tokenizer(f'\n{c}', add_special_tokens=False)["input_ids"]
ids2 += tokenizer(f'{c}', add_special_tokens=False)["input_ids"]
ids = list(set(ids2))
# print(ids2)
# print(ids)
# print([f'`{t}`' for t in tokenizer.batch_decode(ids, skip_special_tokens=True)])
# print([c.strip().startswith(tokenizer.decode(i)) for i in ids])
# only include ones that decode to our original
ids = [i for i in ids if c.strip().startswith(tokenizer.decode(i).strip()) and len(tokenizer.decode(i).strip())]
assert len(ids)
# QC: they should all decode to the same token
decoded_ids = tokenizer.batch_decode(ids)
shortest = sorted(decoded_ids, key=lambda s:len(s))[0]
assert len(shortest)
assert all([decoded_ids[i].strip().startswith(shortest) for i in range(len(decoded_ids))]), f"decoded_ids={decoded_ids}"
# check that we can decode it
c3 = tokenizer.batch_decode(ids)
for c2 in c3:
if not c.strip().startswith(c2.strip()) and len(c2):
print(c, c2, c3)
ids = tokenizer(c, add_special_tokens=False)["input_ids"]
decoded_ids = [tokenizer.decode(i).strip() for i in ids]
print(f"{c}=>{ids}=>{decoded_ids}")
raise AssertionError(f'We should be able to encode and decode the choices, but it failed: tokenizer.decode(tokenizer(`{c}`))==`{c2}`!=`{c}`')
return ids
def choice2ids(all_choices: List[List[str]], tokenizer: PreTrainedTokenizer) -> List[List[int]]:
choices = [list(itertools.chain(*[choice2id(tokenizer, c) for c in choices])) for choices in all_choices]
assert choices[0]!=choices[1], f"choices should be different but were not {all_choices}"
assert choices[0][0]!=choices[1][0], "choices should be different"
return choices
def logits2choice_probs2(logits, choiceids: List[List[int]]):
"""calculate the probability for each group of choices."""
assert logits.ndim==1, f"expected logits to be 1d, got {logits.shape}"
assert logits.sum(0).abs()>2, 'pass in logits, not probs. you may have accidentally passed in a softmaxed values'
choiceids = [list(set(i)) for i in choiceids] # we don't want to double count
probs = torch.softmax(logits, 0) # shape [tokens, inferences)
probs_c = torch.tensor([[probs[cc] for cc in c] for c in choiceids]).sum(1) # sum over alternate choices e.g. [['decrease', 'dec'],['inc', 'increase']]
assert probs_c.sum()<=1.01, f"expected probs to sum to 1, got {probs_c.sum()}"
return probs_c
def row_choice_ids(r, tokenizer):
return choice2ids([c for c in r['answer_choices']], tokenizer)
src/helpers/__init__.py
+8
View File
@@ -0,0 +1,8 @@
def bool2switch(x):
"""[0,1]->[-1,1]"""
return x*2-1
def switch2bool(x):
"""[-1,1]->[0,1]"""
return (x+1)/2
src/helpers/ds.py
+26
View File
@@ -0,0 +1,26 @@
import gc
import torch
from datasets import Dataset
import numpy as np
def ds_keep_cols(ds: Dataset, cols: list) -> Dataset:
cols_all = set(ds.features.keys())
cols_drop = cols_all - set(cols)
return ds.remove_columns(cols_drop)
def clear_mem():
gc.collect()
torch.cuda.empty_cache()
gc.collect()
def shuffle_dataset_by(ds, column):
# ds_tokens = ds.filter(lambda r: r["truncated"] == False)
example_i = np.array(ds["example_i"])
uniq_example_i = np.array(sorted(set(example_i)))
shuffled_indices = np.random.permutation(uniq_example_i)
index = np.arange(len(example_i))
new_inds = np.concatenate([index[example_i == i] for i in shuffled_indices])
return ds.select(new_inds)
src/helpers/lightning.py
+10
View File
@@ -0,0 +1,10 @@
from lightning.pytorch.loggers.csv_logs import CSVLogger
from pathlib import Path
import pandas as pd
def read_metrics_csv(metrics_file_path):
df_hist = pd.read_csv(metrics_file_path)
df_hist["epoch"] = df_hist["epoch"].ffill()
df_histe = df_hist.set_index("epoch").groupby("epoch").mean()
return df_histe
src/helpers/torch.py
+59
View File
@@ -0,0 +1,59 @@
import torch
import numpy as np
import transformers
import random
import gc
import pandas as pd
def get_top_n(scores: torch.Tensor, tokenizer: transformers.PreTrainedTokenizer, n=10) -> pd.Series:
"""Get top n choices and their probabilities given raw logits"""
probs = scores.softmax(-1).squeeze()
assert len(probs.shape)==1
top10 = torch.argsort(probs, dim=-1, descending=True)[:n]
top10_probs = probs[top10]
top10_ext = tokenizer.batch_decode(top10)
return pd.Series(top10_probs, index=top10_ext, name='probs')
def to_numpy(x):
"""
Trys to convert torch to numpy and if possible a single item
"""
if isinstance(x, torch.Tensor):
# note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981
x = x.detach().cpu().float()
if x.squeeze().dim()==0:
return x.item()
return x.numpy()
else:
return x
def set_seeds(n: int) -> None:
transformers.set_seed(n)
torch.manual_seed(n)
np.random.seed(n)
random.seed(n)
def to_item(x):
if isinstance(x, torch.Tensor):
x = x.detach().cpu().item()
return x
def clear_mem():
gc.collect()
torch.cuda.empty_cache()
gc.collect()
def detachcpu(x):
"""
Trys to convert torch if possible a single item
"""
if isinstance(x, torch.Tensor):
# note apache parquet doesn't support half to we go for float https://github.com/huggingface/datasets/issues/4981
x = x.detach().cpu()
if x.squeeze().dim()==0:
return x.item()
return x
else:
return x
src/helpers/typing.py
+17
View File
@@ -0,0 +1,17 @@
from typing import Any, Type, TypeVar, cast
import torch
def float_to_int16(x: torch.Tensor) -> torch.Tensor:
"""Converts a floating point tensor to float16, then reinterprets as int16."""
downcast = x.type(torch.float16)
if not downcast.isfinite().all():
raise ValueError("Cannot convert to 16 bit: values are not finite")
return downcast.view(torch.int16)
def int16_to_float(x: torch.Tensor) -> torch.Tensor:
"""Converts int16 to float16, then reinterprets as float32."""
return x.view(torch.float16).type(torch.float32)
src/models/load.py
+19 -11
View File
@@ -5,13 +5,10 @@ When editing or updating this file check out these resources:
- [LLM-As-Chatbot](https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py)
- [oobabooga](https://github.com/oobabooga/text-generation-webui/blob/main/modules/models.py#L134)
"""
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig, PreTrainedTokenizerBase, PreTrainedTokenizer, GPTQConfig
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig, PreTrainedTokenizerBase, PreTrainedTokenizer, GPTQConfig, BitsAndBytesConfig
import torch
from zmq import has
from src.datasets.dropout import check_for_dropout
from loguru import logger
from typing import Tuple
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
def verbose_change_param(tokenizer, path, after):
@@ -25,7 +22,7 @@ def verbose_change_param(tokenizer, path, after):
return tokenizer
def load_model(model_repo = "microsoft/phi-2", pad_token_id=0, disable_exllama=True) -> Tuple[AutoModelForCausalLM, PreTrainedTokenizerBase]:
def load_model(model_repo = "microsoft/phi-2", pad_token_id=0, disable_exllama=True, device="auto",) -> Tuple[AutoModelForCausalLM, PreTrainedTokenizerBase]:
"""
A uncensored and large coding ones might be best for lying.
@@ -33,12 +30,23 @@ def load_model(model_repo = "microsoft/phi-2", pad_token_id=0, disable_exllama=
# see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py
# gptq_config = GPTQConfig(bits=4, dataset="c4", disable_exllama=False)
model_options = dict(
device_map="auto",
device_map=device,
torch_dtype=torch.float16,
# load_in_8bit=True,
# load_in_8bit=True, # difficult
load_in_4bit=True,
trust_remote_code=True,
# disable_exllama=True,
# flash_attn=True, flash_rotary=True, fused_dense=True
quantization_config=BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16, # bfloat16 is recommended
bnb_4bit_use_double_quant=False,
bnb_4bit_quant_type='nf4',
),
## in the azure phi-repo they use these but you need to install flashattn
# flash_attn=True,
# flash_rotary=True,
# fused_dense=True
)
config = AutoConfig.from_pretrained(model_repo, trust_remote_code=True,)
@@ -58,9 +66,9 @@ def load_model(model_repo = "microsoft/phi-2", pad_token_id=0, disable_exllama=
model = AutoModelForCausalLM.from_pretrained(model_repo, config=config,
**model_options)
if disable_exllama:
from auto_gptq import exllama_set_max_input_length
model = exllama_set_max_input_length(model, max_input_length=5000)
# if disable_exllama:
# from auto_gptq import exllama_set_max_input_length
# model = exllama_set_max_input_length(model, max_input_length=5000)
return model, tokenizer
src/prompts/prompt_loading.py
+1 -1
View File
@@ -26,7 +26,7 @@ import pandas as pd
from loguru import logger
from src.helpers.ds import shuffle_dataset_by
from src.datasets.scores import choice2id, choice2ids, row_choice_ids
from src.extraction.config import ExtractConfig
from src.config import ExtractConfig
from src.models.load import verbose_change_param, AutoConfig, AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase