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 1131e5eee..ee8f68951 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 @@ -39,6 +39,219 @@ def hook(value, hook=None): return hook +@pytest.mark.parametrize("installer", ["hooks", "run_with_hooks", "run_with_cache"]) +@pytest.mark.parametrize("helper", ["hooks", "run_with_hooks", "run_with_cache"]) +@pytest.mark.parametrize("direction", ["fwd", "bwd"]) +@pytest.mark.parametrize("same_point", [False, True]) +def test_retained_hooks_survive_next_helper(installer, helper, direction, same_point): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + events = [] + callback = _record(events, "retained") + install_kwargs = {f"{direction}_hooks": [(name, callback)], "reset_hooks_end": False} + if installer == "hooks": + with bridge.hooks(**install_kwargs): + bridge(tokens).sum().backward() + elif installer == "run_with_hooks": + bridge.run_with_hooks(tokens, **install_kwargs).sum().backward() + else: + bridge.run_with_cache( + tokens, + names_filter=name, + incl_bwd=direction == "bwd", + return_type="loss" if direction == "bwd" else "logits", + reset_hooks_end=False, + ) + retained_fwd = list(hp.fwd_hooks) + retained_bwd = list(hp.bwd_hooks) + assert retained_fwd or retained_bwd + assert bridge.context_level == 0 + target = name if same_point else "blocks.1.hook_out" + temporary = dict(fwd_hooks=[(target, _noop)], bwd_hooks=[(target, _noop)]) + if helper == "hooks": + with bridge.hooks(**temporary): + bridge(tokens).sum().backward() + elif helper == "run_with_hooks": + bridge.run_with_hooks(tokens, **temporary).sum().backward() + else: + bridge.run_with_cache(tokens, names_filter=target, incl_bwd=True, return_type="loss") + assert hp.fwd_hooks == retained_fwd + assert hp.bwd_hooks == retained_bwd + count = len(events) + bridge(tokens).sum().backward() + if installer != "run_with_cache": + assert len(events) > count + assert bridge.context_level == 0 + bridge.reset_hooks() + assert not hp.has_hooks() + + +def test_retained_inner_scope_survives_outer_cleanup_and_next_scope(): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + with bridge.hooks(fwd_hooks=[(name, _noop)]): + with bridge.hooks(fwd_hooks=[(name, _noop)], reset_hooks_end=False): + pass + retained = hp.fwd_hooks[-1] + assert hp.fwd_hooks == [retained] + with bridge.hooks(fwd_hooks=[(name, _noop)]): + with bridge.hooks(fwd_hooks=[(name, _noop)]): + bridge(tokens) + assert hp.fwd_hooks == [retained] + + +@pytest.mark.parametrize("helper", ["hooks", "run_with_hooks", "run_with_cache"]) +def test_retained_hook_survives_later_helper_exception(helper): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + bridge.run_with_hooks(tokens, fwd_hooks=[(name, _noop)], reset_hooks_end=False) + retained = hp.fwd_hooks[0] + + def fail(value, hook): + raise RuntimeError("temporary failure") + + with pytest.raises(RuntimeError, match="temporary failure"): + if helper == "hooks": + with bridge.hooks(fwd_hooks=[(name, fail)]): + bridge(tokens) + elif helper == "run_with_hooks": + bridge.run_with_hooks(tokens, fwd_hooks=[(name, fail)]) + else: + with bridge.hooks(fwd_hooks=[(name, fail)], reset_hooks_end=False): + bridge.run_with_cache(tokens, names_filter=name) + assert retained in hp.fwd_hooks + assert bridge.context_level == 0 + + +def test_cache_setup_failure_removes_only_new_handles(monkeypatch): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + bridge.run_with_hooks(tokens, fwd_hooks=[(name, _noop)], reset_hooks_end=False) + retained = hp.fwd_hooks[0] + other = bridge.hook_dict["blocks.1.hook_out"] + original = other.add_hook + + def reject_backward(*args, **kwargs): + if kwargs.get("dir") == "bwd": + raise RuntimeError("registration failed") + return original(*args, **kwargs) + + monkeypatch.setattr(other, "add_hook", reject_backward) + with pytest.raises(RuntimeError, match="registration failed"): + bridge.run_with_cache( + tokens, names_filter=[name, "blocks.1.hook_out"], incl_bwd=True, return_type="loss" + ) + assert hp.fwd_hooks == [retained] + assert not hp.bwd_hooks + assert not other.has_hooks() + assert bridge.context_level == 0 + + +def test_retained_native_intervention_survives_unrelated_cache(): + from transformer_lens.config import TransformerBridgeConfig + from transformer_lens.model_bridge import TransformerBridge + + cfg = TransformerBridgeConfig( + d_model=16, + d_head=8, + n_heads=2, + n_layers=1, + n_ctx=8, + d_vocab=32, + d_mlp=32, + act_fn="gelu", + seed=2, + ) + bridge = TransformerBridge.boot_native(cfg).eval() + tokens = torch.tensor([[2, 3, 4]]) + with torch.no_grad(): + baseline = bridge(tokens) + controlled = bridge.run_with_hooks( + tokens, + fwd_hooks=[("hook_embed", lambda value, hook: torch.zeros_like(value))], + reset_hooks_end=False, + ) + during, _ = bridge.run_with_cache(tokens, names_filter="unembed.hook_out") + after = bridge(tokens) + assert (controlled - baseline).abs().max() > 1e-4 + torch.testing.assert_close(during, controlled) + torch.testing.assert_close(after, controlled) + bridge.reset_hooks() + with torch.no_grad(): + torch.testing.assert_close(bridge(tokens), baseline) + + +def test_scope_removes_only_owned_handles_at_its_level(): + bridge = _build() + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + with bridge.hooks(fwd_hooks=[(name, _noop)]): + hp.add_hook(_noop, level=bridge.context_level) + caller_handle = hp.fwd_hooks[-1] + assert hp.fwd_hooks == [caller_handle] + + +@pytest.mark.parametrize("helper", ["run_with_hooks", "run_with_cache"]) +@pytest.mark.parametrize("reset_hooks_end", [False, True]) +def test_stop_at_layer_preserves_retained_hook_without_leaking_stop_hook(helper, reset_hooks_end): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + baseline = bridge(tokens) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + bridge.run_with_hooks(tokens, fwd_hooks=[(name, _noop)], reset_hooks_end=False) + retained = hp.fwd_hooks[0] + if helper == "run_with_hooks": + stopped = bridge.run_with_hooks(tokens, stop_at_layer=1, reset_hooks_end=reset_hooks_end) + else: + stopped, _ = bridge.run_with_cache( + tokens, names_filter=name, stop_at_layer=1, reset_hooks_end=reset_hooks_end + ) + assert stopped.shape == (1, 4, 16) + assert retained in hp.fwd_hooks + assert len(hp.fwd_hooks) == 1 + int(helper == "run_with_cache" and not reset_hooks_end) + for point in bridge.hook_points(): + if point is not hp: + assert not point.has_hooks() + assert bridge.context_level == 0 + torch.testing.assert_close(bridge(tokens), baseline) + subsequent, _ = bridge.run_with_cache(tokens, names_filter=name) + torch.testing.assert_close(subsequent, baseline) + + +@pytest.mark.parametrize("helper", ["run_with_hooks", "run_with_cache"]) +def test_retained_mode_removes_stop_hook_after_forward_failure(helper): + bridge = _build() + tokens = torch.randint(0, 64, (1, 4)) + baseline = bridge(tokens) + name = "blocks.0.hook_out" + hp = bridge.hook_dict[name] + + def fail(value, hook): + raise RuntimeError("forward failed before stop") + + bridge.add_hook(name, fail) + retained = hp.fwd_hooks[0] + with pytest.raises(RuntimeError, match="forward failed before stop"): + if helper == "run_with_hooks": + bridge.run_with_hooks(tokens, stop_at_layer=1, reset_hooks_end=False) + else: + bridge.run_with_cache(tokens, names_filter=name, stop_at_layer=1, reset_hooks_end=False) + assert retained in hp.fwd_hooks + assert not bridge.hook_dict["blocks.1.hook_in"].has_hooks() + assert bridge.context_level == 0 + bridge.reset_hooks() + torch.testing.assert_close(bridge(tokens), baseline) + + class TestResetHooks: def test_removes_forward_and_backward_from_every_point(self): """Default direction='both' clears bwd hooks even on non-io points diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index 8acafd672..72d859416 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/transformer_lens/model_bridge/bridge_core.py @@ -27,7 +27,7 @@ from torch.nn import functional as F from transformer_lens.ActivationCache import ActivationCache -from transformer_lens.hook_points import HookPoint +from transformer_lens.hook_points import HookPoint, LensHandle from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.driver_protocol import to_torch from transformer_lens.model_bridge.exceptions import StopAtLayerException @@ -1138,6 +1138,17 @@ def cache_some( remove_batch_dim=remove_batch_dim, ) + @staticmethod + def _remove_temporary_hook( + hook_point: HookPoint, direction: Literal["fwd", "bwd"], handle: LensHandle + ) -> None: + """Remove an invocation's handle, not retained hooks at a reused lexical level.""" + if handle.is_permanent: + return + handle.hook.remove() + handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks + handles[:] = [existing for existing in handles if existing is not handle] + def hooks( self, fwd_hooks: List = [], @@ -1158,7 +1169,7 @@ def hooks( @contextmanager def _hooks_context() -> Iterator["BridgeCore"]: - added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = [] + added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = [] context_level = getattr(self, "context_level", 0) + 1 self.context_level = context_level @@ -1182,7 +1193,8 @@ def add_hook_to_point( ) else: hook_point.add_hook(hook_fn, dir=dir, level=context_level) - added_hooks.append((hook_point, dir)) + handles = hook_point.fwd_hooks if dir == "fwd" else hook_point.bwd_hooks + added_hooks.append((hook_point, dir, handles[-1])) def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None: direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd" @@ -1216,11 +1228,9 @@ def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None: yield self finally: try: - for hook_point, direction in added_hooks: + for hook_point, direction, handle in added_hooks: if reset_hooks_end: - # `level` keeps this to the hooks added above — hooks the - # caller attached before the context survive. - hook_point.remove_hooks(dir=direction, level=context_level) + self._remove_temporary_hook(hook_point, direction, handle) if clear_contexts: hook_point.clear_context() finally: @@ -1263,7 +1273,7 @@ def run_with_hooks( "run_with_hooks() got an unexpected keyword argument 'names_filter'; " "use run_with_cache(names_filter=...) to scope caching." ) - added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = [] + added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = [] effective_stop_layer = None if stop_at_layer is not None and hasattr(self, "blocks"): if stop_at_layer < 0: @@ -1308,7 +1318,8 @@ def add_hook_to_point( ) else: hook_point.add_hook(hook_fn, dir=dir, level=context_level) - added_hooks.append((hook_point, dir)) + handles = hook_point.fwd_hooks if dir == "fwd" else hook_point.bwd_hooks + added_hooks.append((hook_point, dir, handles[-1])) def apply_hooks( hook_list: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool @@ -1394,11 +1405,9 @@ def stop_hook(tensor: Any, *, hook: Any) -> Any: return output finally: try: - for hook_point, direction in added_hooks: + for hook_point, direction, handle in added_hooks: if reset_hooks_end: - # `level` keeps this to the hooks added above — hooks the caller - # attached before the call survive. - hook_point.remove_hooks(dir=direction, level=context_level) + self._remove_temporary_hook(hook_point, direction, handle) if clear_contexts: hook_point.clear_context() finally: @@ -1659,16 +1668,23 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor: ) context_level = getattr(self, "context_level", 0) + 1 self.context_level = context_level + added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = [] + stop_handle: Optional[LensHandle] = None try: for hp, name in hooks: hp.add_hook(make_cache_hook(name), level=context_level) + added_hooks.append((hp, "fwd", hp.fwd_hooks[-1])) if incl_bwd: hp.add_hook(make_grad_cache_hook(name), dir="bwd", level=context_level) + added_hooks.append((hp, "bwd", hp.bwd_hooks[-1])) 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) + stop_handle = stop_hook_point.fwd_hooks[-1] + added_hooks.append((stop_hook_point, "fwd", stop_handle)) except Exception: try: - self.remove_all_hook_fns(level=context_level) + for hp, direction, handle in added_hooks: + self._remove_temporary_hook(hp, direction, handle) finally: self.context_level -= 1 raise @@ -1723,10 +1739,10 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor: raise e finally: try: - if reset_hooks_end: - # `level` keeps this to the hooks added above — hooks the caller - # attached before the call survive. - self.remove_all_hook_fns(level=context_level) + for hp, direction, handle in added_hooks: + # Retain requested cache hooks, but never internal early-stop hooks. + if reset_hooks_end or handle is stop_handle: + self._remove_temporary_hook(hp, direction, handle) if clear_contexts: for hp, _ in hooks: hp.clear_context()