misc

mjc_notes.md

@@ -1546,3 +1546,20 @@ print(pd.Series(ds_tokens['instructed_to_lie']).value_counts()) # should be 50%
 
 
 Ah found it :brain: it was using the same random seed. so I was selecting the Nth each time, which happened to be diff for each dataset. But was the same template and type. OK now I can redo.
+
+
+try with
+- https://huggingface.co/TheBloke/CodeLlama-34B-fp16
+- WizardLM/WizardCoder-Python-13B-V1.0
+
+
+```
+python notebooks/012_make_dataset.py \
+"HuggingFaceH4/starchat-beta" \
+amazon_polarity super_glue:boolq glue:qnli imdb \
+--max_examples 260 260 \
+--max_length=600
+```
+what are fim tokens? Fill-in-the-middle
+
+Fill-in-the-middle uses special tokens to identify the prefix/middle/suffix part of the input and output:

notebooks/012_make_dataset.py

@@ -51,7 +51,16 @@ parser = ArgumentParser(add_help=False)
 parser.add_arguments(ExtractConfig, dest="run")
 
 # argv="""\
-# "WizardLM/WizardCoder-3B-V1.0" \
+# "WizardLM/WizardCoder-Python-13B-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)
+
+# argv="""\
+# "HuggingFaceH4/starchat-beta" \
 # imdb amazon_polarity super_glue:boolq glue:qnli \
 # --max_examples 260 260 \
 # --max_length=600 \
@@ -81,7 +90,7 @@ def load_model(model_repo = "HuggingFaceH4/starchat-beta"):
     model_options = dict(
         device_map="auto",
         # load_in_8bit=True,
-        # load_in_4bit=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,
     )
@@ -273,7 +282,6 @@ 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)
@@ -362,6 +370,22 @@ for ds_name in ds_names:
 
     # ## 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, ...).
+    def expand_choices(choices: List[str]) -> List[str]:
+        """expand out choices by adding versions that are upper, lower, whitespace, etc"""
+        new = []
+        for c in choices:
+            new.append(c)
+            new.append(c.upper())
+            new.append(c.capitalize())
+            new.append(c.lower())
+        return set(new)
+
+
+    left_choices = list(r[0] for r in ds1['answer_choices'])+['no', 'false', 'negative', 'wrong']
+    right_choices = list(r[1] for r in ds1['answer_choices'])+['yes', 'true', 'positive', 'right']
+    left_choices, right_choices = expand_choices(left_choices), expand_choices(right_choices)
+    expanded_choices = [left_choices, right_choices]
+    expanded_choice_ids = choice2ids(expanded_choices, tokenizer)
 
     # this is just based on pairs for that answer...
     add_txt_ans0 = lambda r: {'txt_ans0': tokenizer.decode(r['scores0'].argmax(-1))}
@@ -370,10 +394,12 @@ for ds_name in ds_names:
     add_ans = lambda r: scores2choice_probs(r, row_choice_ids(r), keys=["scores0"])
 
     # Or all expanded choices
+    add_ans_exp = lambda r: scores2choice_probs(r, expanded_choice_ids, prefix="expanded_")
     ds1.set_format(type='numpy')#, columns=['input_ids', 'token_type_ids', 'attention_mask', 'label'])
     ds3 = (
         ds1
         .map(add_ans)
+        .map(add_ans_exp)
         .map(add_txt_ans0)
     )
 

requirements/requirements.txt

@@ -14,6 +14,6 @@ bitsandbytes==0.39.1
 matplotlib
 black
 loguru
-# eleuther-elk==0.1.1
-git+https://github.com/EleutherAI/elk.git@3bbe26c
+eleuther-elk==0.1.1
 # promptsource
+scipy

src/datasets/hs.py

@@ -76,7 +76,6 @@ class ExtractHiddenStates:
         choice_ids: List[torch.Tensor] = None,
         truncation_length=999,
         debug=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
@@ -103,62 +102,70 @@ class ExtractHiddenStates:
 
         # forward pass
         last_token = -1
+        # for WizardLM/WizardCoder-3B-V1.0
         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)]
         
+        # for "WizardLM/WizardCoder-Python-13B-V1.0"
+        HEADS = [f"model.layers.{i}.self_attn" for i in range(self.model.config.num_hidden_layers)]
+        MLPS = [f"model.layers.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
+        
+        layers = HEADS+MLPS
+        module_names = [k for k,v in self.model.named_modules()]
+        layers_not_found = set(layers)-set(module_names)
+        assert len(layers_not_found)==0, f"some layers not found in model: {layers_not_found}. we have {layers}"
+        
         self.model.eval()
         outs = []
         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
-            # 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
-            inputs_embeds = self.model.transformer.wte(input_ids)
-            for _ in range(2):                
-                epsilon=2e-2
-                noise = inputs_embeds.data.new(inputs_embeds.size()).normal_(0, 1) *  epsilon
-                inputs_embeds_w_noise = inputs_embeds + noise
-                model_inputs = self.model.prepare_inputs_for_generation(input_ids=None, inputs_embeds=inputs_embeds_w_noise, 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]
-                    
-                loss = counterfactual_loss(self.model, scores, token_y, token_n)  
-                
-                loss.backward()
-
-                # 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 = tcopy(stack_trace_returns(ret, HEADS))
-                mlp_activation = tcopy(stack_trace_returns(ret, MLPS))
-                residual_stream = head_activation + mlp_activation
-                
-                # select only some layers
-                layers = self.get_layer_selection(outputs)
-                residual_stream = residual_stream[:, layers]
-                hidden_states = hidden_states[:, layers]    
-
-                # collect outputs
-                out = dict(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    scores=outputs["scores"],
-                    layers=layers,
-                    hidden_states=hidden_states,            
-                    residual_stream=residual_stream,
-                )
+            with torch.autocast('cuda', torch.bfloat16):
+                # 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
+                inputs_embeds = self.model.transformer.wte(input_ids)
+                for _ in range(2):                
+                    epsilon=inputs_embeds.abs().mean()*2 # TODO: this worked well for one prompt. Not too differen't, not to simialr. But it's a magic number
+                    noise = inputs_embeds.data.new(inputs_embeds.size()).normal_(0, 1) *  epsilon
+                    inputs_embeds_w_noise = inputs_embeds + noise
+                    model_inputs = self.model.prepare_inputs_for_generation(input_ids=None, inputs_embeds=inputs_embeds_w_noise, 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]
                         
-                if debug:            
-                    out['input_truncated'] = self.tokenizer.batch_decode(input_ids)
-                    out['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].argmax(-1))
-                out = {k: detachcpu(v) for k, v in out.items()}
-                outs.append(out)
+                    # loss = counterfactual_loss(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 = tcopy(stack_trace_returns(ret, HEADS))
+                    mlp_activation = tcopy(stack_trace_returns(ret, MLPS))
+                    residual_stream = head_activation + mlp_activation
+                    
+                    # select only some layers
+                    layers = self.get_layer_selection(outputs)
+                    residual_stream = residual_stream[:, layers]
+                    hidden_states = hidden_states[:, layers]    
+
+                    # collect outputs
+                    out = dict(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        scores=outputs["scores"],
+                        layers=layers,
+                        hidden_states=hidden_states,            
+                        residual_stream=residual_stream,
+                    )
+                            
+                    if debug:            
+                        out['input_truncated'] = self.tokenizer.batch_decode(input_ids)
+                        out['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].softmax(-1).argmax(-1))
+                    out = {k: detachcpu(v) for k, v in out.items()}
+                    outs.append(out)
             
         # I shouldn't have to do this but I get memory leaks
         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
@@ -169,16 +176,22 @@ class ExtractHiddenStates:
 
     def get_layer_selection(self, outputs):
         """Sometimes we don't want to save all layers.
+        
+        We skip the first few (data leakage?). Stride the the middle (could be valuable), and include the last few (possibly high level concepts).
 
-        Typically we can skip some to save space (stride). We might also want to ignore the first and last ones (padding) to avoid data leakage.
-
-        See https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
+        See also https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
         """
-        return torch.arange(
+        # for self.layer_padding, skip the first few
+        strided_layers = torch.arange(
             self.layer_padding,
-            len(outputs["hidden_states"])-1 - self.layer_padding,
-            self.layer_stride,
+            len(outputs["hidden_states"])-1,
+            self.layer_stride-self.layer_padding,
         )
+        # 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(self.layer_padding-self.layer_padding, self.layer_padding)
+        layers = strided_layers+last_few
+        # TODO: check for dups
+        return layers
 
 def detachcpu(x):
     """

src/extraction/config.py

@@ -36,7 +36,7 @@ class ExtractConfig(Serializable):
     """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`."""
     
     layer_padding: InitVar[int] = 4
-    """Clips the first and last layers by this amount"""
+    """Clips the first layers by this amount"""
 
     seed: int = 42
     """Seed to use for prompt randomization. Defaults to 42."""

src/helpers/torch.py

@@ -3,6 +3,16 @@ 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):
     """
@@ -19,7 +29,7 @@ def to_numpy(x):
 
 
 
-def set_seeds(n):
+def set_seeds(n: int) -> None:
     transformers.set_seed(n)
     torch.manual_seed(n)
     np.random.seed(n)