added test and made sure backwards hooks are working (#1058)

* added test and made sure backwards hooks are working

* fixed type issue

Author: bryce13950

tests/integration/model_bridge/test_bridge_integration.py

@@ -10,6 +10,7 @@
 import pytest
 import torch
 
+from transformer_lens import HookedTransformer
 from transformer_lens.ActivationCache import ActivationCache
 from transformer_lens.conversion_utils.conversion_steps.rearrange_hook_conversion import (
     RearrangeHookConversion,
@@ -673,5 +674,58 @@ def test_get_params_multi_query_attention_reshaping():
         original_attn.v.weight.data = original_v_weight
 
 
+def test_TransformerBridge_hooks_backward_hooks():
+    """Test that TransformerBridge.hooks() correctly registers backward hooks.
+
+    This test verifies that TransformerBridge.hooks() properly handles bwd_hooks
+    and registers them correctly, matching the behavior of HookedTransformer.hooks().
+    """
+    # Create both models with the same configuration
+    hooked_model = HookedTransformer.from_pretrained_no_processing("gpt2", device_map="cpu")
+    bridge_model: TransformerBridge = TransformerBridge.boot_transformers("gpt2", device="cpu")  # type: ignore
+    bridge_model.enable_compatibility_mode(no_processing=True)
+
+    # Create a simple backward hook that tracks if it was called
+    hook_called = {"hooked": False, "bridge": False}
+
+    def make_test_hook(model_type):
+        def hook_fn(grad, hook=None):
+            hook_called[model_type] = True
+            # For HookedTransformer, the hook doesn't modify the gradient
+            return None
+
+        return hook_fn
+
+    # Test input
+    test_input = torch.tensor([[1, 2, 3]])
+
+    # Test HookedTransformer - backward hooks should work
+    with hooked_model.hooks(bwd_hooks=[("blocks.0.hook_mlp_out", make_test_hook("hooked"))]):
+        output = hooked_model(test_input)
+        # Check that the backward hook was registered
+        assert (
+            len(hooked_model.blocks[0].hook_mlp_out.bwd_hooks) > 0
+        ), "HookedTransformer should register backward hooks"
+
+        # Trigger backward pass
+        output.sum().backward()
+
+    # Test TransformerBridge - backward hooks should now work correctly
+    # With compatibility mode, TransformerBridge should have the same hook names as HookedTransformer
+    with bridge_model.hooks(bwd_hooks=[("blocks.0.hook_mlp_out", make_test_hook("bridge"))]):
+        output = bridge_model(test_input)
+        # This assertion verifies that backward hooks are now properly registered
+        assert (
+            len(bridge_model.blocks[0].hook_mlp_out.bwd_hooks) > 0
+        ), "TransformerBridge should now register backward hooks correctly"
+
+        # Backward pass should trigger the hook
+        output.sum().backward()
+
+    # Verify the hooks were called appropriately
+    assert hook_called["hooked"], "HookedTransformer backward hook should have been called"
+    assert hook_called["bridge"], "TransformerBridge backward hook should now be called correctly"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

transformer_lens/model_bridge/bridge.py

@@ -1846,8 +1846,10 @@ def run_with_hooks(
         # Store hooks that we add so we can remove them later
         added_hooks: List[Tuple[HookPoint, str]] = []
 
-        def add_hook_to_point(hook_point: HookPoint, hook_fn: Callable, name: str):
-            hook_point.add_hook(hook_fn)
+        def add_hook_to_point(
+            hook_point: HookPoint, hook_fn: Callable, name: str, dir: Literal["fwd", "bwd"] = "fwd"
+        ):
+            hook_point.add_hook(hook_fn, dir=dir)
             added_hooks.append((hook_point, name))
 
         # Add stop_at_layer hook if specified
@@ -1868,10 +1870,11 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
                 block_hook_name = f"blocks.{last_layer_to_process}.hook_out"
                 hook_dict = self.hook_dict
                 if block_hook_name in hook_dict:
-                    add_hook_to_point(hook_dict[block_hook_name], stop_hook, block_hook_name)
+                    add_hook_to_point(hook_dict[block_hook_name], stop_hook, block_hook_name, "fwd")
 
         # Helper function to apply hooks based on name or filter function
         def apply_hooks(hooks: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool):
+            direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
             # Collect aliases for resolving legacy hook names
             aliases = collect_aliases_recursive(self)
 
@@ -1904,13 +1907,15 @@ def wrapped_hook_fn(tensor, hook):
                         actual_hook_name = aliases[hook_name_or_filter]
 
                     if actual_hook_name in hook_dict:
-                        add_hook_to_point(hook_dict[actual_hook_name], hook_fn, actual_hook_name)
+                        add_hook_to_point(
+                            hook_dict[actual_hook_name], hook_fn, actual_hook_name, direction
+                        )
                 else:
                     # Filter function
                     hook_dict = self.hook_dict
                     for name, hook_point in hook_dict.items():
                         if hook_name_or_filter(name):
-                            add_hook_to_point(hook_point, hook_fn, name)
+                            add_hook_to_point(hook_point, hook_fn, name, direction)
 
         try:
             # Apply forward hooks
@@ -2330,10 +2335,18 @@ def _hooks_context():
                 # Add forward hooks
                 for hook_name, hook_fn in fwd_hooks:
                     try:
-                        self.add_hook(hook_name, hook_fn)
+                        self.add_hook(hook_name, hook_fn, dir="fwd")
+                        added_hooks.append((hook_name, hook_fn))
+                    except Exception as e:
+                        print(f"Warning: Failed to add forward hook {hook_name}: {e}")
+
+                # Add backward hooks
+                for hook_name, hook_fn in bwd_hooks:
+                    try:
+                        self.add_hook(hook_name, hook_fn, dir="bwd")
                         added_hooks.append((hook_name, hook_fn))
                     except Exception as e:
-                        print(f"Warning: Failed to add hook {hook_name}: {e}")
+                        print(f"Warning: Failed to add backward hook {hook_name}: {e}")
 
                 yield