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

use mlp too, and ignore head distinction

Browse Source
This commit is contained in:
wassname
2023-10-20 07:41:29 +08:00
parent e3ecc17fcf
commit d8d7a5278c
7 changed files with 64 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/interventions/TheBlokeWizardCoder_Python_13B_V1.0_GPTQ.pkl
BIN
View File
Binary file not shown.
notebooks/027_train_nanda_probe_w_counterfact_70%.ipynb
+3 -2
View File
@@ -77,7 +77,7 @@
"from torch import Tensor\n",
"from torch import optim\n",
"from torch.utils.data import random_split, DataLoader, TensorDataset\n",
"\n",
"from src.helpers.ds import shuffle_dataset_by\n",
"from pathlib import Path\n",
"\n",
"import transformers\n",
@@ -1380,7 +1380,8 @@
" df['y'] = self.y>0\n",
" \n",
"\n",
" # let's create a simple 50/50 train split (the data is already randomized)\n",
" # let's create a simple 50/50 train split (the data is already randomized during gathering) but ordered by example_i so there is little to no overlap\n",
" # FIXME make zero overlap using `shuffle_dataset_by` but rewrite as sort_dataset_by or stratified split in sklearn\n",
" n = len(self.ans)\n",
" self.splits = {\n",
" 'train': (0, int(n * 0.5)),\n",
notebooks/make_dataset.py
+10 -11
View File
@@ -42,7 +42,7 @@ 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 src.helpers.ds import shuffle_dataset_by
from src.datasets.hs import ExtractHiddenStates
from itertools import chain
import functools
@@ -266,13 +266,13 @@ def load_preproc_dataset(ds_name: str, cfg: ExtractConfig, tokenizer: PreTrained
.map(lambda r: {'choice_ids': row_choice_ids(r, tokenizer)}, desc='choice_ids')
)
ds_tokens = ds_tokens.filter(lambda r: r['truncated']==False)
inds = list(range(min(len(ds_tokens), N)))
random.shuffle(inds)
ds_tokens = ds_tokens.select(inds)
ds_tokens = shuffle_dataset_by(ds_tokens, 'example_i')
print('removed truncated rows to leave: num_rows', ds_tokens.num_rows)
return ds_tokens
def row_choice_ids(r, tokenizer):
return choice2ids([[c] for c in r['answer_choices']], tokenizer)
@@ -359,18 +359,17 @@ def create_intervention(ds_name, ds_tokens, model, layer_names, N=10):
activations = np.array(ds_calibration['head_activation']).squeeze(-1)
labels = np.array(ds_calibration["label_true"]).astype(int)==1
num_heads = model.config.num_attention_heads
interventions = get_interventions_dict(activations, labels, layer_names, num_heads)
interventions = get_interventions_dict(activations, labels, layer_names)
return interventions
def load_intervention(ds_name, cfg, model, tokenizer, model_name, N=30):
def load_intervention(ds_name, cfg, model, tokenizer, model_name, N=50):
num_heads = model.config.num_attention_heads
intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'
intervention_f.parent.mkdir(exist_ok=True, parents=True)
if not intervention_f.exists():
layer_names, layer_inds = ExtractHiddenStates(model, tokenizer, layer_stride=cfg.layer_stride, layer_padding=cfg.layer_padding).get_layer_names()
ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer, N=N)
layer_names = ExtractHiddenStates(model, tokenizer, layer_stride=cfg.layer_stride, layer_padding=cfg.layer_padding).get_layer_names()
ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer, N=N*2)
interventions = create_intervention(ds_name, ds_tokens, model, layer_names)
torch.save(interventions, intervention_f)
else:
@@ -431,7 +430,7 @@ if __name__ == "__main__":
# get dataset filename
N = len(ds_tokens)
dataset_name = f"{sanitize(cfg.model)}_{ds_name}_{split_type}_{N}"
f = f"../.ds/{dataset_name}"
f = root_folder / '.ds'/ "{dataset_name}"
ds1 = create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts=intervention, f=f)
src/datasets/dm.py
+4 -1
View File
@@ -6,7 +6,7 @@ from torch.utils.data import DataLoader, TensorDataset
from src.datasets.load import ds2df
from datasets.arrow_dataset import Dataset
from einops import rearrange, reduce, repeat
from src.helpers.ds import shuffle_dataset_by
# def compute_distance(df):
# """distance between ans1 and ans2."""
@@ -18,6 +18,9 @@ from einops import rearrange, reduce, repeat
to_tensor = lambda x: torch.from_numpy(x).float()
to_ds = lambda hs0, y: TensorDataset(to_tensor(hs0), to_tensor(y))
class imdbHSDataModule(pl.LightningDataModule):
def __init__(self,
src/datasets/hs.py
+16 -14
View File
@@ -35,13 +35,13 @@ from src.datasets.intervene import InterventionDict, intervention_meta_fn
from functools import partial
def noise_for_embeds(inputs_embeds, seed=42, std = 2e-2):
B, S, embed_dim = inputs_embeds.shape
with torch.random.fork_rng(devices=[inputs_embeds.device.index]):
torch.manual_seed(seed)
noise = torch.normal(0., std, (embed_dim, ))
noise = repeat(noise, 't -> b s t', b=B, s=S).to(inputs_embeds.device).to(inputs_embeds.dtype)
return noise
# def noise_for_embeds(inputs_embeds, seed=42, std = 2e-2):
# B, S, embed_dim = inputs_embeds.shape
# with torch.random.fork_rng(devices=[inputs_embeds.device.index]):
# torch.manual_seed(seed)
# noise = torch.normal(0., std, (embed_dim, ))
# noise = repeat(noise, 't -> b s t', b=B, s=S).to(inputs_embeds.device).to(inputs_embeds.dtype)
# return noise
def tcopy(x: torch.Tensor):
return x.clone().detach().cpu()
@@ -90,11 +90,9 @@ class ExtractHiddenStates:
# for "WizardLM/WizardCoder-Python-13B-V1.0"
# HACK: depends on model layout
layers_names = [f"model.layers.{i}.self_attn" for i in range(self.model.config.num_hidden_layers)]
module_names = [k for k,v in self.model.named_modules()]
layers_not_found = set(layers_names)-set(module_names)
assert len(layers_not_found)==0, f"some layers not found in model: {layers_not_found}. we have {layers_names}"
return self.get_layer_selection(layers_names)
layers_names_h = [f"model.layers.{i}.self_attn" for i in range(self.model.config.num_hidden_layers)]
layers_names_mlp = [f"model.layers.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
return self.get_layer_selection(layers_names_h) + self.get_layer_selection(layers_names_mlp)
def get_batch_of_hidden_states(
@@ -132,7 +130,7 @@ class ExtractHiddenStates:
# forward pass
last_token = -1
layers_names, layer_inds = self.get_layer_names()
layers_names = self.get_layer_names()
self.model.eval()
@@ -204,6 +202,10 @@ class ExtractHiddenStates:
See also https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
"""
module_names = [k for k,v in self.model.named_modules()]
layers_not_found = set(layer_names)-set(module_names)
assert len(layers_not_found)==0, f"some layers not found in model: {layers_not_found}. we have {layer_names}"
# for self.layer_padding, skip the first few
num_layers = len(layer_names)-1
strided_layers = torch.arange(
@@ -214,7 +216,7 @@ class ExtractHiddenStates:
# for self.layer_padding ALWAYS include the last few. Why, this is based on the intuition that the last layers may be the most valuable
last_few = torch.arange(num_layers-self.layer_padding, num_layers).tolist()
layers_inds = sorted(set(list(strided_layers)+list(last_few)))
return [layer_names[i] for i in layers_inds], layers_inds
return [layer_names[i] for i in layers_inds]
def detachcpu(x):
"""
src/datasets/intervene.py
+17 -108
View File
@@ -8,128 +8,39 @@ from typing import List, Tuple, Dict, Any, Union, NewType
from einops import rearrange, reduce, repeat, asnumpy, parse_shape
import torch
InterventionDict = NewType('InterventionDict', Dict[str, List[Tuple[int, np.ndarray, float]]])
# def get_com_directions(
# num_layers: int,
# num_heads: int,
# head_wise_activations: np.ndarray,
# labels: np.ndarray,
# ) -> np.ndarray:
# """get center of mass direction's for each layer and head."""
# assert max(labels) == 1
# com_directions = []
# for layer in range(num_layers):
# for head in range(num_heads):
# usable_idxs = range(len(head_wise_activations))
# usable_head_wise_activations = np.concatenate(
# [head_wise_activations[i][:, layer, head, :] for i in usable_idxs],
# axis=0,
# )
# true_mass_mean = np.mean(usable_head_wise_activations[labels == 1], axis=0)
# false_mass_mean = np.mean(usable_head_wise_activations[labels == 0], axis=0)
# com_directions.append(true_mass_mean - false_mass_mean)
# com_directions = np.array(com_directions)
# return com_directions
InterventionDict = NewType('InterventionDict', Dict[str, List[Tuple[np.ndarray, float]]])
# def get_magnitude(activations: np.ndarray, labels: np.ndarray) -> np.ndarray:
# """
# refactored to from https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/utils.py#L698
# to use einops and vector ops instead of for loop
# """
# true_mass_mean = reduce(activations[labels], ' b l h -> l h', 'mean')
# false_mass_mean = reduce(activations[~labels], ' b l h -> l h', 'mean')
# direction = true_mass_mean - false_mass_mean
# direction = direction / np.linalg.norm(direction, axis=1, keepdims=True)
# activations = reduce(activations, ' b l h -> l h', 'mean')
# proj_vals = activations * direction
# proj_val_std = reduce(proj_vals, 'l h -> l', np.std)
# return proj_val_std
# def layer_head_to_flattened_idx(layer, head, num_heads):
# return layer * num_heads + head
# def get_interventions_dict(
# probes,
# tuning_activations,
# layer_heads: List[Tuple[int, int]],
# num_heads: int,
# com_directions,
# ) -> InterventionDict:
# """
# Make an intervention dict that works with baukit.TraceDict's edit_output.
# see https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56
# """
# # init
# interventions = InterventionDict({})
# for layer, head in layer_heads:
# interventions[f"model.layers.{layer}.self_attn.head_out"] = []
# std = get_magnitude(activations, labels)
# # work out magnitude of intervention, then record
# for layer, head in layer_heads:
# direction = com_directions[layer_head_to_flattened_idx(layer, head, num_heads)]
# direction = direction / np.linalg.norm(direction)
# activations = tuning_activations[:, layer, head, :] # batch x 128
# proj_vals = activations @ direction.T
# proj_val_std = float(np.std(proj_vals)) # TODO check this is meant to be float
# interventions[f"model.layers.{layer}.self_attn.head_out"].append(
# (head, direction.squeeze(), proj_val_std)
# )
# # sort keys by head index
# for layer, head in layer_heads:
# interventions[f"model.layers.{layer}.self_attn.head_out"] = sorted(
# interventions[f"model.layers.{layer}.self_attn.head_out"],
# key=lambda x: x[0],
# )
# return interventions
def get_magnitude(layer_activations: np.ndarray, labels: np.ndarray, num_heads:int) -> Tuple[np.ndarray,np.ndarray]:
def get_magnitude(activations: np.ndarray, labels: np.ndarray) -> Tuple[np.ndarray,np.ndarray]:
"""
get center of mass direction and magnitude per layer and head
refactored to from https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/utils.py#L698
to use einops and vector ops instead of for loop
"""
# we intervene with statistic per head, like in honest_llama
activations = rearrange(layer_activations, 'b l (h d) -> b l h d', h=num_heads)
true_mass_mean = reduce(activations[labels], ' b l h d -> l h d', 'mean')
false_mass_mean = reduce(activations[~labels], ' b l h d -> l h d', 'mean')
"""
true_mass_mean = reduce(activations[labels], ' b l d -> l d', 'mean')
false_mass_mean = reduce(activations[~labels], ' b l d -> l d', 'mean')
direction = true_mass_mean - false_mass_mean
direction = direction / np.linalg.norm(direction, axis=1, keepdims=True) # sq norm per layer
activations = reduce(activations, ' b l h d -> l h d', 'mean')
activations = reduce(activations, ' b l d -> l d', 'mean')
proj_vals = activations * direction
proj_val_std = reduce(proj_vals, 'l h d -> l h', np.std)
proj_val_std = reduce(proj_vals, 'l d -> l', np.std)
return direction, proj_val_std
def get_interventions_dict(activations:np.ndarray, labels: np.ndarray, layer_names: List[str], num_heads: int) -> InterventionDict:
def get_interventions_dict(activations:np.ndarray, labels: np.ndarray, layer_names: List[str]) -> InterventionDict:
"""
Make an intervention dict that works with baukit.TraceDict's edit_output.
see https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56
"""
direction, proj_val_std = get_magnitude(activations, labels, num_heads)
direction, proj_val_std = get_magnitude(activations, labels)
out = InterventionDict({l:[] for l in layer_names})
for layer_i, ln in enumerate(layer_names):
for head in range(num_heads):
out[ln].append((head, direction[layer_i, head].squeeze(), proj_val_std[layer_i, head]))
out[ln].append((direction[layer_i].squeeze(), proj_val_std[layer_i]))
return out
def intervention_meta_fn(head_outputs: torch.Tensor, layer_name:str, interventions: InterventionDict, num_heads, alpha = 15) -> torch.Tensor:
def intervention_meta_fn(outputs: torch.Tensor, layer_name:str, interventions: InterventionDict, alpha = 15) -> torch.Tensor:
"""see
- honest_llama: https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/validation/validate_2fold.py#L114
- baukit: https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56
@@ -140,11 +51,9 @@ def intervention_meta_fn(head_outputs: torch.Tensor, layer_name:str, interventio
...
"""
head_output, a, b = head_outputs
head_output = rearrange(head_output, 'b s (h d) -> b s h d', h=num_heads)
for head, direction, proj_val_std in interventions[layer_name]:
# head_output: (batch_size, seq_len, num_heads, head_size)
head_output[:, -1:, head, :] += torch.from_numpy(alpha * proj_val_std * direction).to(head_output.device)
head_output = rearrange(head_output, 'b s h d -> b s (h d)')
head_outputs = (head_output, a, b)
return head_outputs
output, a, b = outputs
for direction, proj_val_std in interventions[layer_name]:
# head_output: (batch_size, seq_len, layer_size)
output[:, -1:, :] += torch.from_numpy(alpha * proj_val_std * direction).to(output.device)
outputs = (output, a, b)
return outputs
src/helpers/ds.py
+14 -1
View File
@@ -1,13 +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)
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_tokens["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)