misc

mjc_notes.md

@@ -1447,8 +1447,7 @@ python -m pdbp notebooks/011_make_dataset.py \
 "WizardLM/WizardCoder-3B-V1.0" \
  imdb amazon_polarity super_glue:boolq glue:qnli \
  --max_examples 260 260 \
- --max_length=600 \
- --num_shots=1 
+ --max_length=600
 ```
 
 - [ ] run this exp
@@ -1497,3 +1496,46 @@ outputs2 = model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, out
 # return model
 model.load_state_dict(orig_state_dict)
 ```
+
+
+UPTO:
+- I made a dataset py file
+- I made it collect counterfactual inference
+- now I need to see if a probe that has residual and counterfactual residual does better!
+
+
+counterfactuals?
+- **dice** loss:
+  - [x] hidden_State and hidden_stte counterfactual.... just **overfits**.  acc_lie_lie=0.00% from probe 0.82
+  - [x] with grad it also overfits acc_lie_lie=0.00%? :( this is weird as it has no counterfactuals.. .what's going on??
+    - [ ] :poop:  :bug: this does not make sense!?! why is it sudeently overfitting. this invalidated all these experiments
+  - [x] wioth residual it overfits?  :(
+- [ ] ranking?  
+   - [ ] what about counterfactuals with **ranking**?! 0.87.44$ acc_lie_lie. pretty good?
+- [x] what about nonlinear? nope
+
+Next
+- [ ] with noise on embeddings? this would allow large models agian. I'm really struggling with these small models!?
+
+# 2023-09-22 13:36:27
+
+So there must be bugs in my datasets, as
+- only 3 sysprompts!
+- each dataset seems to only have truth or lie... weird
+- imdb has no coverage
+
+# 2023-09-23 13:13:27
+
+Debugging the datasets
+
+Observations
+- Why are there only X prompts?
+- Why are some only lies, others not??!
+
+```py
+# snippets for debugging chosen prompts
+print(pd.Series(ds_tokens['sys_instr_name']).value_counts())
+print(pd.Series(ds_tokens['template_name']).value_counts())
+print(pd.Series(ds_tokens['label_instructed']).value_counts())
+```
+['ds_string', 'example_i', 'answer', 'question', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'input_ids', 'attention_mask', 'truncated', 'prompt_truncated', 'choice_ids'],

notebooks/011_make_dataset.py

@@ -4,19 +4,12 @@
 # 
 
 # %%
-# import your package
-# %load_ext autoreload
-# %autoreload 2
-
 from loguru import logger
 import sys
 logger.remove()
 logger.add(sys.stderr, format="<level>{message}</level>", level="INFO")
 
 import pandas as pd
-# from matplotlib import pyplot as plt
-# %matplotlib inline
-# plt.style.use('ggplot')
 
 # %%
 import numpy as np
@@ -50,8 +43,11 @@ from src.datasets.load import rows_item
 from src.datasets.batch import batch_hidden_states
 # from src.datasets.scores import choice2ids, scores2choice_probs
 
-# %% [markdown]
-# # Params
+from datasets import disable_caching
+disable_caching()
+import psutil
+max_dataset_memory = f"{psutil.virtual_memory().total //2}"
+os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = max_dataset_memory
 
 # %%
 from simple_parsing import ArgumentParser
@@ -70,27 +66,24 @@ parser.add_arguments(ExtractConfig, dest="run")
 
 args = parser.parse_args()
 cfg = args.run
-cfg
+print(cfg)
 
-# %%
-# Params
 BATCH_SIZE = 1  # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15
 
-# %% [markdown]
-# # Model
-# 
-# Chosing:
-# - https://old.reddit.com/r/LocalLLaMA/wiki/models
-# - https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard
-# - https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json
-# 
-# 
-# A uncensored and large coding ones might be best for lying.
 
 # %%
 from src.models.load import verbose_change_param, AutoConfig, AutoTokenizer, AutoModelForCausalLM
 
 def load_model(model_repo = "HuggingFaceH4/starchat-beta"):
+    """
+    Chosing:
+    - https://old.reddit.com/r/LocalLLaMA/wiki/models
+    - https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard
+    - https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json
+
+
+    A uncensored and large coding ones might be best for lying.
+    """
     # see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py
     model_options = dict(
         device_map="auto",
@@ -113,10 +106,6 @@ def load_model(model_repo = "HuggingFaceH4/starchat-beta"):
     return model, tokenizer
 
 
-
-# %% [markdown]
-# # Load Dataset
-
 # %%
 from itertools import chain
 import functools
@@ -188,7 +177,7 @@ ds_tokens = (
         ),
         batched=True,
     )
-    .map(lambda r: {"truncated": np.sum(r["attention_mask"], -1)<cfg.max_length})
+    .map(lambda r: {"truncated": np.sum(r["attention_mask"], -1)==cfg.max_length})
     .map(
         lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
         batched=True,
@@ -200,7 +189,7 @@ ds_tokens
 # %%
 ds_tokens = ds_tokens.filter(lambda r: r['truncated']==False)
 ds_tokens = ds_tokens.select(range(min(len(ds_tokens), N)))
-ds_tokens.num_rows
+print('removed truncated rows to leave: num_rows', ds_tokens.num_rows)
 
 # %% [markdown]
 # ## Save as Huggingface Dataset
@@ -277,15 +266,12 @@ ds3 = (
 )
 ds3
 
-# %%
-ds3.config_name
-
 # %% [markdown]
 # ## Save to disk
 
 # %%
 ds3.save_to_disk(f)
-f
+print('! f=', f)
 
 # %% [markdown]
 # # QC
@@ -481,3 +467,5 @@ df2= ds2df(ds5)
 df_subset_successull_lies = df2.query("instructed_to_lie==True & (llm_ans==label_instructed)")
 print(f"filtered to {len(df_subset_successull_lies)} num successful lies out of {len(df2)} dataset rows")
 assert len(df_subset_successull_lies)>0, "there should be successful lies in the dataset"
+
+print(f)

notebooks/012_make_dataset.py

@@ -0,0 +1,393 @@
+# # Lets save our data as a huggingface dataset, so it's quick to reuse
+
+
+import psutil, os
+max_dataset_memory = f"{psutil.virtual_memory().total //2}"
+os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = max_dataset_memory
+
+from datasets import disable_caching
+disable_caching()
+
+from loguru import logger
+import sys
+logger.remove()
+logger.add(sys.stderr, format="<level>{message}</level>", level="INFO")
+
+import pandas as pd
+
+import numpy as np
+
+
+from typing import Optional, List, Dict, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+import pickle
+import hashlib
+from pathlib import Path
+
+import transformers
+from datasets import Dataset, DatasetInfo
+from src.datasets.load import load_ds
+from src.models.load import verbose_change_param, AutoConfig, AutoTokenizer, AutoModelForCausalLM
+from tqdm.auto import tqdm
+import os, re, sys, collections, functools, itertools, json
+
+from src.models.load import load_model
+from src.datasets.load import ds2df
+from src.datasets.load import rows_item
+from src.datasets.batch import batch_hidden_states
+# from src.datasets.scores import choice2ids, scores2choice_probs
+
+from simple_parsing import ArgumentParser
+from src.extraction.config import ExtractConfig
+parser = ArgumentParser(add_help=False)
+parser.add_arguments(ExtractConfig, dest="run")
+
+# argv="""\
+# "WizardLM/WizardCoder-3B-V1.0" \
+# imdb amazon_polarity super_glue:boolq glue:qnli \
+# --max_examples 260 260 \
+# --max_length=600 \
+# --num_shots=1 \
+# """.strip().replace('\n','').split()
+# print(argv)
+
+args = parser.parse_args()
+cfg = args.run
+print(cfg)
+
+BATCH_SIZE = 1  # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15
+
+
+
+def load_model(model_repo = "HuggingFaceH4/starchat-beta"):
+    """
+    Chosing:
+    - https://old.reddit.com/r/LocalLLaMA/wiki/models
+    - https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard
+    - https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json
+
+
+    A uncensored and large coding ones might be best for lying.
+    """
+    # see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py
+    model_options = dict(
+        device_map="auto",
+        # load_in_8bit=True,
+        # load_in_4bit=True,
+        torch_dtype=torch.float16, # note because datasets pickles the model into numpy to get the unique datasets name, and because numpy doesn't support bfloat16, we need to use float16
+        # use_safetensors=False,
+    )
+
+    config = AutoConfig.from_pretrained(model_repo, use_cache=False)
+    verbose_change_param(config, 'use_cache', False)
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_repo)
+    verbose_change_param(tokenizer, 'pad_token_id', 0)
+    verbose_change_param(tokenizer, 'padding_side', 'left')
+    verbose_change_param(tokenizer, 'truncation_side', 'left')
+    
+    model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)
+
+    return model, tokenizer
+
+
+def qc_ds(f):
+    ds4 = load_ds(f)
+
+    # QC: check that the dataset is valid
+    for k,v in ds4[0].items():
+        print(k, v.shape, v.dtype)
+        if (isinstance(v, (np.ndarray, np.generic, torch.Tensor)) and (v.dtype in ['float16', 'float32', 'float64', 'int64', 'int32', 'int16', 'int8'])):
+            assert np.isfinite(v).all()
+
+
+    # QC, check which answers are most common
+    common_answers = pd.Series(ds4['txt_ans0']).value_counts()
+    print('Remember it should be binary. Found common LLM answers:', common_answers)
+
+    current_choices = set(list(chain(*ds4['answer_choices'])))
+    unexpected_answers = set(common_answers.head(10).index)-current_choices
+    if len(unexpected_answers):
+        logger.warning(f'found unexpected answers: {unexpected_answers}. You may want to add them to class2choices')
+    
+    mean_prob = ds4['choice_probs0'].sum(-1).mean()
+    print('mean_prob', mean_prob)
+    assert ds4['choice_probs0'].sum(-1).mean()>0.2, f"""
+    Our choices should cover most common answers. But they accounted for a mean probability of {mean_prob:2.2%} (should be >40%). 
+
+    To fix this you might want to improve your prompt or add to your choices
+    """
+
+    df = ds2df(ds4)
+    print(df.head(5))
+
+    # QC check accuracy
+    # it should manage to lie some of the time when asked to lie. Many models wont lie unless very explicitly asked to, but we don't want to do that, we want to leave some ambiguity in the prompt
+
+    d = df.query('instructed_to_lie==True')
+    acc = (d.label_instructed==d.llm_ans).mean()
+    print(f"when the model tries to lie... we get this acc {acc:2.2f}")
+    assert acc>0.1, f"should be acc>0.1 but is acc={acc}"
+
+    # ### QC stats
+    def stats(df):
+        return dict(
+            acc=(df.llm_ans == df.label_instructed).mean(),
+            n=len(df),
+        )
+        
+    def col2statsdf(df, group):
+        return pd.DataFrame(df.groupby(group).apply(stats).to_dict()).T
+        
+        
+    print("how well does it do the simple task of telling the truth, for each template")
+    col2statsdf(df.query('sys_instr_name=="truth"'), 'template_name')
+
+    print("how well does it complete the task for each prompt")
+    # of course getting it to tell the truth is easy, but how effective are the other prompts?
+    col2statsdf(df, 'sys_instr_name')
+
+    # ### QC view row
+
+    # QC by viewing a row
+    r = ds4[0]
+    print(r['prompt_truncated'])
+    print(r['txt_ans0'])
+
+    # # QC: generation
+    # 
+    # Let's a quick generation, so we can QC the output and sanity check that the model can actually do the task
+
+    # r = ds[2]
+    # q = r["prompt_truncated"]
+
+    # pipeline = transformers.pipeline(
+    #     "text-generation",
+    #     model=model,
+    #     tokenizer=tokenizer,
+    # )
+    # sequences = pipeline(
+    #     q.lstrip('<|endoftext|>'),
+    ##     max_length=100,
+    # max_new_tokens=10,
+    #     do_sample=False,
+    #     return_full_text=False,
+    #     eos_token_id=tokenizer.eos_token_id,
+    # )
+
+    # for seq in sequences:
+    #     print("-" * 80)
+    #     print(q)
+    #     print("-" * 80)
+    #     print(f"`{seq['generated_text']}`")
+    #     print("-" * 80)
+    #     print("label", r['label'])
+
+
+    # # QC: linear probe
+
+    from sklearn.preprocessing import RobustScaler
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.metrics import f1_score, roc_auc_score, accuracy_score
+
+    # # just select the question where the model knows the answer. 
+    df = ds2df(ds4)
+    d = df.query('sys_instr_name=="truth"').set_index("example_i")
+
+    # # these are the ones where it got it right when asked to tell the truth
+    m1 = d.llm_ans==d.label_true
+    known_indices = d[m1].index
+    print(f"select rows are {m1.mean():2.2%} based on knowledge")
+    # # convert to row numbers, and use datasets to select
+    known_rows = df['example_i'].isin(known_indices)
+    known_rows_i = df[known_rows].index
+
+    # # also restrict it to significant permutations. That is monte carlo dropout pairs, where the answer changes by more than X%
+    # m = np.abs(df.ans0-df.ans1)>0.05
+    # print(f"selected rows are {m.mean():2.2%} for significance")
+    # significant_rows = m[m].index
+
+    # allowed_rows_i = set(known_rows_i).intersection(significant_rows)
+    # allowed_rows_i = significant_rows
+    ds5 = ds4.select(known_rows_i)
+    df = ds2df(ds5)
+
+
+    large_arrays_keys = [k for k,v in ds4[0].items() if v.ndim>1]
+
+    for k in large_arrays_keys:
+        print('-'*80)
+        print(k)
+        hs = ds5[k]
+        X = hs.reshape(hs.shape[0], -1)
+
+
+        y = df['label_true'] == df['llm_ans']
+
+        # split
+        n = len(y)
+        max_rows = 1000
+        
+        X_train, X_test = X[:n//2], X[n//2:]
+        y_train, y_test = y[:n//2], y[n//2:]
+        X_train = X_train[:max_rows]
+        y_train = y_train[:max_rows]
+        X_test = X_test[:max_rows]
+        y_test = y_test[:max_rows]
+        print('split size', X_train.shape, y_test.shape)
+
+        # scale
+        scaler = RobustScaler()
+        scaler.fit(X_train)
+        X_train2 = scaler.transform(X_train)
+        X_test2 = scaler.transform(X_test)
+
+        lr = LogisticRegression(class_weight="balanced", penalty="l2", max_iter=380)
+        lr.fit(X_train2, y_train>0)
+
+        print("Logistic cls acc: {: 3.2%} [TRAIN]".format(lr.score(X_train2, y_train>0)))
+        print("Logistic cls acc: {: 3.2%} [TEST]".format(lr.score(X_test2, y_test>0)))
+
+
+    # QC: make sure we didn't lose all of the successful lies, which would make the problem trivial
+    df2= ds2df(ds5)
+    df_subset_successull_lies = df2.query("instructed_to_lie==True & (llm_ans==label_instructed)")
+    print(f"filtered to {len(df_subset_successull_lies)} num successful lies out of {len(df2)} dataset rows")
+    assert len(df_subset_successull_lies)>0, "there should be successful lies in the dataset"
+
+    print(f)
+
+
+
+
+from itertools import chain
+import functools
+from src.prompts.prompt_loading import load_prompts
+from src.datasets.scores import scores2choice_probs
+from src.datasets.scores import choice2id, choice2ids
+
+     
+# TODO: loop through all prompts in this dataset
+ds_names = cfg.datasets
+split_type = "train"
+
+model, tokenizer = load_model(cfg.model)
+model.cuda()
+
+def row_choice_ids(r):
+    return choice2ids([[c] for c in r['answer_choices']], tokenizer)
+
+
+for ds_name in ds_names:
+        
+    # TODO: for when we need custom templates....
+    # ds_root_name, _, subset_name = ds_name.partition(":")
+    # template_path = cfg.template_path/ds_root_name
+    # if subset_name:
+    #     template_path = template_path/subset_name
+    # template_path
+
+    N = cfg.max_examples[split_type!="train"]
+    ds_prompts = Dataset.from_generator(
+        load_prompts, 
+        gen_kwargs=dict(
+            ds_string=ds_name, 
+            num_shots=cfg.num_shots,
+            split_type=split_type,
+            # template_path=template_path,
+            seed=cfg.seed,
+            prompt_format='llama',
+            N=N*3,
+        ), 
+        )
+
+
+    # ## Format prompts
+    # The prompt is the thing we most often have to change and debug. So we do it explicitly here.
+    # We do it as transforms on a huggingface dataset.
+    # In this case we use multishot examples from train, and use the test set to generated the hidden states dataset. We will test generalisation on a whole new dataset.
+
+    ds_tokens = (
+        ds_prompts
+        .map(
+            lambda ex: tokenizer(
+                ex["question"], padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True,
+                return_tensors="np",
+                return_attention_mask=True,
+                # return_overflowing_tokens=True,
+            ),
+            batched=True,
+        )
+        .map(lambda r: {"truncated": np.sum(r["attention_mask"], -1)==cfg.max_length})
+        .map(
+            lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
+            batched=True,
+        )
+        .map(lambda r: {'choice_ids': row_choice_ids(r)})
+    )
+    
+    
+    ds_tokens = ds_tokens.filter(lambda r: r['truncated']==False)
+    ds_tokens = ds_tokens.select(range(min(len(ds_tokens), N)))
+    print('removed truncated rows to leave: num_rows', ds_tokens.num_rows)
+
+    # ## Save as Huggingface Dataset
+    # get dataset filename
+    sanitize = lambda s:s.replace('/', '').replace('-', '_') if s is not None else s
+    dataset_name = f"{sanitize(cfg.model)}_{ds_name}_{split_type}_{N}"
+    f = f"../.ds/{dataset_name}"
+
+    gen_kwargs = dict(
+        model=model,
+        tokenizer=tokenizer,
+        data=ds_tokens,
+        batch_size=BATCH_SIZE,
+        layer_padding=cfg.layer_padding,
+        layer_stride=cfg.layer_stride,
+    )
+
+    info_kwargs = dict(extract_cfg=cfg.to_dict(), ds_name=ds_name, split_type=split_type, f=f, date=pd.Timestamp.now().isoformat(),)
+
+    # [DatasetInfo](https://github.com/huggingface/datasets/blob/9b21e181b642bd55b3ef68c1948bfbcd388136d6/src/datasets/info.py#L94)
+    ds1 = Dataset.from_generator(
+        generator=batch_hidden_states,
+        info=DatasetInfo(
+            description=json.dumps(info_kwargs, indent=2),
+            config_name=f,
+        ),
+        gen_kwargs=gen_kwargs,
+        num_proc=1,
+        
+    )
+
+    # ## Add labels
+    # For our probe. 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, ...).
+
+    # this is just based on pairs for that answer...
+    add_txt_ans0 = lambda r: {'txt_ans0': tokenizer.decode(r['scores0'].argmax(-1))}
+
+    # Either just use the template choices
+    add_ans = lambda r: scores2choice_probs(r, row_choice_ids(r), keys=["scores0"])
+
+    # Or all expanded choices
+    ds1.set_format(type='numpy')#, columns=['input_ids', 'token_type_ids', 'attention_mask', 'label'])
+    ds3 = (
+        ds1
+        .map(add_ans)
+        .map(add_txt_ans0)
+    )
+
+    ds3.save_to_disk(f)
+    print('! saved f=', f)
+    
+    try:
+        qc_ds(f)
+    except Exception as e:
+        print('QC failed', e)
+        # raise e
+

notebooks/102b_scratch_extract_grads_simpler.ipynb

@@ -746,6 +746,48 @@
     "print(f\"loss={l}, pos={score_y2}, neg={score_n2}\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'67078602752'"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import psutil \n",
+    "max_dataset_memory = f\"{psutil.virtual_memory().total}\"\n",
+    "max_dataset_memory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "67.078602752"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "67078602752/1e9"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

src/datasets/dm.py

@@ -8,12 +8,12 @@ from datasets.arrow_dataset import Dataset
 from einops import rearrange, reduce, repeat
 
 
-def compute_distance(df):
-    """distance between ans1 and ans2."""
-    true_switch_sign = df.label_true*2-1 # switch sign to desired answer. with this we ask which is more true
-    # otherwise we ask which is more positive
-    distance = (df.ans1-df.ans0) * true_switch_sign
-    return distance
+# def compute_distance(df):
+#     """distance between ans1 and ans2."""
+#     true_switch_sign = df.label_true*2-1 # switch sign to desired answer. with this we ask which is more true
+#     # otherwise we ask which is more positive
+#     distance = (df.ans1-df.ans0) * true_switch_sign
+#     return distance
 
 to_tensor = lambda x: torch.from_numpy(x).float()
 to_ds = lambda hs0, y: TensorDataset(to_tensor(hs0), to_tensor(y))

src/datasets/hs.py

@@ -76,7 +76,7 @@ class ExtractHiddenStates:
         choice_ids: List[torch.Tensor] = None,
         truncation_length=999,
         debug=False,
-        counterfactual_fwd=False,
+        counterfactual_fwd=True,
     ):
         """
         Given a decoder model and a batch of texts, gets a pair of hidden states (in a given layer) on that input texts
@@ -106,6 +106,8 @@ class ExtractHiddenStates:
         HEADS = [f"transformer.h.{i}.attn.c_proj" for i in range(self.model.config.num_hidden_layers)]
         MLPS = [f"transformer.h.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
         
+        orig_state_dict = self.model.state_dict()
+        optimizer = torch.optim.SGD(self.model.parameters(),lr=.00002)
         self.model.eval()
         with TraceDict(self.model, HEADS+MLPS, retain_grad=True, detach=True) as ret:
             # with torch.autocast('cuda', torch.bfloat16): # FIXME not reccomended for backwards pass
@@ -162,30 +164,40 @@ class ExtractHiddenStates:
         
         residual_stream = head_activation_and_grad + mlp_activation_and_grad
         
-        if counterfactual_fwd:
-            with TraceDict(self.model, HEADS+MLPS, detach=True) as ret2:
-                orig_state_dict = self.model.state_dict()
-                optimizer = torch.optim.SGD(self.model.parameters(),lr=.00002)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-                outputs2 = self.model(**model_inputs, 
-                    output_hidden_states=True, return_dict=True)
+        if counterfactual_fwd:            
+            
+            # optimizer.zero_grad()
+            # loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
                 
+            with TraceDict(self.model, HEADS+MLPS, detach=True) as ret2:
+                # counterfactual forward pass
+                with torch.no_grad():
+                    outputs2 = self.model(**model_inputs, 
+                        output_hidden_states=True, return_dict=True)
+                    scores2 = outputs2["scores"] = outputs2.logits[:, last_token, :].float()
+                
+                # record info
                 head_activation2 = tcopy(stack_trace_returns(ret2, HEADS))
                 mlp_activation2 = tcopy(stack_trace_returns(ret2, MLPS))
                 residual_stream2 = head_activation2 + mlp_activation2
-                residual_stream2 = residual_stream2[:, layers]
+                residual_stream2 = residual_stream2[:, layers].float()
                 
                 # stack
                 hidden_states2 = list(outputs2.hidden_states)
                 hidden_states2 = rearrange(hidden_states2, 'lyrs b seq hs -> b lyrs seq hs')[:, :, last_token]
-                hidden_states2 = hidden_states2[:, layers]
+                hidden_states2 = hidden_states2[:, layers].float()
+                
+                
             # reset
             self.model.load_state_dict(orig_state_dict)
             optimizer.zero_grad()
+        else:
+            loss.backward()
             
         self.model.eval()
+        
 
 
         # collect outputs
@@ -210,17 +222,18 @@ class ExtractHiddenStates:
             # w_grads_mlp_cfc=w_grads_mlp_cfc,
             # w_grads_attn=w_grads_attn,
         )
-        out = {k: detachcpu(v) for k, v in out.items()}
         if debug:            
             out['input_truncated'] = self.tokenizer.batch_decode(input_ids)
             out['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].argmax(-1))
         
         if counterfactual_fwd:
-            out['residual_stream2'] = residual_stream2
-            out['hidden_states2'] = hidden_states2
+            out['scores2'] = outputs2["scores"]
+            out['hidden_states2'] = hidden_states2.float()
+            out['residual_stream2'] = residual_stream2.float()
+        
+        out = {k: detachcpu(v) for k, v in out.items()}
             
         # I shouldn't have to do this but I get memory leaks
-        self.model.load_state_dict(orig_state_dict)
         outputs = hidden_states = hidden_states2 = loss = orig_state_dict = scores = token_y = token_n = input_ids = attention_mask = choice_ids = residual_stream = residual_stream2 = None
         clear_mem()
             
@@ -237,7 +250,7 @@ class ExtractHiddenStates:
         """
         return torch.arange(
             self.layer_padding,
-            len(outputs["hidden_states"]) - self.layer_padding,
+            len(outputs["hidden_states"])-1 - self.layer_padding,
             self.layer_stride,
         )
 

src/extraction/config.py

@@ -32,10 +32,10 @@ class ExtractConfig(Serializable):
     """Indices of layers to extract hidden states from. We follow the HF convention, so
     0 is the embedding, and 1 is the output of the first transformer layer."""
 
-    layer_stride: InitVar[int] = 1
+    layer_stride: InitVar[int] = 4
     """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`."""
     
-    layer_padding: InitVar[int] = 0
+    layer_padding: InitVar[int] = 4
     """Clips the first and last layers by this amount"""
 
     seed: int = 42