From e4de522353e94ac0502c70dc5bfa6b6eddc6770f Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:16:04 +0800 Subject: [PATCH 1/2] Fix TransformerBridge temporary hook cleanup --- .../test_hook_orchestration_parity.py | 180 ++++++++++++++++++ transformer_lens/model_bridge/bridge_core.py | 102 ++++++---- 2 files changed, 250 insertions(+), 32 deletions(-) diff --git a/tests/unit/model_bridge/supported_architectures/test_hook_orchestration_parity.py b/tests/unit/model_bridge/supported_architectures/test_hook_orchestration_parity.py index 7cd17ba1f..1131e5eee 100644 --- a/tests/unit/model_bridge/supported_architectures/test_hook_orchestration_parity.py +++ b/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() diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index 620c7dab5..2bd1619be 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/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"]]] = [] + self.context_level += 1 + context_level = self.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) + self.context_level += 1 + context_level = self.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 + self.context_level += 1 + context_level = self.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(): From ede493b7a200bf8a1dbe146ffcc7d6908d30e382 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:56:12 +0800 Subject: [PATCH 2/2] Handle lightweight bridge hook contexts --- transformer_lens/model_bridge/bridge_core.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index 2bd1619be..b5797ebd6 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/transformer_lens/model_bridge/bridge_core.py @@ -930,8 +930,8 @@ def hooks( @contextmanager def _hooks_context() -> Iterator["BridgeCore"]: added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = [] - self.context_level += 1 - context_level = self.context_level + context_level = getattr(self, "context_level", 0) + 1 + self.context_level = context_level def add_hook_to_point( hook_point: HookPoint, @@ -1104,8 +1104,8 @@ 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) - self.context_level += 1 - context_level = self.context_level + 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: @@ -1359,8 +1359,8 @@ 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 - self.context_level += 1 - context_level = self.context_level + 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)