added cache layer for hook collection (#1032)

* added cache layer for hook collection

* added hook registry

* merged setattr

* fixed type issue

* made sure aliases were used during hook registration from generalized components

* resolved aliased hooks proplery

* resolved remaining hook alias issues

Author: bryce13950

tests/integration/model_bridge/test_bridge_integration.py

@@ -110,6 +110,9 @@ def test_cache():
     model_name = "gpt2"  # Use a smaller model for testing
     bridge = TransformerBridge.boot_transformers(model_name)
 
+    # Enable compatibility mode to include hook aliases
+    bridge.enable_compatibility_mode(disable_warnings=True)
+
     if bridge.tokenizer.pad_token is None:
         bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
 

tests/integration/model_bridge/test_qkv_hook_compatibility.py

@@ -120,72 +120,3 @@ def test_hook_aliases_work_correctly(self):
         assert qkv_bridge.q_hook_out is bridge.blocks[0].attn.q.hook_out, "Q property should work"
         assert qkv_bridge.k_hook_out is bridge.blocks[0].attn.k.hook_out, "K property should work"
         assert qkv_bridge.v_hook_out is bridge.blocks[0].attn.v.hook_out, "V property should work"
-
-    def test_head_ablation_hook_works_correctly(self):
-        """Test that head ablation hook works correctly with TransformerBridge."""
-        # Load GPT-2 in TransformerBridge
-        bridge = TransformerBridge.boot_transformers("gpt2", device="cpu")
-
-        # Turn on compatibility mode
-        bridge.enable_compatibility_mode(disable_warnings=True)
-
-        # Create test tokens (same as in the demo)
-        gpt2_tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
-
-        layer_to_ablate = 0
-        head_index_to_ablate = 8
-
-        # Test both hook names
-        hook_names_to_test = [
-            "blocks.0.attn.hook_v",  # Compatibility mode alias
-            "blocks.0.attn.v.hook_out",  # Direct property access
-        ]
-
-        for hook_name in hook_names_to_test:
-            print(f"\nTesting hook name: {hook_name}")
-
-            # Track if the hook was called
-            hook_called = False
-            mutation_applied = False
-
-            # We define a head ablation hook
-            def head_ablation_hook(value, hook):
-                nonlocal hook_called, mutation_applied
-                hook_called = True
-                print(f"Shape of the value tensor: {value.shape}")
-
-                # Apply the ablation (out-of-place to avoid view modification error)
-                result = value.clone()
-                result[:, :, head_index_to_ablate, :] = 0.0
-
-                # Check if the mutation was applied (the result should be zero for the ablated head)
-                if torch.all(result[:, :, head_index_to_ablate, :] == 0.0):
-                    mutation_applied = True
-
-                return result
-
-            # Get original loss
-            original_loss = bridge(gpt2_tokens, return_type="loss")
-
-            # Run with head ablation hook
-            ablated_loss = bridge.run_with_hooks(
-                gpt2_tokens, return_type="loss", fwd_hooks=[(hook_name, head_ablation_hook)]
-            )
-
-            print(f"Original Loss: {original_loss.item():.3f}")
-            print(f"Ablated Loss: {ablated_loss.item():.3f}")
-
-            # Assert that the hook was called
-            assert hook_called, f"Head ablation hook should have been called for {hook_name}"
-
-            # Assert that the mutation was applied
-            assert (
-                mutation_applied
-            ), f"Mutation should have been applied to the value tensor for {hook_name}"
-
-            # Assert that ablated loss is higher than original loss (ablation should hurt performance)
-            assert (
-                ablated_loss.item() > original_loss.item()
-            ), f"Ablated loss should be higher than original loss for {hook_name}"
-
-            print(f"✅ Hook {hook_name} works correctly!")

transformer_lens/model_bridge/bridge.py

@@ -35,7 +35,7 @@
 )
 from transformer_lens.model_bridge.hook_point_wrapper import HookPointWrapper
 from transformer_lens.model_bridge.types import ComponentMapping
