Fix TransformerBridge temporary hook cleanup

tests/unit/model_bridge/supported_architectures/test_hook_orchestration_parity.py

@@ -31,6 +31,14 @@ def _noop(value, hook=None):
     return value
 
 
+def _record(events, label):
+    def hook(value, hook=None):
+        events.append(label)
+        return value
+
+    return hook
+
+
 class TestResetHooks:
     def test_removes_forward_and_backward_from_every_point(self):
         """Default direction='both' clears bwd hooks even on non-io points
@@ -178,3 +186,175 @@ def recording(
         with bridge.hooks(fwd_hooks=[("blocks.0.hook_out", _noop)]):
             pass
         assert "blocks.0.hook_out" in seen
+
+
+class TestTemporaryHookScopes:
+    @pytest.mark.parametrize("helper", ["hooks", "run_with_hooks", "run_with_cache"])
+    def test_preexisting_forward_hook_survives_helper_cleanup(self, helper):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        events: list[str] = []
+        hook_point.add_hook(_record(events, "existing"))
+
+        if helper == "hooks":
+            with bridge.hooks(fwd_hooks=[(hook_name, _record(events, "temporary"))]):
+                bridge(tokens)
+        elif helper == "run_with_hooks":
+            bridge.run_with_hooks(
+                tokens,
+                fwd_hooks=[(hook_name, _record(events, "temporary"))],
+            )
+        else:
+            bridge.run_with_cache(tokens, names_filter=hook_name)
+
+        events_after_helper = events.copy()
+        assert len(hook_point.fwd_hooks) == 1
+
+        bridge(tokens)
+
+        assert events == events_after_helper + ["existing"]
+
+    def test_nested_context_removes_only_inner_hooks(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        events: list[str] = []
+
+        with bridge.hooks(fwd_hooks=[(hook_name, _record(events, "outer"))]):
+            with bridge.hooks(fwd_hooks=[(hook_name, _record(events, "inner"))]):
+                bridge(tokens)
+            assert len(hook_point.fwd_hooks) == 1
+            bridge(tokens)
+
+        assert events == ["outer", "inner", "outer"]
+        assert not hook_point.has_hooks()
+
+    def test_nested_run_with_cache_preserves_outer_hook(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        events: list[str] = []
+
+        with bridge.hooks(fwd_hooks=[(hook_name, _record(events, "outer"))]):
+            bridge.run_with_cache(tokens, names_filter=hook_name)
+            assert len(hook_point.fwd_hooks) == 1
+            bridge(tokens)
+
+        assert events == ["outer", "outer"]
+        assert not hook_point.has_hooks()
+
+    def test_run_with_cache_exception_preserves_preexisting_hook(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+
+        def raising_hook(value, hook=None):
+            raise RuntimeError("existing hook failed")
+
+        hook_point.add_hook(raising_hook)
+
+        with pytest.raises(RuntimeError, match="existing hook failed"):
+            bridge.run_with_cache(tokens, names_filter=hook_name)
+
+        assert len(hook_point.fwd_hooks) == 1
+
+    def test_preexisting_backward_hook_survives_context_cleanup(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        events: list[str] = []
+
+        def existing_hook(gradient, hook=None):
+            events.append("existing")
+
+        def temporary_hook(gradient, hook=None):
+            events.append("temporary")
+
+        hook_point.add_hook(existing_hook, dir="bwd")
+        with bridge.hooks(bwd_hooks=[(hook_name, temporary_hook)]):
+            bridge(tokens).sum().backward()
+
+        events_after_context = events.copy()
+        assert len(hook_point.bwd_hooks) == 1
+
+        bridge.zero_grad()
+        bridge(tokens).sum().backward()
+
+        assert events == events_after_context + ["existing"]
+
+    def test_run_with_cache_incl_bwd_preserves_preexisting_hooks(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        hook_point.add_hook(_noop, dir="fwd")
+        hook_point.add_hook(_noop, dir="bwd")
+
+        bridge.run_with_cache(
+            tokens,
+            names_filter=hook_name,
+            incl_bwd=True,
+            return_type="loss",
+        )
+
+        assert len(hook_point.fwd_hooks) == 1
+        assert len(hook_point.bwd_hooks) == 1
+
+    def test_permanent_hook_survives_temporary_scope_cleanup(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+        bridge.add_perma_hook(hook_name, _noop)
+
+        bridge.run_with_hooks(tokens, fwd_hooks=[(hook_name, _noop)])
+
+        assert len(hook_point.fwd_hooks) == 1
+        assert hook_point.fwd_hooks[0].is_permanent
+
+    @pytest.mark.parametrize("helper", ["hooks", "run_with_hooks"])
+    def test_reset_hooks_end_false_retains_temporary_hooks(self, helper):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        hook_name = "blocks.0.hook_out"
+        hook_point = bridge._hook_registry[hook_name]
+
+        if helper == "hooks":
+            with bridge.hooks(fwd_hooks=[(hook_name, _noop)], reset_hooks_end=False):
+                bridge(tokens)
+        else:
+            bridge.run_with_hooks(
+                tokens,
+                fwd_hooks=[(hook_name, _noop)],
+                reset_hooks_end=False,
+            )
+
+        assert bridge.context_level == 0
+        assert len(hook_point.fwd_hooks) == 1
+
+    def test_callable_filter_adds_one_hook_for_canonical_and_alias_names(self):
+        bridge = _build()
+        tokens = torch.randint(0, 64, (1, 4))
+        canonical_name = "blocks.0.mlp.hook_out"
+        alias_name = "blocks.0.hook_mlp_out"
+        hook_point = bridge.hook_dict[canonical_name]
+        events: list[str] = []
+
+        with bridge.hooks(
+            fwd_hooks=[
+                (
+                    lambda name: name in {canonical_name, alias_name},
+                    _record(events, "temporary"),
+                )
+            ]
+        ):
+            bridge(tokens)
+
+        assert events == ["temporary"]
+        assert not hook_point.has_hooks()

transformer_lens/model_bridge/bridge_core.py

@@ -88,6 +88,7 @@ def __init__(
         self._hook_registry_initialized = False
         self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}
         self._property_alias_registry: Dict[str, str] = {}
+        self.context_level = 0
         self._driver = driver
         if not hasattr(adapter, "component_mapping") or adapter.component_mapping is None:
             raise ValueError("Adapter must have a component_mapping attribute")
@@ -929,6 +930,8 @@ def hooks(
         @contextmanager
         def _hooks_context() -> Iterator["BridgeCore"]:
             added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = []
+            context_level = getattr(self, "context_level", 0) + 1
+            self.context_level = context_level
 
             def add_hook_to_point(
                 hook_point: HookPoint,
@@ -942,9 +945,14 @@ def add_hook_to_point(
                     if hook_point.name is not None:
                         alias_names_list.append(hook_point.name)
                     alias_names_list.append(name)
-                    hook_point.add_hook(hook_fn, dir=dir, alias_names=alias_names_list)
+                    hook_point.add_hook(
+                        hook_fn,
+                        dir=dir,
+                        level=context_level,
+                        alias_names=alias_names_list,
+                    )
                 else:
-                    hook_point.add_hook(hook_fn, dir=dir)
+                    hook_point.add_hook(hook_fn, dir=dir, level=context_level)
                 added_hooks.append((hook_point, dir))
 
             def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None:
@@ -973,9 +981,12 @@ def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None:
                 apply_hooks(bwd_hooks, False)
                 yield self
             finally:
-                if reset_hooks_end:
-                    for hook_point, direction in added_hooks:
-                        hook_point.remove_hooks(dir=direction)
+                try:
+                    if reset_hooks_end:
+                        for hook_point, direction in added_hooks:
+                            hook_point.remove_hooks(dir=direction, level=context_level)
+                finally:
+                    self.context_level -= 1
 
         return _hooks_context()
 
@@ -1045,25 +1056,16 @@ def add_hook_to_point(
                 if hook_point.name is not None:
                     alias_names_list.append(hook_point.name)
                 alias_names_list.append(name)
-                hook_point.add_hook(hook_fn, dir=dir, alias_names=alias_names_list)
+                hook_point.add_hook(
+                    hook_fn,
+                    dir=dir,
+                    level=context_level,
+                    alias_names=alias_names_list,
+                )
             else:
-                hook_point.add_hook(hook_fn, dir=dir)
+                hook_point.add_hook(hook_fn, dir=dir, level=context_level)
             added_hooks.append((hook_point, dir))
 
-        if stop_at_layer is not None and hasattr(self, "blocks"):
-            if stop_at_layer < 0:
-                stop_at_layer = len(self.blocks) + stop_at_layer
-            if stop_at_layer >= 0 and stop_at_layer < len(self.blocks):
-
-                def stop_hook(tensor: Any, *, hook: Any) -> Any:
-                    raise StopAtLayerException(tensor)
-
-                # Stop at the beginning of the specified block, not at the end of the previous block
-                block_hook_name = f"blocks.{stop_at_layer}.hook_in"
-                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, "fwd")
-
         def apply_hooks(
             hook_list: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool
         ) -> None:
@@ -1102,7 +1104,25 @@ def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
                             hook_name_to_use = hook_point.name if hook_point.name else n
                             add_hook_to_point(hook_point, hook_fn, hook_name_to_use, direction)
 
+        context_level = getattr(self, "context_level", 0) + 1
+        self.context_level = context_level
         try:
+            if stop_at_layer is not None and hasattr(self, "blocks"):
+                if stop_at_layer < 0:
+                    stop_at_layer = len(self.blocks) + stop_at_layer
+                if stop_at_layer >= 0 and stop_at_layer < len(self.blocks):
+
+                    def stop_hook(tensor: Any, *, hook: Any) -> Any:
+                        raise StopAtLayerException(tensor)
+
+                    # Stop at the beginning of the specified block, not at the end of the previous block
+                    block_hook_name = f"blocks.{stop_at_layer}.hook_in"
+                    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, "fwd"
+                        )
+
             apply_hooks(fwd_hooks, True)
             apply_hooks(bwd_hooks, False)
             if start_at_layer is not None:
@@ -1115,9 +1135,12 @@ def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
                 output = e.layer_output
             return output
         finally:
-            if reset_hooks_end:
-                for hook_point, direction in added_hooks:
-                    hook_point.remove_hooks(dir=direction)
+            try:
+                if reset_hooks_end:
+                    for hook_point, direction in added_hooks:
+                        hook_point.remove_hooks(dir=direction, level=context_level)
+            finally:
+                self.context_level -= 1
 
     # ---- high-level execution: run_with_cache ----
 
@@ -1202,6 +1225,8 @@ def create_names_filter_fn(filter_input):
         cache: Dict[str, torch.Tensor] = {}
         hooks: List[Tuple[HookPoint, str]] = []
         visited: set[int] = set()
+        stop_hook_point: Optional[HookPoint] = None
+        stop_hook_fn: Optional[Callable] = None
 
         # None → no-op .to(None), tensors stay on their current device.
         cache_device = kwargs.pop("device", None)
@@ -1283,10 +1308,6 @@ def grad_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
 
             return grad_hook
 
-        for hp, name in hooks:
-            hp.add_hook(make_cache_hook(name))
-            if incl_bwd:
-                hp.add_hook(make_grad_cache_hook(name), dir="bwd")
         processed_args = [input]
         # Driver-aware input placement: torch drivers move input_ids to the model's
         # device; remote drivers (no local parameters) leave them as-is.
@@ -1320,8 +1341,8 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
                 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
                 hook_dict = self.hook_dict
                 if block_hook_name in hook_dict:
-                    hook_dict[block_hook_name].add_hook(stop_hook)
-                    hooks.append((hook_dict[block_hook_name], block_hook_name))
+                    stop_hook_point = hook_dict[block_hook_name]
+                    stop_hook_fn = stop_hook
         filtered_kwargs = kwargs.copy()
         # ``cache_device`` is honored by ``make_cache_hook`` above (``tensor.detach().to(cache_device)``);
         # the model and inputs stay where the caller put them, matching ``ActivationCache.to``.
@@ -1338,6 +1359,21 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
             )
         if start_at_layer is not None:
             filtered_kwargs["start_at_layer"] = start_at_layer
+        context_level = getattr(self, "context_level", 0) + 1
+        self.context_level = context_level
+        try:
+            for hp, name in hooks:
+                hp.add_hook(make_cache_hook(name), level=context_level)
+                if incl_bwd:
+                    hp.add_hook(make_grad_cache_hook(name), dir="bwd", level=context_level)
+            if stop_hook_point is not None and stop_hook_fn is not None:
+                stop_hook_point.add_hook(stop_hook_fn, level=context_level)
+        except Exception:
+            try:
+                self.remove_all_hook_fns(level=context_level)
+            finally:
+                self.context_level -= 1
+            raise
         try:
             if (
                 "output_attentions" not in filtered_kwargs
@@ -1377,8 +1413,10 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
         except Exception as e:
             raise e
         finally:
-            for hp, _ in hooks:
-                hp.remove_hooks(dir="both" if incl_bwd else "fwd")
+            try:
+                self.remove_all_hook_fns(level=context_level)
+            finally:
+                self.context_level -= 1
         if self.compatibility_mode == True:
             reverse_aliases = {}
             for old_name, new_name in aliases.items():