Repository navigation
fix(bridge): discover non-standard block lists in _enumerate_blocks - #1855
Conversation
_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 TransformerLensOrg#1791
koriyoshi2041
left a comment
koriyoshi2041
left a comment
There was a problem hiding this comment.
The fallback only checks the first element, so a non-block ModuleList whose first item happens to be a residual-stream component is accepted wholesale. On exact head a741e86e, ModuleList([OpaqueBlockBridge(...), nn.Linear(2, 2)]) makes _enumerate_blocks() return both entries, despite the documented “whose items are residual-stream blocks” contract. This can leak unrelated modules into every direct consumer of _enumerate_blocks(). The four focused tests pass locally; an all-items predicate plus a mixed-list regression should preserve the intended structural boundary.
…k-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.
|
Thanks @koriyoshi2041, good catch. I reproduced it on a741e86: ModuleList([OpaqueBlockBridge(...), nn.Linear(2, 2)]) returned both entries. The fallback now requires every item to carry the marker, and a mixed list is skipped whole rather than partly enumerated. Added a regression test covering both orderings (non-block first and last); it fails on the old check and passes now. |
|
Rechecked the updated head 569e738. The focused block-list discovery suite passes 5/5 locally, including both mixed-list orderings, and the all-items predicate keeps the fallback fail-closed. The original structural-boundary issue is resolved. |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for putting this together @LightWork666! Looks solid, just a couple comments below, let me know if you have any questions
| 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 |
There was a problem hiding this comment.
The fallback accepts a narrower set than the named path it extends: T5BlockBridge, AltUpBlockBridge, the T5Gemma decoder blocks, and DeepSeek V4's blocks are all enumerated today when registered under a known name, yet none of them set hook_out_is_single_residual_stream. The same stack under a new adapter's name silently enumerates zero blocks, now depending on the name chosen. Every item of a stack that setup_components builds carries is_list_item, which matches exactly what the named path accepts (the vision-layer lists that also set it live nested inside vision_encoder, not in the top-level _modules this loop scans). Can you investigate and decide whether discovery should accept what the named path accepts, rather than only single-stream blocks?
| assert result[6][1].name == "transformer.coda.0" | ||
|
|
||
|
|
||
| def test_enumerate_blocks_ignores_modulelists_without_the_marker(): |
There was a problem hiding this comment.
None of these tests pin the discovery criterion: swapping the hook_out_is_single_residual_stream check for is_list_item passes the whole file, so the deliberate exclusion of marker-less stacks is unguarded. A list of GeneralizedComponents without the marker, asserted either included or excluded, would pin whichever semantics you land on.
|
Hi @LightWork666! No rush on this, but the next release will be going up in the next couple days, let me know if you will have time to address these questions before then, if not it will have to wait for 4.4 |
|
Hi @jlarson4, thanks for the heads-up. Yes, I can address both before the release. I checked the first point and you're right: |
… 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.
|
Pushed
Checked the existing adapters: Raven is the only one with top-level block lists outside the known names, and the vision-layer lists are nested inside their encoder component, so nothing unintended is picked up. |
|
@LightWork666 Thanks for getting this together for me! Looks great, merging now |
Fixes #1791.
_enumerate_blocks()walked the fixed_BLOCK_LIST_ATTRStuple ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks"). An architecture that registers its block lists under any other name enumerated zero blocks, with no error. Raven/Huginn registersprelude,core_blockandcoda, so on a Raven bridgeblocks_with(),composition_layer_indices()andattn_head_labelssilently returned empty results. This is the same hardcoded-name fragility as #1769, on the enumeration side.Fix. The known names are still walked first, in the same fixed order. Afterwards, any other
nn.ModuleListin_moduleswhose items sethook_out_is_single_residual_streamis appended in registration order, so a new adapter's block lists are picked up without a name to add. Lists without the marker are ignored, so unrelated ModuleLists are not read as layers._modulesis read throughself.__dict__, matching_has_registered_blocks().Correction to the issue: I wrote that
all_composition_scores()already checks this marker. It does not. The existing consumer isJacobianLens.validate_model. The docstrings here describe it correctly.Scope. This fixes enumeration only. Whether Raven's fused-QKV attention exposes per-head
W_Q/W_K/W_V/W_Ofor the stacking helpers is separate and unchanged here.Testing
tests/unit/model_bridge/test_block_list_discovery.py(4 tests, no model download): Raven-shaped lists are discovered in order; marker-less ModuleLists are ignored; known names keep priority and order;composition_layer_indices()reports every attention block on a Raven-shaped bridge._enumerate_blocks()compared on 14 real bridges (12 tiny-checkpoint architectures, real gpt2, pythia-70m) and a tiny real T5 encoder-decoder. Same pairs, order and object identity in every case, and no extra ModuleLists were picked up.uv run mypy ., black and isort are clean.tests/unittier (not slow): 6323 passed, 5 failed. All 5 failures are worker crashes (macOS-arm64 SIGBUS) ingenerate()tests (test_bridge_generate_bos.py,test_generation_compatibility). They crash the same way on unmodifieddev-4.x(which tests crash varies from run to run), and_enumerate_blocks()is not reachable fromgenerate().Known limits and what is not verified
OpaqueBlockBridge.all_composition_scores()previously failed early with "No attention layers found". With Raven's blocks now enumerated, it proceeds further and may fail differently, since Raven's fused QKV attention may not exposeW_Q/W_K/W_V/W_O. That is not addressed here.hook_out_is_single_residual_stream. A block class that doesn't set it (asT5BlockBridgedoesn't) is only found if its list uses one of the known names. Multi-stream block lists (e.g. DeepSeek V4) are excluded on purpose.test_analysis_methods.pyskip locally (no goldens dataset). Thegenerate()tests that hit the macOS-arm64 SIGBUS crash identically on unmodifieddev-4.x.