[Bug Report] stop_at_layer silently no-ops on RavenArchitectureAdapter (Huginn depth-recurrent decoder)

[LightWork666]

Describe the bug

TransformerBridge.forward(..., stop_at_layer=N) (and anything built on it, e.g. input_to_embed()) silently does nothing on a Raven/Huginn bridge -- no error, no early stop. The forward pass runs to full completion as if stop_at_layer had never been passed.

Two independent gaps combine to cause this:

1. Top-level wiring only knows about self.blocks. In TransformerBridge.forward(), the stop_at_layer handling explicitly rejects a fixed list of known non-standard block-list names (L_blocks, H_blocks, encoder_blocks, decoder_blocks) with a clear NotImplementedError, then sets _stop_at_layer_idx on every block only if hasattr(self, "blocks"). RavenArchitectureAdapter.component_mapping has no "blocks" key at all -- Huginn's three block lists are prelude, core_block, and coda (matching its prelude / weight-tied recurrent core / coda structure). Raven isn't in the reject-list either, so the request silently falls through both checks: no error, and _stop_at_layer_idx never gets set on anything.
2. Even if it were set, OpaqueBlockBridge._check_stop_at_layer()'s layer-index regex wouldn't recognize it anyway. It only matches \.layers\.(\d+) or blocks\.(\d+); Raven's per-block names (transformer.core_block.0, transformer.prelude.0, transformer.coda.0) match neither. The standard BlockBridge class already has a more general _extract_layer_idx() ((?:^|\.)(?:blocks|h|layers)\.(\d+), covering GPT-2's h.N too) that OpaqueBlockBridge doesn't share or reuse -- and even that wouldn't match core_block/prelude/coda either way.

Surveyed every other architecture using OpaqueBlockBridge/SSMBlockBridge (mamba, mamba2, nemotron_h, rwkv7, zamba2): all name their list backbone.layers or model.layers, which matches fine. Raven is the only affected architecture currently in the adapter set.

Reproduction (no model download needed -- pure component-level and adapter-level verification)

import torch
import torch.nn as nn
from transformer_lens.model_bridge.generalized_components.opaque_block import OpaqueBlockBridge
from transformer_lens.model_bridge.exceptions import StopAtLayerException

class MockBlock(nn.Module):
    def forward(self, hidden_states, **kwargs):
        return hidden_states

# Exactly how component_setup.py names Raven's core_block list items.
bridge = OpaqueBlockBridge(name="transformer.core_block", submodules={})
bridge.name = "transformer.core_block.0"
bridge.set_original_component(MockBlock())
bridge._stop_at_layer_idx = 0

try:
    bridge(torch.randn(1, 3, 4))
    print("BUG: no StopAtLayerException raised")
except StopAtLayerException:
    print("correctly raised")
# -> prints "BUG: no StopAtLayerException raised"

from types import SimpleNamespace
from transformer_lens.model_bridge.supported_architectures.raven import RavenArchitectureAdapter

cfg = SimpleNamespace(d_model=8, d_head=4, n_heads=2, n_layers=1, n_ctx=16, d_vocab=32,
                       d_mlp=16, act_fn="gelu", normalization_type="RMS", default_prepend_bos=False)
adapter = RavenArchitectureAdapter(cfg)
print("blocks" in adapter.component_mapping)  # -> False

System Info

• Installed from source, dev-4.x
• OS-independent (pure Python logic, no weights involved); the real checkpoint (tomg-group-umd/huginn-0125, ~14GB) was not downloaded for this report -- the component-level and adapter-level checks above fully isolate and confirm the bug without it.

Additional context

This is the same general failure class as #1632/#1633 (stop_at_layer never stopping on the native Bridge) -- a different, non-standard block-list shape falling outside what the stop_at_layer wiring accounts for. HRM-Text's L_blocks/H_blocks and encoder-decoder's encoder_blocks/decoder_blocks already got an explicit reject-with-clear-error carve-out at some point; Raven's three-way prelude/core_block/coda split was missed.

Expected behaviour & fix pointers

At minimum, Raven should join the explicit reject-list so stop_at_layer raises a clear NotImplementedError instead of silently no-op'ing -- matching the existing L_blocks/H_blocks/encoder_blocks/decoder_blocks precedent. Actually supporting it would need the top-level wiring to walk all of prelude/core_block/coda (not just a single self.blocks), and OpaqueBlockBridge._check_stop_at_layer() to recognize those list names -- complicated further by core_block being a recurrent list (each of its 4 blocks runs N times per forward), so "stop at core_block layer i" would need a defined semantic for which recurrence step to stop at.

Acceptance:

• stop_at_layer on a Raven/Huginn bridge either works correctly or raises a clear, explicit error -- never silently no-ops
• input_to_embed() (which calls stop_at_layer=0 internally) does not silently return full logits instead of the post-embedding residual for Raven
• A fast regression test (matching the reproduction above, no checkpoint download) pins whichever behavior is chosen
• make unit-test and uv run mypy . pass

Checklist

• I have checked that there is no similar issue in the repo (required)

[jlarson4]

Thanks for reporting this @LightWork666! Great investigation and documentation. I want to slightly adjust the approach to fixing this. The reject-list is fragile, and we should move away from using it where possible. It requires a new entry every time someone adds a non-standard adapter, which creates an extra step that is very easy for other contributors and maintainers to miss.

There is an example of the solution I'd like to use about ten lines below, in the sibling start_at_layer path at transformer_bridge.py line 2370:

if not hasattr(self, "blocks"):
    raise NotImplementedError("start_at_layer requires a 'blocks' stack.")

Inverting stop_at_layer to that allowlist covers Raven and every future adapter. It also allows us to delete the four-name reject-list.

An additional note, don't use hasattr for the guard. TransformerBridge.__getattr__ falls through to original_model, so hasattr(self, "blocks") can read True when the wrapped HF model happens to expose a .blocks attribute. Components are registered via add_module, checking for "blocks" in self._modules is a more precise test. The existing start_at_layer guard carries the same latent weakness, and it would be great if it could be strengthened while this change is being made.

If you'd like to take this on, let me know and I can assign you!

In the future, it would be good to build out actual support for this via top-level wiring, but for now the rejection is good enough. Consider that a follow up, I will leave the issue open after the error is in place so that we don't lose track of that.