From 586ac39b42462ebfea6d65e2cd2fbcb24a0f4a4c Mon Sep 17 00:00:00 2001 From: kyleyinxu Date: Fri, 31 Jul 2026 01:42:40 -0400 Subject: [PATCH] Drop optional unused aliases --- demos/Jacobian_Lens_Demo.ipynb | 258 +----------------- tests/unit/test_optional_submodule.py | 32 +++ .../model_bridge/component_setup.py | 37 +++ 3 files changed, 73 insertions(+), 254 deletions(-) diff --git a/demos/Jacobian_Lens_Demo.ipynb b/demos/Jacobian_Lens_Demo.ipynb index d2053f351..c17bd1c94 100644 --- a/demos/Jacobian_Lens_Demo.ipynb +++ b/demos/Jacobian_Lens_Demo.ipynb @@ -121,7 +121,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "ff26ab5da82d4829b81bfd328fe48612", + "model_id": "d66ac557d99c4599bbf1fd8700ce3613", "version_major": 2, "version_minor": 0 }, @@ -1544,9 +1544,8 @@ "source": [ "### Boot Qwen3.5-4B through TransformerBridge\n", "\n", - "Use the same `device` and `model_dtype` as above. Some hybrid layers may emit Hook-alias\n", - "warnings for missing `attn.*` targets on linear-attention blocks; that's expected, and\n", - "J-lens only needs `blocks.{i}.hook_out`.\n" + "Use the same `device` and `model_dtype` as above. J-lens only needs residual\n", + "`blocks.{i}.hook_out` on every layer (full-attention and GatedDeltaNet alike).\n" ] }, { @@ -1558,7 +1557,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "1896e4ed8ad544c8b8a59f035bc5b207", + "model_id": "126582e447b24107b68a0cc34417d03b", "version_major": 2, "version_minor": 0 }, @@ -1569,252 +1568,6 @@ "metadata": {}, "output_type": "display_data" }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.0') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.0') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.0') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.0') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.0') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.1') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.1') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.1') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.1') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.1') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.4') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.4') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.4') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.4') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.4') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.5') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.5') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.5') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.5') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.5') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.6') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.6') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.6') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.6') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.6') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.8') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.8') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.8') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.8') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.8') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.9') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.9') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.9') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.9') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.9') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.10') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.10') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.10') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.10') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.10') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.12') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.12') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.12') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.12') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.12') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.13') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.13') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.13') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.13') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.13') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.14') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.14') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.14') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.14') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.14') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.16') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.16') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.16') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.16') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.16') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.17') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.17') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.17') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.17') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.17') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.18') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.18') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.18') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.18') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.18') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.20') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.20') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.20') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.20') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.20') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.21') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.21') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.21') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.21') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.21') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.22') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.22') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.22') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.22') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.22') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.24') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.24') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.24') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.24') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.24') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.25') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.25') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.25') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.25') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.25') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.26') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.26') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.26') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.26') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.26') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.28') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.28') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.28') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.28') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.28') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.29') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.29') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.29') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.29') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.29') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_in' -> 'attn.hook_attn_in' on BlockBridge(name='model.language_model.layers.30') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.30') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_q_input' -> 'attn.hook_q_input' on BlockBridge(name='model.language_model.layers.30') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_k_input' -> 'attn.hook_k_input' on BlockBridge(name='model.language_model.layers.30') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n", - "/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:501: UserWarning: Hook alias 'hook_v_input' -> 'attn.hook_v_input' on BlockBridge(name='model.language_model.layers.30') did not resolve; this hook will not be accessible.\n", - " getattr(module, \"_register_aliases\")()\n" - ] - }, { "name": "stdout", "output_type": "stream", @@ -1825,9 +1578,6 @@ ], "source": [ "# NBVAL_IGNORE_OUTPUT\n", - "# Hybrid Qwen3.5 emits expected Hook-alias warnings on linear-attn layers\n", - "# (aliases point at attn.* which is absent there). Harmless for J-lens,\n", - "# which only needs blocks.{i}.hook_out.\n", "QWEN_MODEL = \"Qwen/Qwen3.5-4B\"\n", "\n", "model = TransformerBridge.boot_transformers(\n", diff --git a/tests/unit/test_optional_submodule.py b/tests/unit/test_optional_submodule.py index 80610540c..0238503b3 100644 --- a/tests/unit/test_optional_submodule.py +++ b/tests/unit/test_optional_submodule.py @@ -2,6 +2,7 @@ import copy import logging +import warnings import pytest import torch @@ -157,6 +158,37 @@ def test_non_optional_raises(self): with pytest.raises(AttributeError): setup_submodules(block, adapter, model.layers[3]) + def test_prunes_hook_aliases_for_skipped_attn(self): + """Absent optional attn drops HT attn aliases; present layers keep them.""" + blocks = _setup_blocks(HybridModel(), AttnAdapter()) + attn_aliases = ( + "hook_attn_out", + "hook_attn_in", + "hook_q_input", + "hook_k_input", + "hook_v_input", + ) + for i in range(3): + for alias in attn_aliases: + assert alias in blocks[i].hook_aliases + for alias in attn_aliases: + assert alias not in blocks[3].hook_aliases + # Residual / MLP aliases are independent of attn and must survive. + assert "hook_resid_pre" in blocks[3].hook_aliases + assert "hook_mlp_out" in blocks[3].hook_aliases + + def test_pruned_aliases_do_not_warn_on_register(self): + blocks = _setup_blocks(HybridModel(), AttnAdapter()) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + blocks[3]._register_aliases() + unresolved = [ + w + for w in caught + if "did not resolve" in str(w.message) and "hook_attn" in str(w.message) + ] + assert unresolved == [] + # -- Tests: blocks_with() ----------------------------------------------------- diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index fc21f99ea..ba0141515 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -150,6 +150,43 @@ def setup_submodules( # Clean up so architecture_adapter traversal won't find stale entries for name in skipped_optional: component.submodules.pop(name, None) + if skipped_optional: + _prune_hook_aliases_for_skipped(component, skipped_optional) + + +def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None: + """Drop hook aliases whose target path starts at a skipped optional submodule. + + Hybrid layers share a BlockBridge template that declares HT aliases like + ``hook_attn_out -> attn.hook_out``. When ``attn`` is absent on a GatedDeltaNet + (or Mamba) layer, those aliases cannot resolve and would warn at + ``_register_aliases``. Prune them here so only inapplicable aliases disappear + on the affected layer — full-attention siblings keep theirs. + """ + aliases = getattr(component, "hook_aliases", None) + if not aliases: + return + skipped_set = set(skipped) + # Deepcopied blocks still share the class-level dict until mutated. + if aliases is type(component).hook_aliases: + aliases = dict(aliases) + component.hook_aliases = aliases + + def _first_segment(path: str) -> str: + return path.split(".", 1)[0] + + to_drop: list[str] = [] + for alias_name, target in aliases.items(): + if isinstance(target, list): + kept = [t for t in target if _first_segment(t) not in skipped_set] + if not kept: + to_drop.append(alias_name) + elif len(kept) != len(target): + aliases[alias_name] = kept + elif _first_segment(target) in skipped_set: + to_drop.append(alias_name) + for alias_name in to_drop: + aliases.pop(alias_name, None) def setup_components(