map hook_pos_embed to rotary_emb, allow hook_aliases to be a list (#1034)

* map hook_pos_embed to rotary_emb, allow hook_aliases to be a list

* ran format

* removed neuroscope

* added full coverage for hook aliases change

* updated cache to accept list

---------

Co-authored-by: Bryce Meyer <[email protected]>

Author: hijohnnylin

.github/workflows/checks.yml

@@ -152,7 +152,7 @@ jobs:
           - "Exploratory_Analysis_Demo"
           # - "Grokking_Demo"
           # - "Head_Detector_Demo"
-          - "Interactive_Neuroscope"
+          # - "Interactive_Neuroscope"
           # - "LLaMA"
           # - "LLaMA2_GPU_Quantized"
           - "Main_Demo"

transformer_lens/model_bridge/bridge.py

@@ -51,9 +51,10 @@ class TransformerBridge(nn.Module):
 
     # Top-level hook aliases for legacy TransformerLens names
     # Placing these on the main bridge ensures aliases like 'hook_embed' are available
-    hook_aliases = {
+    hook_aliases: Dict[str, Union[str, List[str]]] = {
         "hook_embed": "embed.hook_out",
-        "hook_pos_embed": "pos_embed.hook_out",
+        # rotary style models use rotary_emb.hook_out, but gpt2-style models use pos_embed.hook_out
+        "hook_pos_embed": ["pos_embed.hook_out", "rotary_emb.hook_out"],
         "hook_unembed": "unembed.hook_out",
     }
 
@@ -131,11 +132,25 @@ def _initialize_hook_registry(self) -> None:
 
         # 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():
+            for alias_name, target 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
+                # Convert to Dict[str, str] for resolve_alias if target_name is a list
+                if isinstance(target, list):
+                    # For list targets, try each one until one works
+                    for single_target in target:
+                        try:
+                            target_hook = resolve_alias(
+                                self, alias_name, {alias_name: single_target}
+                            )
+                            if target_hook is not None:
+                                self._hook_registry[alias_name] = target_hook
+                                break
+                        except AttributeError:
+                            continue
+                else:
+                    target_hook = resolve_alias(self, alias_name, {alias_name: target})
+                    if target_hook is not None:
+                        self._hook_registry[alias_name] = target_hook
 
         self._hook_registry_initialized = True
 
@@ -213,9 +228,17 @@ def hook_dict(self) -> dict[str, HookPoint]:
 
         # 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]
+            for alias_name, target in self.hook_aliases.items():
+                # Handle both string and list target names
+                if isinstance(target, list):
+                    # For list targets, find the first one that exists in hooks
+                    for single_target in target:
+                        if single_target in hooks:
+                            hooks[alias_name] = hooks[single_target]
+                            break
+                else:
+                    if target in hooks:
+                        hooks[alias_name] = hooks[target]
 
         return hooks
 
@@ -239,9 +262,16 @@ def __getattr__(self, name: str) -> Any:
 
         # 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]
+            target = self.hook_aliases[name]
+            # Handle both string and list target names
+            if isinstance(target, list):
+                # For list targets, find the first one that exists in the registry
+                for single_target in target:
+                    if single_target in self._hook_registry:
+                        return self._hook_registry[single_target]
+            else:
+                if target in self._hook_registry:
+                    return self._hook_registry[target]
 
         return super().__getattr__(name)
 
@@ -1040,7 +1070,15 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
             # If compatibility mode is enabled, we need to handle aliases
             # Create duplicate cache entries for TransformerLens compatibility
             # Use the aliases collected from components (reverse mapping: new -> old)
-            reverse_aliases = {new_name: old_name for old_name, new_name in aliases.items()}
+            # Handle the case where some alias values might be lists
+            reverse_aliases = {}
+            for old_name, new_name in aliases.items():
+                if isinstance(new_name, list):
+                    # For list values, create a mapping for each item in the list
+                    for single_new_name in new_name:
+                        reverse_aliases[single_new_name] = old_name
+                else:
+                    reverse_aliases[new_name] = old_name
 
             # Create duplicate entries in cache
             cache_items_to_add = {}
@@ -1056,8 +1094,16 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
 
             # Add cache entries for all aliases (both hook and cache aliases)
             for alias_name, target_name in aliases.items():
-                if target_name in cache and alias_name not in cache:
-                    cache[alias_name] = cache[target_name]
+                # Handle both string and list target names
+                if isinstance(target_name, list):
+                    # For list targets, find the first one that exists in cache
+                    for single_target in target_name:
+                        if single_target in cache and alias_name not in cache:
+                            cache[alias_name] = cache[single_target]
+                            break
+                else:
+                    if target_name in cache and alias_name not in cache:
+                        cache[alias_name] = cache[target_name]
 
         if return_cache_object:
             cache_obj = ActivationCache(cache, self, has_batch_dim=not remove_batch_dim)

transformer_lens/utilities/aliases.py

@@ -30,17 +30,48 @@ def resolve_alias(
                 stacklevel=3,  # Adjusted for utility function call
             )
 
-        target_name_split = target_name.split(".")
-
-        if len(target_name_split) > 1:
-            current_attr = target_object
-            for i in range(len(target_name_split) - 1):
-                current_attr = getattr(current_attr, target_name_split[i])
-                next_attr = getattr(current_attr, target_name_split[i + 1])
-            return next_attr
+        def _resolve_single_target(target_name: str) -> Any:
+            """Helper function to resolve a single target name."""
+            target_name_split = target_name.split(".")
+            # there are multiple target names, so we need to check all of them
+            # this is the case for hook_pos_embed, which can be either pos_embed.hook_out (gpt2-style) or rotary_emb.hook_out (gemma/etc-style)
+            if len(target_name_split) > 1:
+                current_attr = target_object
+                for i in range(len(target_name_split) - 1):
+                    if not hasattr(current_attr, target_name_split[i]):
+                        continue
+                    current_attr = getattr(current_attr, target_name_split[i])
+
+                # Check if the final attribute exists
+                if not hasattr(current_attr, target_name_split[-1]):
+                    raise AttributeError(
+                        f"'{type(current_attr).__name__}' object has no attribute '{target_name_split[-1]}'"
+                    )
+                next_attr = getattr(current_attr, target_name_split[-1])
+                return next_attr
+            else:
+                # Check if the target attribute exists before getting it
+                if not hasattr(target_object, target_name):
+                    raise AttributeError(
+                        f"'{type(target_object).__name__}' object has no attribute '{target_name}'"
+                    )
+                # Return the target hook
+                return getattr(target_object, target_name)
+
+        # if the target_name is a list, we check all elements
+        if isinstance(target_name, list):
+            for target_name_item in target_name:
+                try:
+                    result = _resolve_single_target(target_name_item)
+                    return result
+                except AttributeError:
+                    continue
+            # If we get here, none of the targets in the list were found
+            raise AttributeError(
+                f"None of the target names {target_name} could be resolved on '{type(target_object).__name__}' object"
+            )
         else:
-            # Return the target hook
-            return getattr(target_object, target_name)
+            return _resolve_single_target(target_name)
     return None