sum aggregation works now; refactoring some stuff

Author: saprmarks

attribution.py

@@ -2,58 +2,14 @@
 import torch as t
 from tqdm import tqdm
 from numpy import ndindex
+from loading_utils import Submodule
 from activation_utils import SparseAct
-from dataclasses import dataclass
 from nnsight.envoy import Envoy
 from dictionary_learning.dictionary import Dictionary
 from typing import Callable
 
-DEBUGGING = False
-
-if DEBUGGING:
-    tracer_kwargs = {'validate' : True, 'scan' : True}
-else:
-    tracer_kwargs = {'validate' : False, 'scan' : False}
-
 EffectOut = namedtuple('EffectOut', ['effects', 'deltas', 'grads', 'total_effect'])
 
-@dataclass(frozen=True)
-class Submodule:
-    name: str
-    submodule: Envoy
-    use_input: bool = False
-    is_tuple: bool = False
-
-    def __hash__(self):
-        return hash(self.name)
-
-    def get_activation(self):
-        if self.use_input:
-            return self.submodule.input # TODO make sure I didn't break for pythia
-        else:
-            if self.is_tuple:
-                return self.submodule.output[0]
-            else:
-                return self.submodule.output
-
-    def set_activation(self, x):
-        if self.use_input:
-            self.submodule.input[:] = x
-        else:
-            if self.is_tuple:
-                self.submodule.output[0][:] = x
-            else:
-                self.submodule.output = x
-
-    def stop_grad(self):
-        if self.use_input:
-            self.submodule.input.grad = t.zeros_like(self.submodule.input)
-        else:
-            if self.is_tuple:
-                self.submodule.output[0].grad = t.zeros_like(self.submodule.output[0])
-            else:
-                self.submodule.output.grad = t.zeros_like(self.submodule.output)
-
 
 def _pe_attrib(
         clean,
@@ -66,7 +22,7 @@ def _pe_attrib(
 ):
     hidden_states_clean = {}
     grads = {}
-    with model.trace(clean, **tracer_kwargs):
+    with model.trace(clean):
         for submodule in submodules:
             dictionary = dictionaries[submodule]
             x = submodule.get_activation()
@@ -90,7 +46,7 @@ def _pe_attrib(
         total_effect = None
     else:
         hidden_states_patch = {}
-        with model.trace(patch, **tracer_kwargs), t.inference_mode():
+        with t.no_grad(), model.trace(patch):
             for submodule in submodules:
                 dictionary = dictionaries[submodule]
                 x = submodule.get_activation()
@@ -126,7 +82,7 @@ def _pe_ig(
     metric_kwargs=dict(),
 ):
     hidden_states_clean = {}
-    with t.no_grad(), model.trace(clean, **tracer_kwargs):
+    with t.no_grad(), model.trace(clean):
         for submodule in submodules:
             dictionary = dictionaries[submodule]
             x = submodule.get_activation()
@@ -144,7 +100,7 @@ def _pe_ig(
         total_effect = None
     else:
         hidden_states_patch = {}
-        with t.no_grad(), model.trace(patch, **tracer_kwargs):
+        with t.no_grad(), model.trace(patch):
             for submodule in submodules:
                 dictionary = dictionaries[submodule]
                 x = submodule.get_activation()
@@ -163,7 +119,7 @@ def _pe_ig(
         dictionary = dictionaries[submodule]
         clean_state = hidden_states_clean[submodule]
         patch_state = hidden_states_patch[submodule]
-        with model.trace(**tracer_kwargs) as tracer:
+        with model.trace() as tracer:
             metrics = []
             fs = []
             for step in range(steps):
@@ -172,7 +128,7 @@ def _pe_ig(
                 f.act.requires_grad_().retain_grad()
                 f.res.requires_grad_().retain_grad()
                 fs.append(f)
-                with tracer.invoke(clean, scan=tracer_kwargs['scan']):
+                with tracer.invoke(clean):
                     submodule.set_activation(dictionary.decode(f.act) + f.res)
                     metrics.append(metric_fn(model, **metric_kwargs))
             metric = sum([m for m in metrics])
@@ -200,7 +156,7 @@ def _pe_exact(
     metric_fn,
 ):
     hidden_states_clean = {}
-    with model.trace(clean, **tracer_kwargs), t.inference_mode():
+    with t.no_grad(), model.trace(clean):
         for submodule in submodules:
             dictionary = dictionaries[submodule]
             x = submodule.get_activation()
@@ -218,7 +174,7 @@ def _pe_exact(
         total_effect = None
     else:
         hidden_states_patch = {}
-        with model.trace(patch, **tracer_kwargs), t.inference_mode():
+        with t.no_grad(), model.trace(patch):
             for submodule in submodules:
                 dictionary = dictionaries[submodule]
                 x = submodule.get_activation()
@@ -241,24 +197,22 @@ def _pe_exact(
         # iterate over positions and features for which clean and patch differ
         idxs = t.nonzero(patch_state.act - clean_state.act)
         for idx in tqdm(idxs):
-            with t.inference_mode():
-                with model.trace(clean, **tracer_kwargs):
-                    f = clean_state.act.clone()
-                    f[tuple(idx)] = patch_state.act[tuple(idx)]
-                    x_hat = dictionary.decode(f)
-                    submodule.set_activation(x_hat + clean_state.res)
-                    metric = metric_fn(model).save()
-                effect.act[tuple(idx)] = (metric.value - metric_clean.value).sum()
+            with t.no_grad(), model.trace(clean):
+                f = clean_state.act.clone()
+                f[tuple(idx)] = patch_state.act[tuple(idx)]
+                x_hat = dictionary.decode(f)
+                submodule.set_activation(x_hat + clean_state.res)
+                metric = metric_fn(model).save()
+            effect.act[tuple(idx)] = (metric.value - metric_clean.value).sum()
 
         for idx in list(ndindex(effect.resc.shape)): # type: ignore
-            with t.inference_mode():
-                with model.trace(clean, **tracer_kwargs):
-                    res = clean_state.res.clone()
-                    res[tuple(idx)] = patch_state.res[tuple(idx)] # type: ignore
-                    x_hat = dictionary.decode(clean_state.act)
-                    submodule.set_activation(x_hat + res)
-                    metric = metric_fn(model).save()
-                effect.resc[tuple(idx)] = (metric.value - metric_clean.value).sum() # type: ignore
+            with t.no_grad(), model.trace(clean):
+                res = clean_state.res.clone()
+                res[tuple(idx)] = patch_state.res[tuple(idx)] # type: ignore
+                x_hat = dictionary.decode(clean_state.act)
+                submodule.set_activation(x_hat + res)
+                metric = metric_fn(model).save()
+            effect.resc[tuple(idx)] = (metric.value - metric_clean.value).sum() # type: ignore
 
         effects[submodule] = effect
         deltas[submodule] = patch_state - clean_state
@@ -311,8 +265,7 @@ def jvp(
     vjv_values = {}
 
     downstream_feature_idxs = downstream_features.to_tensor().nonzero()
-    with model.trace(input, **tracer_kwargs):
-
+    with model.trace(input):
         # forward pass modifications
         x = upstream_submod.get_activation()
         x_hat, f = upstream_dict.hacked_forward_for_sfc(x) # hacking around an nnsight bug

circuit.py

@@ -11,19 +11,11 @@
 from attribution import patching_effect, jvp
 from circuit_plotting import plot_circuit, plot_circuit_posaligned
 from dictionary_learning import AutoEncoder
-from loading_utils import load_examples, load_examples_nopair
+from data_loading_utils import load_examples, load_examples_nopair
 from dictionary_loading_utils import load_saes_and_submodules
 from nnsight import LanguageModel
 from coo_utils import sparse_reshape
 
-DEBUGGING = True
-
-if DEBUGGING:
-    tracer_kwargs = {"validate": True, "scan": True}
-else:
-    tracer_kwargs = {"validate": False, "scan": False}
-
-
 def get_circuit(
     clean,
     patch,
@@ -35,11 +27,10 @@ def get_circuit(
     dictionaries,
     metric_fn,
     metric_kwargs=dict(),
-    aggregation="sum",  # or 'none' for not aggregating across sequence position
+    aggregation="sum",  # or "none" for not aggregating across sequence position
     nodes_only=False,
     parallel_attn=False,
     node_threshold=0.1,
-    edge_threshold=0.01,
 ):
     all_submods = ([embed] if embed is not None else []) + [
         submod for layer_submods in zip(attns, mlps, resids) for submod in layer_submods
@@ -109,7 +100,7 @@ def N(upstream, downstream, midstream=[]):
         edges[f"mlp_{layer}"][f"resid_{layer}"] = MR_effect
         edges[f"attn_{layer}"][f"resid_{layer}"] = AR_effect
 
-        if parallel_attn:
+        if not parallel_attn:
             AM_effect = N(attn, mlp)
             edges[f"attn_{layer}"][f"mlp_{layer}"] = AM_effect
 
@@ -160,13 +151,13 @@ def N(upstream, downstream, midstream=[]):
 
         # aggregate across batch dimension
         for child in edges:
-            bc, fc = nodes[child].act.shape
+            bc, _ = nodes[child].act.shape
             for parent in edges[child]:
                 weight_matrix = edges[child][parent]
                 if parent == "y":
                     weight_matrix = weight_matrix.sum(dim=0) / bc
                 else:
-                    bp, fp = nodes[parent].act.shape
+                    bp, _ = nodes[parent].act.shape
                     assert bp == bc
                     weight_matrix = weight_matrix.sum(dim=(0, 2)) / bc
                 edges[child][parent] = weight_matrix
@@ -465,8 +456,8 @@ def metric_fn(model):
         "google/gemma-2-2b": 26,
     }[args.model]
     parallel_attn = {
-        "EleutherAI/pythia-70m-deduped": False,
-        "google/gemma-2-2b": True,
+        "EleutherAI/pythia-70m-deduped": True,
+        "google/gemma-2-2b": False,
     }[args.model]
     include_embed = {
         "EleutherAI/pythia-70m-deduped": True,
@@ -478,7 +469,7 @@ def metric_fn(model):
     }[args.model]
 
     if args.model == "EleutherAI/pythia-70m-deduped":
-        model = LanguageModel(args.model, device_map=device, dispatch=True)
+        model = LanguageModel(args.model, device_map=device, dispatch=True, torch_dtype=dtype)
     elif args.model == "google/gemma-2-2b":
         model = LanguageModel(
             args.model,
@@ -498,17 +489,17 @@ def metric_fn(model):
         )
     else:
         data_path = f"data/{args.dataset}.json"
-        if args.aggregation == "sum":
-            raise NotImplementedError(
-                "Sum aggregation is not yet implemented for new data loading."
-            )
-            examples = load_examples(
-                data_path, args.num_examples, model, pad_to_length=args.example_length
-            )
-        else:
-            examples = load_examples(
-                data_path, args.num_examples, model, use_min_length_only=True
-            )
+        # if args.aggregation == "sum":
+        #     raise NotImplementedError(
+        #         "Sum aggregation is not yet implemented for new data loading."
+        #     )
+        #     examples = load_examples(
+        #         data_path, args.num_examples, model, pad_to_length=args.example_length
+        #     )
+        # else:
+        examples = load_examples(
+            data_path, args.num_examples, model, use_min_length_only=True
+        )
 
     num_examples = min([args.num_examples, len(examples)])
     if num_examples < args.num_examples:  # warn the user
@@ -616,7 +607,6 @@ def metric_fn(model):
                 nodes_only=args.nodes_only,
                 aggregation=args.aggregation,
                 node_threshold=args.node_threshold,
-                edge_threshold=args.edge_threshold,
                 parallel_attn=parallel_attn,
             )
 
@@ -688,11 +678,9 @@ def metric_fn(model):
             annotations=annotations,
             save_dir=f"{args.plot_dir}/{save_base}_node{args.node_threshold}_edge{args.edge_threshold}",
             gemma_mode=(args.model == "google/gemma-2-2b"),
+            parallel_attn=parallel_attn,
         )
     else:
-        raise NotImplementedError(
-            "Sum aggregation is not yet implemented for new data loading."
-        )
         plot_circuit(
             nodes,
             edges,
@@ -701,5 +689,7 @@ def metric_fn(model):
             edge_threshold=args.edge_threshold,
             pen_thickness=args.pen_thickness,
             annotations=annotations,
-            save_dir=f"{args.plot_dir}/{save_basename}_dict{args.dict_id}_node{args.node_threshold}_edge{args.edge_threshold}_n{num_examples}_agg{args.aggregation}",
+            save_dir=f"{args.plot_dir}/{save_base}_node{args.node_threshold}_edge{args.edge_threshold}_n{num_examples}_agg{args.aggregation}",
+            gemma_mode=(args.model == "google/gemma-2-2b"),
+            parallel_attn=parallel_attn,
         )