80% on linear probe... hmm

mjc_notes.md

@@ -1193,3 +1193,13 @@ then they only take the top 48 attentions heads, and only direction
 So I have it making a dataset. I still need to work out the label. And I need a large sample.
 
 But the worst thing is I've got from 15B models with a batch of 10, to 3B models with a batch of 1 (1/50x). Just by going from 4bit to 16 (4x) but also adding gradient/graph (contributed 10x?).
+
+# 2023-09-04 06:29:58
+
+A quick linear probe gets 77% acc on test set. hmm
+
+
+hmm los of interesting stuff in this code https://github.com/likenneth/honest_llama/blob/master/utils.py#L17
+- datasets
+- models
+- getting hidden states

src/datasets/batch.py

@@ -48,7 +48,9 @@ def batch_hidden_states(model, tokenizer, data: Dataset, batch_size=2, mcdropout
                 # int16 makes our storage much smaller
                 hs0=float_to_int16(torch.from_numpy(hs0['hidden_states'][j])),
                 scores0=hs0["scores"][j],
-                grads0=hs0['grads'][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],                    

src/datasets/hs.py

@@ -107,7 +107,13 @@ class ExtractHiddenStates:
         grads_all = get_gradients(self.model, outputs, token_y, token_n)
         p = ".+mlp.c_proj.weight" # get the last weight of each layer (ignore bias)
         # p = ".+mlp.c_proj.bias" # get the last weight of each layer
-        grads = torch.stack([g.mean(1).float() for k,g in grads_all.items() if re.match(p, k)])
+        grads_mlp = torch.stack([g.mean(1).float() for k,g in grads_all.items() if re.match(p, k)])
+        
+        p = ".+attn.c_proj.weight" # get the last weight of each layer (ignore bias)
+        grads_attn = torch.stack([g.mean(0).float() for k,g in grads_all.items() if re.match(p, k)])
+        
+        p = ".+mlp.c_fc.weight" # get the last weight of each layer (ignore bias)
+        grads_mlp_cfc = torch.stack([g.mean(0).float() for k,g in grads_all.items() if re.match(p, k)])
         
         hidden_states = torch.stack(
             [outputs["hidden_states"][i] for i in layers], 1
@@ -122,7 +128,9 @@ class ExtractHiddenStates:
             scores=outputs["scores"],
             input_ids=input_ids,
             layers=layers,
-            grads = grads,
+            grads_attn = grads_attn,
+            grads_mlp=grads_mlp,
+            grads_mlp_cfc=grads_mlp_cfc,
         )
         out = {k: to_numpy(v) for k, v in out.items()}
         if debug: