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

wip

Browse Source
This commit is contained in:
deep1
2023-09-09 08:08:11 +08:00
parent 6ffae33c49
commit 4aed182abe
5 changed files with 348 additions and 436 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjc_notes.md
+5
View File
@@ -1230,3 +1230,8 @@ mlp 76%
attn 75%
previouslly I was extracting the grad on the weights. now it's the grad on the outputs/activations which seems better although perhaps harder to classify!
# 2023-09-09 08:06:54
- [ ] run probe on some data
- [ ] add a mlp one too
notebooks/03_make_dataset.ipynb
+277 -375
View File
File diff suppressed because one or more lines are too long
src/datasets/batch.py
+4 -10
View File
@@ -27,7 +27,6 @@ def batch_hidden_states(model, tokenizer, data: Dataset, batch_size=2, mcdropout
ds_t_subset.set_format(type='torch')
ds_p_subset = data.remove_columns(torch_cols)
# TODO check it has a few critical ones in
dl = DataLoader(ds_t_subset, batch_size=batch_size, shuffle=False)
for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
@@ -46,19 +45,14 @@ def batch_hidden_states(model, tokenizer, data: Dataset, batch_size=2, mcdropout
large_arrays_keys = [k for k,v in hs0.items() if v.ndim>2]
large_arrays_as_int16 = {
k:float_to_int16(torch.from_numpy(hs0[k][j]))
# k:float_to_int16(hs0[k][j])
k:hs0[k][j]
for k in large_arrays_keys}
yield dict(
large_arrays_keys=large_arrays_keys,
scores0=hs0["scores"][j],
# grads_mlp0=hs0['grads_mlp'][j],
# grads_mlp_cfc0=hs0['grads_mlp_cfc'][j],
# grads_attn0=hs0['grads_attn'][j],
# hs1=float_to_int16(torch.from_numpy(hs1['hidden_states'][j])),
# scores1=hs1["scores"][j],
# large_arrays_keys=large_arrays_keys,
scores0=hs0["scores"][j],
ds_index=index[j],
src/datasets/hs.py
+59 -49
View File
@@ -25,10 +25,14 @@ from datasets import Dataset
import numpy as np
import torch
import torch.nn.functional as F
from baukit import Trace, TraceDict
from baukit.nethook import Trace, TraceDict, recursive_copy
from einops import rearrange, reduce, repeat
from src.datasets.scores import choice2id, choice2ids
def tcopy(x: torch.Tensor):
return x.clone().detach().cpu()
def counterfactual_backwards(model, scores, token_y, token_n):
"""do a backwards pass where the loss is the distance to the opposite scores"""
model.zero_grad()
@@ -44,7 +48,7 @@ def stack_trace_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
def stack_trace_grad_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
hs = [ret[h].output.grad for h in names]
hs = [ret[h].output.grad.detach() for h in names]
return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
def select_weight_grads(weight_grads: Dict[str, torch.Tensor], pattern:str= ".+attn.c_proj.weight", mean_axis:int=1):
@@ -57,8 +61,8 @@ class ExtractHiddenStates:
model: PreTrainedModel
tokenizer: PreTrainedTokenizer
layer_stride: int = 1
layer_padding: int = 2
layer_stride: int = 8
layer_padding: int = 3
def get_batch_of_hidden_states(
@@ -99,50 +103,42 @@ class ExtractHiddenStates:
MLPS = [f"transformer.h.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
self.model.train()
with TraceDict(self.model, HEADS+MLPS, retain_grad=True) as ret:
with torch.autocast('cuda'): # FIXME not reccomended for backwards pass
# Forward for one step is the same as greedy generation for one step
# https://github.com/huggingface/transformers/blob/234cfefbb083d2614a55f6093b0badfb2efc3b45/src/transformers/generation_utils.py#L1528
model_inputs = self.model.prepare_inputs_for_generation(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
outputs = self.model.forward(
**model_inputs,
return_dict=True,
output_hidden_states=True,
)
scores = outputs["scores"] = outputs.logits[:, last_token, :]
token_n = choice_ids[:, 0] # [batch, tokens]
token_y = choice_ids[:, 1]
# with torch.autocast('cuda', torch.bfloat16): # FIXME not reccomended for backwards pass
# Forward for one step is the same as greedy generation for one step
# https://github.com/huggingface/transformers/blob/234cfefbb083d2614a55f6093b0badfb2efc3b45/src/transformers/generation_utils.py#L1528
model_inputs = self.model.prepare_inputs_for_generation(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
outputs = self.model.forward(
**model_inputs,
return_dict=True,
output_hidden_states=True,
)
scores = outputs["scores"] = outputs.logits[:, last_token, :].float()
token_n = choice_ids[:, 0] # [batch, tokens]
token_y = choice_ids[:, 1]
counterfactual_backwards(self.model, scores, token_y, token_n)
# stack
hidden_states = list(outputs.hidden_states)
hidden_states = rearrange(hidden_states, 'lyrs b seq hs -> b lyrs seq hs')[:, :, last_token]
## from ret, we get the layer activation and the grads on them
head_activation = stack_trace_returns(ret, HEADS)
mlp_activation = stack_trace_returns(ret, MLPS)
head_activation_grads = tcopy(stack_trace_grad_returns(ret, HEADS))
mlp_activation_grads = tcopy(stack_trace_grad_returns(ret, MLPS))
## we also get the gradients on weights, as this might be a lower dimensional space than the grads on activations
ret = None
ps = self.model.named_parameters()
weight_grads = {n:g.grad.detach().float().cpu()[None, :] for n,g in ps if g.grad is not None}
weight_grads = {
n: tcopy(g.grad)[None, :]
for n,g in ps if g.grad is not None}
w_grads_mlp = select_weight_grads(weight_grads, pattern= ".+attn.c_proj.weight", mean_axis=1)
w_grads_attn = select_weight_grads(weight_grads, pattern= ".+attn.c_attn.weight", mean_axis=0)
w_grads_mlp_cfc = select_weight_grads(weight_grads, pattern= ".+mlp.c_fc.weight", mean_axis=0)
weight_grads = None
self.model.zero_grad()
# stack
hidden_states = list(outputs.hidden_states)
hidden_states = rearrange(hidden_states, 'lyrs b seq hs -> b lyrs seq hs')[:, :, last_token]
## from ret, we get the layer activation and the grads on them
head_activation = stack_trace_returns(ret, HEADS)
mlp_activation = stack_trace_returns(ret, MLPS)
head_activation_grads = stack_trace_grad_returns(ret, HEADS)
mlp_activation_grads = stack_trace_grad_returns(ret, MLPS)
## we also get the gradients on weights, as this might be a lower dimensional space than the grads on activations
p = ".+mlp.c_proj.weight" # get the last weight of each layer (ignore bias)
# rearrange([g.mean(1).float() for k,g in weight_grads.items() if re.match(p, k)])
# w_grads_mlp = torch.stack([g.mean(1).float() for k,g in weight_grads.items() if re.match(p, k)])
w_grads_mlp = select_weight_grads(weight_grads, pattern= ".+attn.c_proj.weight", mean_axis=1)
w_grads_attn = select_weight_grads(weight_grads, pattern= ".+attn.c_attn.weight", mean_axis=0)
w_grads_mlp_cfc = select_weight_grads(weight_grads, pattern= ".+mlp.c_fc.weight", mean_axis=0)
# p = ".+attn.c_proj.weight" # get the last weight of each layer (ignore bias)
# w_grads_attn = torch.stack([g.mean(0).float() for k,g in weight_grads.items() if re.match(p, k)])
# p = ".+mlp.c_fc.weight" # get the last weight of each layer (ignore bias)
# w_grads_mlp_cfc = torch.stack([g.mean(0).float() for k,g in weight_grads.items() if re.match(p, k)])
# select only some layers
layers = self.get_layer_selection(outputs)
@@ -165,21 +161,23 @@ class ExtractHiddenStates:
hidden_states=hidden_states,
head_activation=head_activation,
mlp_activation=mlp_activation,
# mlp_activation=mlp_activation,
head_activation_grads = head_activation_grads,
mlp_activation_grads=mlp_activation_grads,
# mlp_activation_grads=mlp_activation_grads,
w_grads_mlp=w_grads_mlp,
w_grads_mlp_cfc=w_grads_mlp_cfc,
# w_grads_mlp=w_grads_mlp,
# w_grads_mlp_cfc=w_grads_mlp_cfc,
w_grads_attn=w_grads_attn,
)
out = {k: to_numpy(v) for k, v in out.items()}
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))
return out
def get_layer_selection(self, outputs):
"""Sometimes we don't want to save all layers.
@@ -194,3 +192,15 @@ class ExtractHiddenStates:
self.layer_stride,
)
def detachcpu(x):
"""
Trys to convert torch 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
else:
return x
src/datasets/load.py
+3 -2
View File
@@ -38,5 +38,6 @@ def ds2df(ds, cols=None):
def load_ds(f):
ds = load_from_disk(f)
ks = ds['large_arrays_keys'][0]
return ds.map(lambda x: {k: int16_to_float(torch.from_numpy(ds[k])) for k in ks})
ks = [k for k,v in ds[0].items() if (v.dtype=='int64') and k not in ['ds_index']]
# ds = ds.map(lambda x: {k: int16_to_float(torch.from_numpy(ds[k]).long()) for k in ks})
return ds