From a741e86e76289e24fd3166436eef6103ca2ce346 Mon Sep 17 00:00:00 2001 From: lightyagami6669 Date: Tue, 6 Oct 2026 10:42:31 +0530 Subject: [PATCH 1/3] fix(bridge): discover non-standard block lists in _enumerate_blocks _enumerate_blocks() walked a fixed tuple of block-list names, so an architecture that registers its blocks under any other name (Raven/Huginn's prelude, core_block and coda) enumerated zero blocks. blocks_with(), composition_layer_indices() and attn_head_labels then silently returned empty results instead of the real blocks or an error. After the known names, also pick up any other nn.ModuleList whose items set hook_out_is_single_residual_stream, in registration order. Known-name order and behavior are unchanged. Fixes #1791 --- .../model_bridge/test_block_list_discovery.py | 137 ++++++++++++++++++ .../model_bridge/transformer_bridge.py | 24 ++- 2 files changed, 160 insertions(+), 1 deletion(-) create mode 100644 tests/unit/model_bridge/test_block_list_discovery.py diff --git a/tests/unit/model_bridge/test_block_list_discovery.py b/tests/unit/model_bridge/test_block_list_discovery.py new file mode 100644 index 0000000000..8b5721c88a --- /dev/null +++ b/tests/unit/model_bridge/test_block_list_discovery.py @@ -0,0 +1,137 @@ +"""_enumerate_blocks() discovery of non-standard block-list names (#1791). + +``_BLOCK_LIST_ATTRS`` is a fixed tuple of known block-list attribute names +("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks"). +Architectures that register their blocks under any other name -- e.g. +Raven/Huginn's ``prelude`` / ``core_block`` / ``coda`` -- silently enumerated +zero blocks, which in turn made ``blocks_with()``, ``composition_layer_indices()``, +and ``attn_head_labels`` silently return empty results instead of the real +per-block data, with no error. + +``_enumerate_blocks()`` now falls back to structural discovery for any +``nn.ModuleList`` in ``_modules`` not already covered by the known names, +picking up ``GeneralizedComponent``s marked ``hook_out_is_single_residual_stream`` +(the marker ``JacobianLens.validate_model`` already uses to require a +single-stream block stack) -- so a new adapter's block lists are found +automatically. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from transformer_lens.model_bridge.generalized_components.linear import LinearBridge +from transformer_lens.model_bridge.generalized_components.opaque_block import ( + OpaqueBlockBridge, +) +from transformer_lens.model_bridge.transformer_bridge import TransformerBridge + +_BORROWED_HELPERS = ( + "_enumerate_blocks", + "blocks_with", + "_resolve_submodule_name", + "composition_layer_indices", + "_reject_encoder_decoder_composition", +) + + +def _stub(modules: dict) -> SimpleNamespace: + stub = SimpleNamespace(_modules=dict(modules), cfg=SimpleNamespace(n_heads=2)) + for helper in _BORROWED_HELPERS: + setattr(stub, helper, getattr(TransformerBridge, helper).__get__(stub)) + for name, module in modules.items(): + setattr(stub, name, module) + return stub + + +class _MockAttn(nn.Module): + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return hidden_states + + +def _opaque_block_with_attn(name: str) -> OpaqueBlockBridge: + """An OpaqueBlockBridge with a bridged "attn" submodule registered, matching + how a real adapter (e.g. Raven's) wires a residual-stream block with + attention -- _resolve_submodule_name only needs "attn" present in + block._modules, not a fully wired/functional submodule.""" + block = OpaqueBlockBridge(name=name, submodules={}) + block.set_original_component(_MockAttn()) + block.add_module("attn", LinearBridge(name="Wqkv", config=None)) + return block + + +def test_enumerate_blocks_discovers_non_standard_block_lists(): + """Raven's shape: three separate OpaqueBlockBridge lists named prelude / + core_block / coda -- none of them "blocks", "encoder_blocks", etc.""" + stub = _stub( + { + "prelude": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.prelude.{i}") for i in range(2)] + ), + "core_block": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.core_block.{i}") for i in range(4)] + ), + "coda": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.coda.{i}") for i in range(2)] + ), + } + ) + + result = stub._enumerate_blocks() + + assert [idx for idx, _ in result] == list(range(8)) + # Registration order preserved: prelude, then core_block, then coda. + assert result[0][1].name == "transformer.prelude.0" + assert result[2][1].name == "transformer.core_block.0" + assert result[6][1].name == "transformer.coda.0" + + +def test_enumerate_blocks_ignores_modulelists_without_the_marker(): + """A ModuleList that isn't a residual-stream block list (no + hook_out_is_single_residual_stream) must not be swept up as one -- + otherwise any incidental ModuleList on the bridge would get misread as + layers.""" + stub = _stub({"some_unrelated_list": nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])}) + + assert stub._enumerate_blocks() == [] + + +def test_enumerate_blocks_known_names_keep_priority_and_order(): + """Known names (here just "blocks") are still enumerated via the fixed + priority list, ahead of any structurally-discovered extra list, and the + existing decoder-only behavior is unchanged.""" + known = nn.ModuleList([_opaque_block_with_attn(f"blocks.{i}") for i in range(2)]) + extra = nn.ModuleList([_opaque_block_with_attn(f"extra.{i}") for i in range(2)]) + stub = _stub({"extra": extra, "blocks": known}) + + result = stub._enumerate_blocks() + + assert [block.name for _, block in result] == [ + "blocks.0", + "blocks.1", + "extra.0", + "extra.1", + ] + + +def test_composition_layer_indices_finds_raven_shaped_attention_blocks(): + """End-to-end through blocks_with()/composition_layer_indices(): previously + this silently returned [] for a Raven-shaped bridge; it must now report + every attention-bearing block.""" + stub = _stub( + { + "prelude": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.prelude.{i}") for i in range(2)] + ), + "core_block": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.core_block.{i}") for i in range(4)] + ), + "coda": nn.ModuleList( + [_opaque_block_with_attn(f"transformer.coda.{i}") for i in range(2)] + ), + } + ) + + assert stub.composition_layer_indices() == list(range(8)) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index d5db979561..7f92a6edff 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1292,14 +1292,36 @@ def _enumerate_blocks(self) -> List[Tuple[int, Any]]: plain layer indices. Encoder-decoder models register ``encoder_blocks`` and ``decoder_blocks`` instead; those are concatenated into one index space, matching ``HookedEncoderDecoder``'s ``chain(encoder, decoder)``. + + Non-standard architectures (e.g. Raven/Huginn's ``prelude`` / + ``core_block`` / ``coda``) register block lists under other names + entirely, so ``_BLOCK_LIST_ATTRS`` alone would silently enumerate zero + blocks for them (#1791). After the known names, any other + ``nn.ModuleList`` in ``_modules`` whose items are residual-stream + blocks -- ``GeneralizedComponent``s with ``hook_out_is_single_residual_stream`` + set, the marker ``JacobianLens.validate_model`` already uses to require a + single-stream block stack -- is discovered structurally and appended in + registration order, so a new adapter's block lists are picked up + automatically without a name to add here. """ pairs: List[Tuple[int, Any]] = [] + modules = self.__dict__.get("_modules") or {} + known_names = set(_BLOCK_LIST_ATTRS) for list_name in _BLOCK_LIST_ATTRS: - block_list = self._modules.get(list_name) + block_list = modules.get(list_name) if not isinstance(block_list, nn.ModuleList): continue for block in block_list: pairs.append((len(pairs), block)) + for list_name, block_list in modules.items(): + if list_name in known_names or not isinstance(block_list, nn.ModuleList): + continue + if len(block_list) == 0 or not getattr( + block_list[0], "hook_out_is_single_residual_stream", False + ): + continue + for block in block_list: + pairs.append((len(pairs), block)) return pairs def _resolve_submodule_name(self, block: Any, submodule: str) -> Optional[str]: From 569e73871b6148543ee52b09782ae93cb71ea674 Mon Sep 17 00:00:00 2001 From: lightyagami6669 Date: Tue, 6 Oct 2026 11:14:55 +0530 Subject: [PATCH 2/3] fix(bridge): require every item to be a residual-stream block in block-list discovery The structural fallback in _enumerate_blocks() only checked the first item of a ModuleList for hook_out_is_single_residual_stream, so a mixed list such as [block, nn.Linear] was accepted whole and the non-block module leaked into every consumer. Require all items to carry the marker; a mixed list is skipped instead of partly enumerated. Adds a regression test for mixed lists with the non-block first and last. --- tests/unit/model_bridge/test_block_list_discovery.py | 11 +++++++++++ transformer_lens/model_bridge/transformer_bridge.py | 9 +++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/unit/model_bridge/test_block_list_discovery.py b/tests/unit/model_bridge/test_block_list_discovery.py index 8b5721c88a..7f65db7686 100644 --- a/tests/unit/model_bridge/test_block_list_discovery.py +++ b/tests/unit/model_bridge/test_block_list_discovery.py @@ -98,6 +98,17 @@ def test_enumerate_blocks_ignores_modulelists_without_the_marker(): assert stub._enumerate_blocks() == [] +def test_enumerate_blocks_skips_mixed_lists_instead_of_leaking_non_blocks(): + """A ModuleList is only a block list if *every* item is a residual-stream + block. Checking just the first item let a mixed list through whole, leaking + unrelated modules into every consumer of _enumerate_blocks().""" + first_is_block = nn.ModuleList([_opaque_block_with_attn("mixed.0"), nn.Linear(2, 2)]) + last_is_block = nn.ModuleList([nn.Linear(2, 2), _opaque_block_with_attn("mixed2.1")]) + stub = _stub({"first_is_block": first_is_block, "last_is_block": last_is_block}) + + assert stub._enumerate_blocks() == [] + + def test_enumerate_blocks_known_names_keep_priority_and_order(): """Known names (here just "blocks") are still enumerated via the fixed priority list, ahead of any structurally-discovered extra list, and the diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 7f92a6edff..a8ff784db7 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1297,12 +1297,13 @@ def _enumerate_blocks(self) -> List[Tuple[int, Any]]: ``core_block`` / ``coda``) register block lists under other names entirely, so ``_BLOCK_LIST_ATTRS`` alone would silently enumerate zero blocks for them (#1791). After the known names, any other - ``nn.ModuleList`` in ``_modules`` whose items are residual-stream + ``nn.ModuleList`` in ``_modules`` whose items are *all* residual-stream blocks -- ``GeneralizedComponent``s with ``hook_out_is_single_residual_stream`` set, the marker ``JacobianLens.validate_model`` already uses to require a single-stream block stack -- is discovered structurally and appended in registration order, so a new adapter's block lists are picked up - automatically without a name to add here. + automatically without a name to add here. A list mixing blocks with + other modules is skipped whole rather than partly enumerated. """ pairs: List[Tuple[int, Any]] = [] modules = self.__dict__.get("_modules") or {} @@ -1316,8 +1317,8 @@ def _enumerate_blocks(self) -> List[Tuple[int, Any]]: for list_name, block_list in modules.items(): if list_name in known_names or not isinstance(block_list, nn.ModuleList): continue - if len(block_list) == 0 or not getattr( - block_list[0], "hook_out_is_single_residual_stream", False + if len(block_list) == 0 or not all( + getattr(block, "hook_out_is_single_residual_stream", False) for block in block_list ): continue for block in block_list: From c884c7d33fcdff5403ce04b9ff2a82cea287b40b Mon Sep 17 00:00:00 2001 From: lightyagami6669 Date: Fri, 23 Oct 2026 19:08:36 +0530 Subject: [PATCH 3/3] fix(bridge): discover block lists by is_list_item, matching the named path The structural fallback in _enumerate_blocks only accepted lists whose items set hook_out_is_single_residual_stream. The named path (_BLOCK_LIST_ATTRS) accepts every block type, including T5BlockBridge, AltUpBlockBridge, the T5Gemma decoder blocks and DeepSeek V4's blocks, none of which set that marker. The same stack under a new adapter's name therefore silently enumerated zero blocks. Accept lists whose items all set is_list_item instead, which is what setup_components puts on every block of a stack it builds. Mixed lists are still skipped whole, and Raven's prelude/core_block/coda lists are still discovered. Vision-layer lists also set is_list_item but are nested inside their encoder component, not top-level modules. Tests pin both sides of the criterion: a list of list-item components without the single-stream marker is included, and a list of components that are not list items is excluded. Swapping the check back to the marker now fails the new inclusion test. --- .../model_bridge/test_block_list_discovery.py | 53 +++++++++++++++---- .../model_bridge/transformer_bridge.py | 17 +++--- 2 files changed, 54 insertions(+), 16 deletions(-) diff --git a/tests/unit/model_bridge/test_block_list_discovery.py b/tests/unit/model_bridge/test_block_list_discovery.py index 7f65db7686..6748b93651 100644 --- a/tests/unit/model_bridge/test_block_list_discovery.py +++ b/tests/unit/model_bridge/test_block_list_discovery.py @@ -10,10 +10,11 @@ ``_enumerate_blocks()`` now falls back to structural discovery for any ``nn.ModuleList`` in ``_modules`` not already covered by the known names, -picking up ``GeneralizedComponent``s marked ``hook_out_is_single_residual_stream`` -(the marker ``JacobianLens.validate_model`` already uses to require a -single-stream block stack) -- so a new adapter's block lists are found -automatically. +picking up lists whose items are all ``is_list_item`` components (what +``setup_components`` sets on every block of a stack it builds, and exactly what +the named path accepts) -- so a new adapter's block lists are found +automatically, including block types that never set +``hook_out_is_single_residual_stream`` (T5, AltUp, T5Gemma, DeepSeek V4). """ from __future__ import annotations @@ -22,6 +23,9 @@ import torch import torch.nn as nn +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) from transformer_lens.model_bridge.generalized_components.linear import LinearBridge from transformer_lens.model_bridge.generalized_components.opaque_block import ( OpaqueBlockBridge, @@ -51,6 +55,14 @@ def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: return hidden_states +class _MultiStreamBlock(GeneralizedComponent): + """A list-item block that is not a single-residual-stream block: it sets + ``is_list_item`` but not ``hook_out_is_single_residual_stream``, like + ``T5BlockBridge``, ``AltUpBlockBridge`` and DeepSeek V4's blocks.""" + + is_list_item = True + + def _opaque_block_with_attn(name: str) -> OpaqueBlockBridge: """An OpaqueBlockBridge with a bridged "attn" submodule registered, matching how a real adapter (e.g. Raven's) wires a residual-stream block with @@ -88,16 +100,39 @@ def test_enumerate_blocks_discovers_non_standard_block_lists(): assert result[6][1].name == "transformer.coda.0" -def test_enumerate_blocks_ignores_modulelists_without_the_marker(): - """A ModuleList that isn't a residual-stream block list (no - hook_out_is_single_residual_stream) must not be swept up as one -- - otherwise any incidental ModuleList on the bridge would get misread as - layers.""" +def test_enumerate_blocks_ignores_modulelists_that_are_not_list_items(): + """A ModuleList whose items aren't list-item components must not be swept + up as a block list -- otherwise any incidental ModuleList on the bridge + would get misread as layers.""" stub = _stub({"some_unrelated_list": nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])}) assert stub._enumerate_blocks() == [] +def test_enumerate_blocks_accepts_list_items_without_the_single_stream_marker(): + """Pins the discovery criterion: ``is_list_item``, not + ``hook_out_is_single_residual_stream``. The named path enumerates T5 / AltUp / + T5Gemma / DeepSeek V4 blocks, none of which set the single-stream marker, so + the same stack under an unknown name must enumerate too instead of silently + yielding zero blocks.""" + items = [_MultiStreamBlock(name=f"layers.{i}") for i in range(3)] + assert not any(b.hook_out_is_single_residual_stream for b in items) + stub = _stub({"layers": nn.ModuleList(items)}) + + result = stub._enumerate_blocks() + + assert [block.name for _, block in result] == ["layers.0", "layers.1", "layers.2"] + + +def test_enumerate_blocks_ignores_components_that_are_not_list_items(): + """The other side of the criterion: a ``GeneralizedComponent`` list that does + not set ``is_list_item`` (e.g. a stack of non-layer helpers) is excluded.""" + items = [GeneralizedComponent(name=f"helpers.{i}") for i in range(2)] + stub = _stub({"helpers": nn.ModuleList(items)}) + + assert stub._enumerate_blocks() == [] + + def test_enumerate_blocks_skips_mixed_lists_instead_of_leaking_non_blocks(): """A ModuleList is only a block list if *every* item is a residual-stream block. Checking just the first item let a mixed list through whole, leaking diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index a8ff784db7..e5b5bb453c 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1297,13 +1297,16 @@ def _enumerate_blocks(self) -> List[Tuple[int, Any]]: ``core_block`` / ``coda``) register block lists under other names entirely, so ``_BLOCK_LIST_ATTRS`` alone would silently enumerate zero blocks for them (#1791). After the known names, any other - ``nn.ModuleList`` in ``_modules`` whose items are *all* residual-stream - blocks -- ``GeneralizedComponent``s with ``hook_out_is_single_residual_stream`` - set, the marker ``JacobianLens.validate_model`` already uses to require a - single-stream block stack -- is discovered structurally and appended in + ``nn.ModuleList`` in ``_modules`` whose items are *all* list-item + components (``is_list_item``, which ``setup_components`` sets on every + block of a stack it builds) is discovered structurally and appended in registration order, so a new adapter's block lists are picked up - automatically without a name to add here. A list mixing blocks with - other modules is skipped whole rather than partly enumerated. + automatically without a name to add here. This is the same criterion the + named path accepts, so block types that never set + ``hook_out_is_single_residual_stream`` (T5, AltUp, T5Gemma, DeepSeek V4) + are enumerated regardless of the name they are registered under. A list + mixing blocks with other modules is skipped whole rather than partly + enumerated. """ pairs: List[Tuple[int, Any]] = [] modules = self.__dict__.get("_modules") or {} @@ -1318,7 +1321,7 @@ def _enumerate_blocks(self) -> List[Tuple[int, Any]]: if list_name in known_names or not isinstance(block_list, nn.ModuleList): continue if len(block_list) == 0 or not all( - getattr(block, "hook_out_is_single_residual_stream", False) for block in block_list + getattr(block, "is_list_item", False) for block in block_list ): continue for block in block_list: