From 9b31d4659446669fc45f8aaf4cc22dcca96134fd Mon Sep 17 00:00:00 2001 From: Arnav Bendre Date: Thu, 17 Sep 2026 11:31:59 +0000 Subject: [PATCH 1/3] fix: raise on stop_at_layer when no 'blocks' stack is registered Replace the stop_at_layer reject-list with an allowlist guard and check self._modules instead of hasattr for both stop_at_layer and start_at_layer. Addresses #1769 --- .../model_bridge/test_stop_at_layer_guard.py | 36 ++++++++++++++++ .../model_bridge/transformer_bridge.py | 41 +++++++++++-------- 2 files changed, 61 insertions(+), 16 deletions(-) create mode 100644 tests/unit/model_bridge/test_stop_at_layer_guard.py diff --git a/tests/unit/model_bridge/test_stop_at_layer_guard.py b/tests/unit/model_bridge/test_stop_at_layer_guard.py new file mode 100644 index 0000000000..2d2cc60189 --- /dev/null +++ b/tests/unit/model_bridge/test_stop_at_layer_guard.py @@ -0,0 +1,36 @@ +import pytest +import torch +import torch.nn as nn + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge + + +def _bare_bridge(**block_lists: nn.Module) -> TransformerBridge: + """A TransformerBridge with only the given block lists registered (no HF model).""" + bridge = TransformerBridge.__new__(TransformerBridge) + nn.Module.__init__(bridge) + bridge.cfg = TransformerBridgeConfig( + d_model=8, + d_head=4, + n_layers=1, + n_ctx=16, + d_vocab=32, + d_mlp=16, + n_heads=2, + architecture="RavenForCausalLM", + ) + for name, module in block_lists.items(): + bridge.add_module(name, module) + return bridge + + +def test_stop_at_layer_raises_without_blocks_stack() -> None: + """Raven-style prelude/core_block/coda lists must not silently ignore stop_at_layer.""" + bridge = _bare_bridge( + prelude=nn.ModuleList([nn.Identity()]), + core_block=nn.ModuleList([nn.Identity()]), + coda=nn.ModuleList([nn.Identity()]), + ) + with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"): + bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 0f555cab16..cf0702e70c 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -3,6 +3,7 @@ This module provides the bridge components that wrap remote model components and provide a consistent interface for accessing their weights and performing operations. """ + import inspect import logging import re @@ -452,6 +453,16 @@ def n_params_total(self) -> int: """ return self._n_params_total + def _has_registered_blocks(self) -> bool: + """Whether a ``blocks`` stack is registered as a submodule on this bridge. + + Checks ``_modules`` directly rather than ``hasattr``: ``__getattr__`` falls + through to the wrapped HF model, so ``hasattr(self, "blocks")`` can be True + for a model that merely exposes its own ``.blocks`` attribute. + """ + modules = self.__dict__.get("_modules") or {} + return "blocks" in modules + def __getattr__(self, name: str) -> Any: """Provide a clear error message for missing attributes.""" # Re-invoke original_model's property so its descriptive AttributeError @@ -2128,25 +2139,20 @@ def forward( if start_at_layer is not None: input = self._setup_start_at_layer(input, start_at_layer) + # change this to have an allowlist + # Set stop_at_layer flag on all blocks if requested if stop_at_layer is not None: - if ( - hasattr(self, "L_blocks") - or hasattr(self, "H_blocks") - or hasattr(self, "encoder_blocks") - or hasattr(self, "decoder_blocks") - ): + if not self._has_registered_blocks(): raise NotImplementedError( - "stop_at_layer is not supported on non-standard block list " - "names (L_blocks, H_blocks, encoder_blocks, decoder_blocks). " - "The bridge only supports stop_at_layer on 'blocks'." + "stop_at_layer requires a 'blocks' stack; this architecture " + "does not register one." ) - if hasattr(self, "blocks"): - effective_stop_at_layer = ( - len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer - ) - for block in self.blocks: - block._stop_at_layer_idx = effective_stop_at_layer + effective_stop_at_layer = ( + len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer + ) + for block in self.blocks: + block._stop_at_layer_idx = effective_stop_at_layer # Map HookedEncoderDecoder-style kwargs to HF-compatible names if "decoder_input" in kwargs: @@ -2441,8 +2447,11 @@ def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor "start_at_layer is only supported on the standard 'blocks' stack, " f"not {alt!r}." ) - if not hasattr(self, "blocks"): + + # guard for the block stack + if not self._has_registered_blocks(): raise NotImplementedError("start_at_layer requires a 'blocks' stack.") + if not (isinstance(input, torch.Tensor) and input.is_floating_point()): raise ValueError( "start_at_layer requires a residual-stream tensor [batch, pos, d_model]; " From ca4fa17a943120dede97b4ca1654f9e0d37cf31c Mon Sep 17 00:00:00 2001 From: Arnav Bendre Date: Fri, 18 Sep 2026 07:51:51 +0000 Subject: [PATCH 2/3] docs: note NotImplementedError in stop_at_layer docstring test: cover a wrapped model exposing its own .blocks Addresses #1769 --- .../model_bridge/transformer_bridge.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 7d582fdaea..d5db979561 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -2106,7 +2106,10 @@ def forward( output is discarded when block k swaps in the residual) but are excluded from ``run_with_cache`` output. Requires an HF model that accepts ``inputs_embeds``; only supported on the standard ``blocks`` stack. - stop_at_layer: Layer to stop forward pass at + stop_at_layer: Layer to stop forward pass at. Only supported on the + standard ``blocks`` stack; architectures that register no ``blocks`` + (e.g. Raven's ``prelude``/``core_block``/``coda``) raise + ``NotImplementedError`` rather than running to completion. pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3) and vision models (eg. ViT, DeiT). The tensor is passed directly to the underlying HuggingFace model. @@ -2146,9 +2149,6 @@ def forward( if start_at_layer is not None: input = self._setup_start_at_layer(input, start_at_layer) - - # change this to have an allowlist - # Set stop_at_layer flag on all blocks if requested if stop_at_layer is not None: if not self._has_registered_blocks(): @@ -2455,8 +2455,6 @@ def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor "start_at_layer is only supported on the standard 'blocks' stack, " f"not {alt!r}." ) - - # guard for the block stack if not self._has_registered_blocks(): raise NotImplementedError("start_at_layer requires a 'blocks' stack.") @@ -2833,18 +2831,22 @@ def _generate_tokens( temperature=temperature, freq_penalty=freq_penalty, repetition_penalty=repetition_penalty, - tokens=penalty_tokens - if _generate_from_embeds - else (decoder_tokens if is_encoder_decoder else current_tokens), + tokens=( + penalty_tokens + if _generate_from_embeds + else (decoder_tokens if is_encoder_decoder else current_tokens) + ), ).to(self.cfg.device) else: sampled_tokens = utils.sample_logits( final_logits, temperature=0.0, repetition_penalty=repetition_penalty, - tokens=penalty_tokens - if _generate_from_embeds - else (decoder_tokens if is_encoder_decoder else current_tokens), + tokens=( + penalty_tokens + if _generate_from_embeds + else (decoder_tokens if is_encoder_decoder else current_tokens) + ), ).to(self.cfg.device) # Freeze rows that finished on an earlier step so they stop emitting From c82e99e95c2843616dc33ff23b31dda1f66a73b0 Mon Sep 17 00:00:00 2001 From: Arnav Bendre Date: Fri, 18 Sep 2026 08:02:12 +0000 Subject: [PATCH 3/3] test: cover a wrapped model exposing its own .blocks Addresses #1769 --- .../unit/model_bridge/test_stop_at_layer_guard.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/model_bridge/test_stop_at_layer_guard.py b/tests/unit/model_bridge/test_stop_at_layer_guard.py index 2d2cc60189..e0e5db99a0 100644 --- a/tests/unit/model_bridge/test_stop_at_layer_guard.py +++ b/tests/unit/model_bridge/test_stop_at_layer_guard.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest import torch import torch.nn as nn @@ -34,3 +36,15 @@ def test_stop_at_layer_raises_without_blocks_stack() -> None: ) with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"): bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0) + + +def test_blocks_guard_ignores_unregistered_blocks_attribute() -> None: + """A wrapped HF model exposing `.blocks` must not satisfy the guard via __getattr__.""" + bridge = _bare_bridge() + bridge.__dict__["original_model"] = SimpleNamespace(blocks=[object()]) + assert hasattr(bridge, "blocks") # the trap: __getattr__ falls through to the HF model + assert not bridge._has_registered_blocks() + with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"): + bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0) + with pytest.raises(NotImplementedError, match="start_at_layer requires a 'blocks' stack"): + bridge.forward(torch.zeros(1, 3, 8), start_at_layer=0)