Merge remote-tracking branch 'origin/dev-3.x' into add_support_for_gpt_oss

Author: bryce13950

tests/integration/model_bridge/test_bridge_root_module_cache_compatibility.py

@@ -11,6 +11,8 @@
     },
 )
 
+bridge.enable_compatibility_mode(disable_warnings=False)
+
 act_names_in_cache = [
     # "hook_embed",
     # "hook_pos_embed",

tests/integration/model_bridge/test_cache_hook_equality.py

@@ -7,6 +7,7 @@
 
 prompt = "Hello World!"
 bridge = TransformerBridge.boot_transformers(MODEL, device="cpu")
+bridge.enable_compatibility_mode(disable_warnings=False)
 hooked_transformer = HookedTransformer.from_pretrained(MODEL, device="cpu")
 
 act_names_in_cache = [

tests/integration/model_bridge/test_hook_compatibility.py

@@ -71,6 +71,7 @@ def test_required_hooks_available(self, transformer_bridge):
         def hook_exists_on_model(model, hook_path: str) -> bool:
             """Check if a hook path exists on the model by traversing attributes."""
             parts = hook_path.split(".")
+            model.enable_compatibility_mode(disable_warnings=False)
             current = model
 
             try:

tests/unit/utilities/test_aliases.py

@@ -20,6 +20,7 @@ def test_resolve_existing_alias(self):
         mock_target = Mock()
         mock_hook = Mock()
         mock_target.actual_hook = mock_hook
+        mock_target.disable_warnings = False  # Ensure warnings are enabled
 
         hook_aliases = {"old_hook": "actual_hook"}
 

transformer_lens/model_bridge/bridge.py

@@ -65,6 +65,7 @@ def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: An
         self.adapter = adapter
         self.cfg = adapter.cfg
         self.tokenizer = tokenizer
+        self.compatibility_mode = False
 
         # Add device information to config from the loaded model
         if not hasattr(self.cfg, "device"):
@@ -196,6 +197,29 @@ def __str__(self) -> str:
         lines.extend(self._format_component_mapping(mapping, indent=1))
         return "\n".join(lines)
 
+    def enable_compatibility_mode(self, disable_warnings: bool = False) -> None:
+        """Enable compatibility mode for the bridge.
+
+        This sets up the bridge to work with legacy HookedTransformer components/hooks.
+        It will also disable warnings about the usage of legacy components/hooks if specified.
+
+        Args:
+            disable_warnings: Whether to disable warnings about legacy components/hooks
+        """
+        # Avoid circular import
+        from transformer_lens.utilities.bridge_components import (
+            apply_fn_to_all_components,
+        )
+
+        self.compatibility_mode = True
+
+        def set_compatibility_mode(component: Any) -> None:
+            """Set compatibility mode on a component."""
+            component.compatibility_mode = True
+            component.disable_warnings = disable_warnings
+
+        apply_fn_to_all_components(self, set_compatibility_mode)
+
     # ==================== TOKENIZATION METHODS ====================
 
     def to_tokens(
@@ -807,26 +831,28 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
             for hp, _ in hooks:
                 hp.remove_hooks()
 
-        # 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()}
-
-        # Create duplicate entries in cache
-        cache_items_to_add = {}
-        for cache_name, cached_value in cache.items():
-            # Check if this cache name should have an alias
-            for new_name, old_name in reverse_aliases.items():
-                if cache_name == new_name:
-                    cache_items_to_add[old_name] = cached_value
-                    break
-
-        # Add the aliased entries to the cache
-        cache.update(cache_items_to_add)
-
-        # 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]
+        if self.compatibility_mode == True:
+            # 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()}
+
+            # Create duplicate entries in cache
+            cache_items_to_add = {}
+            for cache_name, cached_value in cache.items():
+                # Check if this cache name should have an alias
+                for new_name, old_name in reverse_aliases.items():
+                    if cache_name == new_name:
+                        cache_items_to_add[old_name] = cached_value
+                        break
+
+            # Add the aliased entries to the cache
+            cache.update(cache_items_to_add)
+
+            # 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]
 
         if return_cache_object:
             cache_obj = ActivationCache(cache, self, has_batch_dim=not remove_batch_dim)

transformer_lens/model_bridge/generalized_components/base.py

@@ -27,6 +27,11 @@ class GeneralizedComponent(nn.Module):
     # Class attribute indicating whether this component represents a list item (like blocks)
     is_list_item: bool = False
 
+    # Compatibility mode that can be activated/deactivated for legacy components/hooks
+    compatibility_mode: bool = False
+    # Whether to disable warnings about deprecated hooks
+    disable_warnings: bool = False
+
     # Dictionary mapping deprecated hook names to their new equivalents
     # Subclasses can override this to define their own aliases
     hook_aliases: Dict[str, str] = {}
@@ -175,7 +180,7 @@ def _getattr_helper(self, name: str) -> Any:
 
         # Check if this is a deprecated hook alias
         resolved_hook = resolve_alias(self, name, self.hook_aliases)
-        if resolved_hook is not None:
+        if resolved_hook is not None and self.compatibility_mode == True:
             return resolved_hook
 
         # Avoid recursion by checking if we're looking for original_component
@@ -248,7 +253,7 @@ def __getattr__(self, name: str):
             # If we reach here, we can resolve the alias normally
             resolved_property = resolve_alias(self, name, self.property_aliases)
 
-            if resolved_property is not None:
+            if resolved_property is not None and self.compatibility_mode == True:
                 return resolved_property
 
         # If an internal call or no alias was found, just regularly get the attribute

transformer_lens/utilities/aliases.py

@@ -5,7 +5,9 @@
 
 
 def resolve_alias(
-    target_object: Any, requested_name: str, aliases: Dict[str, str]
+    target_object: Any,
+    requested_name: str,
+    aliases: Dict[str, str],
 ) -> Optional[Any]:
     """Resolve a hook alias to the actual hook object.
 
@@ -19,12 +21,14 @@ def resolve_alias(
     """
     if requested_name in aliases:
         target_name = aliases[requested_name]
-        warnings.warn(
-            f"Hook '{requested_name}' is deprecated and will be removed in a future version. "
-            f"Use '{target_name}' instead.",
-            FutureWarning,
-            stacklevel=3,  # Adjusted for utility function call
-        )
+
+        if target_object.disable_warnings == False:
+            warnings.warn(
+                f"Hook '{requested_name}' is deprecated and will be removed in a future version. "
+                f"Use '{target_name}' instead.",
+                FutureWarning,
+                stacklevel=3,  # Adjusted for utility function call
+            )
 
         target_name_split = target_name.split(".")
 

transformer_lens/utilities/bridge_components.py

@@ -0,0 +1,109 @@
+"""Utilities for traversing and applying functions to every component in a TransformerBridge model."""
+
+from typing import Any, Callable
+
+import torch.nn as nn
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+
+def collect_all_submodules_of_component(
+    model: TransformerBridge,
+    component: GeneralizedComponent,
+    submodules: dict,
+    block_prefix: str = "",
+) -> dict:
+    """Recursively collects all submodules of a component in a TransformerBridge model.
+    Args:
+        model: The TransformerBridge model to collect submodules from
+        component: The component to collect submodules from
+        submodules: A dictionary to populate with submodules (modified in-place)
+        block_prefix: Prefix for the block name, needed for components that are part of a block bridge
+    Returns:
+        Dictionary mapping submodule names to their respective submodules
+    """
+    for component_submodule in component.submodules.values():
+        submodules[block_prefix + component_submodule.name] = component_submodule
+
+        # If the component is a list item, we need to collect all submodules of the block bridge
+        if component_submodule.is_list_item:
+            submodules = collect_components_of_block_bridge(model, component_submodule, submodules)
+
+        # If the component has submodules, we need to collect them recursively
+        if component_submodule.submodules:
+            submodules = collect_all_submodules_of_component(
+                model, component_submodule, submodules, block_prefix
+            )
+    return submodules
+
+
+def collect_components_of_block_bridge(
+    model: TransformerBridge, component: GeneralizedComponent, components: dict
+) -> dict:
+    """Collects all components of a BlockBridge component.
+    Args:
+        model: The TransformerBridge model to collect components from
+        component: The BlockBridge component to collect components from
+        components: A dictionary to populate with components (modified in-place)
+    Returns:
+        Dictionary mapping component names to their respective components
+    """
+
+    # Retrieve the remote component list from the adapter (we need a ModuleList to iterate over)
+    remote_module_list = model.adapter.get_remote_component(model.original_model, component.name)
+
+    # Make sure the remote component is a ModuleList
+    if isinstance(remote_module_list, nn.ModuleList):
+        for block in remote_module_list:
+            components[block.name] = block
+            components = collect_all_submodules_of_component(model, block, components, block.name)
+    return components
+
+
+def collect_all_components(model: TransformerBridge, components: dict) -> dict:
+    """Collects all components in a TransformerBridge inside a dictionary.
+    The keys are the component names, and the values are the components themselves.
+    Args:
+        model: The TransformerBridge model to collect components from
+        components: A dictionary to populate with components (modified in-place)
+    Returns:
+        Dictionary mapping component names to their respective components
+    """
+
+    # Iterate through all components in component mapping
+    for component in model.adapter.get_component_mapping().values():
+        components[component.name] = component
+        components = collect_all_submodules_of_component(model, component, components)
+
+        # We need to enable compatibility mode for all different blocks of the component if the component is a list item
+        if component.is_list_item:
+            components = collect_components_of_block_bridge(model, component, components)
+    return components
+
+
+def apply_fn_to_all_components(
+    model: TransformerBridge,
+    fn: Callable[[GeneralizedComponent], Any],
+    components: dict | None = None,
+) -> dict[str, Any]:
+    """Applies a function to all components in the TransformerBridge model.
+    Args:
+        model: The TransformerBridge model to apply the function to
+        fn: The function to apply to each component
+        components: Optional dictionary of components to apply the function to, if None, all components are collected
+    Returns:
+        return_values: A dictionary mapping component names to the return values of the function
+    """
+
+    if components is None:
+        components = collect_all_components(model, {})
+
+    return_values = {}
+
+    for component in components.values():
+        return_values[component.name] = fn(component)
+
+    return return_values