transformer_lens.model_bridge.supported_architectures.recurrent_gemma module

RecurrentGemma (Griffin) architecture adapter.

Supports RecurrentGemmaForCausalLM (e.g. google/recurrentgemma-2b, google/recurrentgemma-9b), the open instance of the Griffin architecture.

Architecture overview: - Heterogeneous layers defined by config.block_types (default pattern

("recurrent", "recurrent", "attention") repeated over num_hidden_layers). Each model.layers.{i} is a RecurrentGemmaDecoderLayer whose temporal_block is either:

  • RecurrentGemmaRecurrentBlock — the RG-LRU real-gated linear recurrence (linear_x / linear_y / conv_1d / rg_lru / linear_out), or

  • RecurrentGemmaSdpaAttention — local sliding-window GQA attention with partial rotary (q_proj / k_proj / v_proj / o_proj / rotary_emb).

  • Every layer additionally has temporal_pre_norm (pre-norm before the temporal block), channel_pre_norm (pre-norm before the MLP), and a gated MLP mlp_block (gate_proj / up_proj / down_proj, GELU-tanh).

  • model.final_norm + lm_head.

Gemma-family numerics (shared with Gemma-2/3): RMSNorm applies (1.0 + weight) (rmsnorm_uses_offset), token embeddings are scaled by sqrt(hidden_size) at runtime inside the HF forward, and the final logits are tanh-soft-capped at config.logits_soft_cap (30.0).

Key adapter decision (mirrors Lfm2MoeArchitectureAdapter): because the temporal_block substructure varies per layer (recurrent vs. attention), we wrap each decoder layer as a whole with residual-stream hooks only, rather than pretending every layer has a homogeneous attention/MLP substructure. This keeps execution correct on both layer types. Finer-grained RG-LRU state hooks (hook_ssm_write / hook_ssm_state style) are a natural follow-up.

applicable_phases = [4]: the whole-layer bridge exposes only residual hooks, so the component-level comparisons in phases 1-3 do not apply; phase 4 (generation + text quality) does.

class transformer_lens.model_bridge.supported_architectures.recurrent_gemma.RecurrentGemmaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RecurrentGemmaForCausalLM (Griffin).

Hybrid RG-LRU recurrence + local sliding-window attention. The temporal-block type per layer is determined by config.block_types[layer_idx % len(block_types)].

__init__(cfg: Any) → None

Initialize the RecurrentGemma architecture adapter.

applicable_phases: list[int] = [4]
component_mapping: ComponentMapping | None
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None
class transformer_lens.model_bridge.supported_architectures.recurrent_gemma.RecurrentGemmaBlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)

Bases: BlockBridge

Whole-layer RecurrentGemma bridge exposing only residual-stream hooks.

RecurrentGemma interleaves RG-LRU recurrent layers and local-attention layers. Wrapping the HF decoder layer as a whole preserves correct execution while avoiding unresolved standard attention/MLP aliases on the recurrent layers (which have no q/k/v/o or gate/up/down substructure).

hook_aliases: Dict[str, str | List[str]] = {'hook_resid_post': 'hook_out', 'hook_resid_pre': 'hook_in'}