diff --git a/data/interventions/TheBlokeWizardCoder_Python_13B_V1.0_GPTQ.pkl b/data/interventions/TheBlokeWizardCoder_Python_13B_V1.0_GPTQ.pkl deleted file mode 100644 index c826826..0000000 Binary files a/data/interventions/TheBlokeWizardCoder_Python_13B_V1.0_GPTQ.pkl and /dev/null differ diff --git a/notebooks/027_train_nanda_probe_w_counterfact_70%.ipynb b/notebooks/027_train_nanda_probe_w_counterfact_70%.ipynb index 5ac1182..a25cbb3 100644 --- a/notebooks/027_train_nanda_probe_w_counterfact_70%.ipynb +++ b/notebooks/027_train_nanda_probe_w_counterfact_70%.ipynb @@ -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", diff --git a/notebooks/make_dataset.py b/notebooks/make_dataset.py index 529bd75..26b127f 100644 --- a/notebooks/make_dataset.py +++ b/notebooks/make_dataset.py @@ -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) diff --git a/src/datasets/dm.py b/src/datasets/dm.py index 7a3cb9a..a9bd927 100644 --- a/src/datasets/dm.py +++ b/src/datasets/dm.py @@ -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, diff --git a/src/datasets/hs.py b/src/datasets/hs.py index d1b7d1f..4ec7e67 100644 --- a/src/datasets/hs.py +++ b/src/datasets/hs.py @@ -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): """ diff --git a/src/datasets/intervene.py b/src/datasets/intervene.py index 9d8b7ac..4b85043 100644 --- a/src/datasets/intervene.py +++ b/src/datasets/intervene.py @@ -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 diff --git a/src/helpers/ds.py b/src/helpers/ds.py index d6b3b81..0386684 100644 --- a/src/helpers/ds.py +++ b/src/helpers/ds.py @@ -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)