-from transformer_lens.utilities.aliases import collect_aliases_recursive
+from transformer_lens.utilities.aliases import collect_aliases_recursive, resolve_alias
 
 if TYPE_CHECKING:
     from transformer_lens.ActivationCache import ActivationCache
@@ -54,6 +54,7 @@ class TransformerBridge(nn.Module):
     hook_aliases = {
         "hook_embed": "embed.hook_out",
         "hook_pos_embed": "pos_embed.hook_out",
+        "hook_unembed": "unembed.hook_out",
     }
 
     def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: Any):
@@ -70,6 +71,11 @@ def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: An
         self.cfg = adapter.cfg
         self.tokenizer = tokenizer
         self.compatibility_mode = False
+        self._hook_cache = None  # Cache for hook discovery results
+        self._hook_registry: Dict[
+            str, HookPoint
+        ] = {}  # Dynamic registry of hook names to HookPoints
+        self._hook_registry_initialized = False  # Track if registry has been initialized
 
         # Add device information to config from the loaded model
         if not hasattr(self.cfg, "device"):
@@ -84,68 +90,159 @@ def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: An
         # Set original components on the pre-created bridge components
         set_original_components(self, self.adapter, self.original_model)
 
-    @property
-    def hook_dict(self) -> dict[str, HookPoint]:
-        """Get all HookPoint objects in the model for compatibility with HookedTransformer."""
-        hooks = {}
-        visited = set()  # Move visited set outside the recursive function
-
-        def collect_hookpoints(module: nn.Module, prefix: str = "") -> None:
-            """Recursively collect all HookPoint objects."""
-            obj_id = id(module)
+        # Initialize hook registry after components are set up
+        self._initialize_hook_registry()
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Override setattr to track HookPoint objects dynamically."""
+        # Call parent setattr first
+        super().__setattr__(name, value)
+
+        # Check if this is a HookPoint being set
+        if isinstance(value, HookPoint):
+            # Set the name on the HookPoint
+            value.name = name
+            # Add to registry
+            self._hook_registry[name] = value
+        elif isinstance(value, HookPointWrapper):
+            # Handle HookPointWrapper objects
+            hook_in_name = f"{name}.hook_in"
+            hook_out_name = f"{name}.hook_out"
+            value.hook_in.name = hook_in_name
+            value.hook_out.name = hook_out_name
+            self._hook_registry[hook_in_name] = value.hook_in
+            self._hook_registry[hook_out_name] = value.hook_out
+        elif hasattr(value, "get_hooks") and callable(getattr(value, "get_hooks")):
+            # This is a GeneralizedComponent being set
+            # We need to register its hooks with the appropriate prefix
+            component_hooks = value.get_hooks()
+            for hook_name, hook in component_hooks.items():
+                full_name = f"{name}.{hook_name}"
+                hook.name = full_name
+                self._hook_registry[full_name] = hook
+
+    def _initialize_hook_registry(self) -> None:
+        """Initialize the hook registry by scanning existing components."""
+        if self._hook_registry_initialized:
+            return
+
+        # Scan existing components for hooks
+        self._scan_existing_hooks(self, "")
+
+        # Add bridge aliases if compatibility mode is enabled
+        if self.compatibility_mode and self.hook_aliases:
+            for alias_name, target_name in self.hook_aliases.items():
+                # Use the existing alias system to resolve the target hook
+                target_hook = resolve_alias(self, alias_name, self.hook_aliases)
+                if target_hook is not None:
+                    self._hook_registry[alias_name] = target_hook
+
+        self._hook_registry_initialized = True
+
+    def _scan_existing_hooks(self, module: nn.Module, prefix: str = "") -> None:
+        """Scan existing modules for hooks and add them to registry."""
+        visited = set()
+
+        def scan_module(mod: nn.Module, path: str = "") -> None:
+            obj_id = id(mod)
             if obj_id in visited:
                 return
             visited.add(obj_id)
 
-            for attr_name in dir(module):
-                if attr_name.startswith("_"):
-                    continue
-                # Skip original_component to avoid deep traversal
-                if attr_name == "original_component":
-                    continue
+            # Check if this is a GeneralizedComponent with its own hook registry
+            if hasattr(mod, "get_hooks") and callable(getattr(mod, "get_hooks")):
+                # Use the component's own hook registry
                 try:
-                    attr = getattr(module, attr_name)
+                    component_hooks = mod.get_hooks()  # type: ignore
+                    if isinstance(component_hooks, dict):
+                        # Type cast to help mypy understand this is a dict of hooks
+                        hooks_dict = cast(Dict[str, HookPoint], component_hooks)  # type: ignore
+                        for hook_name, hook in hooks_dict.items():  # type: ignore
+                            full_name = f"{path}.{hook_name}" if path else hook_name
+                            hook.name = full_name
+                            self._hook_registry[full_name] = hook
                 except Exception:
-                    continue
+                    # If get_hooks() fails, fall through to the else block
+                    pass
+            else:
+                # Fall back to scanning attributes for non-GeneralizedComponent modules
+                for attr_name in dir(mod):
+                    if attr_name.startswith("_"):
+                        continue
+                    if attr_name == "original_component":
+                        continue
 
-                name = f"{prefix}.{attr_name}" if prefix else attr_name
-                if isinstance(attr, HookPoint):
-                    # Set the name on the HookPoint so it can be used in caching
-                    attr.name = name
-                    hooks[name] = attr
-                elif hasattr(attr, "hook_in") and hasattr(attr, "hook_out"):
-                    # Handle HookPointWrapper objects
-                    if isinstance(attr, HookPointWrapper):
-                        # Add hook_in and hook_out from the wrapper
+                    try:
+                        attr = getattr(mod, attr_name)
+                    except Exception:
+                        continue
+
+                    name = f"{path}.{attr_name}" if path else attr_name
+
+                    if isinstance(attr, HookPoint):
+                        attr.name = name
+                        self._hook_registry[name] = attr
+                    elif isinstance(attr, HookPointWrapper):
                         hook_in_name = f"{name}.hook_in"
                         hook_out_name = f"{name}.hook_out"
                         attr.hook_in.name = hook_in_name
                         attr.hook_out.name = hook_out_name
-                        hooks[hook_in_name] = attr.hook_in
-                        hooks[hook_out_name] = attr.hook_out
-                elif isinstance(attr, nn.Module) and attr is not module:
-                    collect_hookpoints(attr, name)
-                elif isinstance(attr, (list, tuple)):
-                    for i, item in enumerate(attr):
-                        if isinstance(item, nn.Module):
-                            collect_hookpoints(item, f"{name}[{i}]")
-
-            # Also traverse named_children() to catch ModuleList and other containers
-            for child_name, child_module in module.named_children():
-                # Skip original_component and _original_component to avoid deep traversal
+                        self._hook_registry[hook_in_name] = attr.hook_in
+                        self._hook_registry[hook_out_name] = attr.hook_out
+                    elif isinstance(attr, nn.Module) and attr is not mod:
+                        scan_module(attr, name)
+                    elif isinstance(attr, (list, tuple)):
+                        for i, item in enumerate(attr):
+                            if isinstance(item, nn.Module):
+                                scan_module(item, f"{name}[{i}]")
+
+            # Check named children
+            for child_name, child_module in mod.named_children():
                 if child_name == "original_component" or child_name == "_original_component":
                     continue
-                child_path = f"{prefix}.{child_name}" if prefix else child_name
-                collect_hookpoints(child_module, child_path)
+                child_path = f"{path}.{child_name}" if path else child_name
+                scan_module(child_module, child_path)
+
+        scan_module(module, prefix)
+
+    @property
+    def hook_dict(self) -> dict[str, HookPoint]:
+        """Get all HookPoint objects in the model for compatibility with HookedTransformer."""
+        # Start with the current registry
+        hooks = self._hook_registry.copy()
+
+        # Add aliases if compatibility mode is enabled
+        if self.compatibility_mode:
+            for alias_name, target_name in self.hook_aliases.items():
+                if target_name in hooks:
+                    hooks[alias_name] = hooks[target_name]
 
-        collect_hookpoints(self, "")
         return hooks
 
+    def _discover_hooks(self) -> dict[str, HookPoint]:
+        """Get all HookPoint objects from the registry (deprecated, use hook_dict)."""
+        return self._hook_registry.copy()
+
+    def clear_hook_cache(self) -> None:
+        """Clear the cached hook discovery results (deprecated, kept for compatibility)."""
+        pass  # No longer needed since we don't use caching
+
+    def clear_hook_registry(self) -> None:
+        """Clear the hook registry and force re-initialization."""
+        self._hook_registry.clear()
+        self._hook_registry_initialized = False
+
     def __getattr__(self, name: str) -> Any:
         """Provide a clear error message for missing attributes."""
         if name in self.__dict__:
             return self.__dict__[name]
 
+        # Check if this is a hook alias when compatibility mode is enabled
+        if self.compatibility_mode and name in self.hook_aliases:
+            target_name = self.hook_aliases[name]
+            if target_name in self._hook_registry:
+                return self._hook_registry[target_name]
+
         return super().__getattr__(name)
 
     def _get_nested_attr(self, path: str) -> Any:
@@ -247,6 +344,10 @@ def set_compatibility_mode(component: Any) -> None:
 
         apply_fn_to_all_components(self, set_compatibility_mode)
 
+        # Re-initialize the hook registry to include aliases from components
+        self.clear_hook_registry()
+        self._initialize_hook_registry()
+
     # ==================== TOKENIZATION METHODS ====================
 
     def to_tokens(
@@ -738,56 +839,14 @@ def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
 
             return cache_hook
 
-        # Recursively collect all HookPoint objects
-        def collect_hookpoints(module: nn.Module, prefix: str = "") -> None:
-            obj_id = id(module)
-            if obj_id in visited:
-                return
-            visited.add(obj_id)
-
-            for attr_name in dir(module):
-                if attr_name.startswith("_"):
-                    continue
-                # Skip the original_model to avoid collecting hooks from HuggingFace model
-                if attr_name == "original_model" or attr_name == "original_component":
-                    continue
-                try:
-                    attr = getattr(module, attr_name)
-                except Exception:
-                    continue
-
-                def add_hook_to_list(hook: HookPoint, name: str):
-                    # Set the name on the HookPoint so it can be used in caching
-                    hook.name = name
-
-                    # Only add hook if it passes the names filter
-                    if names_filter_fn(name):
-                        hooks.append((hook, name))
-
-                name = f"{prefix}.{attr_name}" if prefix else attr_name
-                if isinstance(attr, HookPoint):
-                    add_hook_to_list(attr, name)
-                elif isinstance(attr, HookPointWrapper):
-                    # Add hooks for the wrapped hook points (hook_in and hook_out)
-                    add_hook_to_list(attr.hook_in, f"{name}.hook_in")
-                    add_hook_to_list(attr.hook_out, f"{name}.hook_out")
-                elif isinstance(attr, nn.Module):
-                    collect_hookpoints(attr, name)
-                elif isinstance(attr, (list, tuple)):
-                    for i, item in enumerate(attr):
-                        if isinstance(item, nn.Module):
-                            collect_hookpoints(item, f"{name}[{i}]")
-
-            # Also traverse named_children() to catch ModuleList and other containers
-            for child_name, child_module in module.named_children():
-                child_path = f"{prefix}.{child_name}" if prefix else child_name
-                # Skip the original_model module
-                if child_name == "original_model" or child_name == "original_component":
-                    continue
-                collect_hookpoints(child_module, child_path)
+        # Use cached hooks instead of re-discovering them
+        hook_dict = self.hook_dict
 
-        # Collect hooks from bridge components (these have the clean TransformerLens paths)
-        collect_hookpoints(self, "")
+        # Filter hooks based on names_filter
+        for hook_name, hook in hook_dict.items():
+            # Only add hook if it passes the names filter
+            if names_filter_fn(hook_name):
+                hooks.append((hook, hook_name))
 
         # Register hooks
         for hp, name in hooks: