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..e0e5db99a0 --- /dev/null +++ b/tests/unit/model_bridge/test_stop_at_layer_guard.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +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) + + +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) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 4c224b8499..d5db979561 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 @@ -460,6 +461,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 @@ -2095,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. @@ -2135,26 +2149,18 @@ def forward( if start_at_layer is not None: input = self._setup_start_at_layer(input, start_at_layer) - # 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'." - ) - if hasattr(self, "blocks"): - effective_stop_at_layer = ( - len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer + "stop_at_layer requires a 'blocks' stack; this architecture " + "does not register one." ) - 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: @@ -2449,8 +2455,9 @@ 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"): + 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]; " @@ -2824,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