transformer_lens.model_bridge.supported_architectures package

Submodules

Module contents

Supported architecture adapters.

This module contains all the supported architecture adapters for different model architectures.

class transformer_lens.model_bridge.supported_architectures.AfmoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for AfmoeForCausalLM models.

__init__(cfg: Any) → None

Initialize the AFMoE architecture adapter.

supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.ApertusArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Apertus models.

Apertus uses a pre-norm architecture with RMSNorm, Q/K normalization in attention, rotary position embeddings (RoPE with LLaMA-3 scaling), grouped query attention (GQA), non-gated MLP (XiELU activation), and no biases on any projections.

Similar to Qwen3 (pre-norm RMSNorm, QK-norm, GQA, RoPE) but uses a non-gated MLP (up_proj -> XiELU -> down_proj) instead of gated MLP.

Note: Apertus uses different layer norm names than most Llama-family models: - attention_layernorm (instead of input_layernorm) - feedforward_layernorm (instead of post_attention_layernorm)

__init__(cfg: Any) → None

Initialize the Apertus architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch XIELUActivation to defer eager .item() calls for meta tensor compat.

Transformers v5 uses meta tensors during from_pretrained, but XIELUActivation.__init__ eagerly calls .item() on beta/eps buffers to precompute _beta_scalar/_eps_scalar for the CUDA kernel path. This fails on meta device. Once upstream fixes this (transformers PR #43473), this patch can be removed.

Instead of reimplementing __init__, we wrap it to catch the meta tensor failure and defer scalar computation to forward() time.

class transformer_lens.model_bridge.supported_architectures.ArceeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Arcee models (ArceeForCausalLM / AFM-4.5B).

Arcee is a Llama-style dense decoder: pre-norm RMSNorm, rotary position embeddings (RoPE), grouped query attention (GQA), and no biases on any projection. The single distinguishing feature is the MLP: an ungated feed-forward block (up_proj -> ReLU^2 -> down_proj) using the squared-ReLU activation (HF hidden_act = "relu2") instead of the gated SiLU/GeLU used by Llama. The post-activation neurons are exposed via the MLP bridge’s hook_post (mlp.out.hook_in), which is useful for inspecting the sparse activation structure ReLU^2 produces.

Structurally identical to Llama except for the ungated ReLU^2 MLP; unlike Apertus it uses standard input_layernorm / post_attention_layernorm names and has no Q/K normalization.

Optional Parameters (may not exist in state_dict):

Arcee models do NOT have biases on attention or MLP projections (attention_bias = false, mlp_bias = false):

  • blocks.{i}.attn.b_Q / b_K / b_V / b_O - No bias on attention projections

  • blocks.{i}.mlp.b_in - No bias on MLP input (up_proj)

  • blocks.{i}.mlp.b_out - No bias on MLP output (down_proj)

  • blocks.{i}.ln1.b / ln2.b / ln_final.b - RMSNorm has no bias

Weight processing handles these missing biases gracefully via ProcessWeights._safe_get_tensor().

__init__(cfg: Any) → None

Initialize the Arcee architecture adapter.

class transformer_lens.model_bridge.supported_architectures.AudioFlamingo3ArchitectureAdapter(cfg: Any)

Bases: Qwen2AudioArchitectureAdapter

Architecture adapter for AudioFlamingo3ForConditionalGeneration models.

class transformer_lens.model_bridge.supported_architectures.BD3LMArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for BD3LM (Block Diffusion LM, ICLR 2025).

BD3LM uses adaLN conditioning on diffusion timesteps, a custom Rotary embedding, joint QKV projections, and non-causal block-diffusion masking. Because adaLN modulation varies per-timestep, it cannot be folded into weights — the adapter uses DelegatedAttentionBlockBridge to delegate each DDiTBlock.forward() wholesale to the original HF module. Hooks fire at block boundaries and on mapped submodules.

applicable_phases: list[int] = [1, 2, 3]
convert_weights() → dict[str, Tensor]

Return empty dict — delegation means no weight rearrangement.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch BD3LM dynamic class before from_pretrained runs.

Modeling code has a custom __getattr__ that fails to delegate back to PreTrainedModel, raising AttributeError on all_tied_weights_keys.

prepare_model(hf_model: Any) → None

Patch BD3LM quirks that prevent standard bridge construction.

Three issues must be fixed before the bridge can wrap the model:

  1. vocab_embed is an nn.Parameter, not nn.Embedding, so it lacks a .weight attribute that EmbeddingBridge expects.

  2. The flex attention backend crashes on CPU; fall back to sdpa and regenerate block_diff_mask for the new backend.

  3. The HF forward() does not accept output_attentions and other kwargs the bridge unconditionally injects; patch at runtime because no other hook point allows filtering them before HF’s forward call.

supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.BaichuanArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Baichuan models (v1 and v2).

Baichuan uses combined QKV via W_pack (nn.Linear(h, 3*h)) with RoPE, RMSNorm, and gated MLP (SwiGLU). Per-layer rotary embeddings.

Optional Parameters (may not exist in state_dict):

Baichuan models do NOT have biases on any projection:

  • blocks.{i}.attn.b_Q / b_K / b_V / b_O — no bias

  • blocks.{i}.mlp.b_gate / b_in / b_out — no bias

  • blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch transformers v5 incompatibilities before from_pretrained runs.

prepare_model(hf_model: Any) → None

Fix rotary caches and normalize NormHead weights before bridge creation.

RotaryEmbedding differs between v1 and v2: - v1 (Baichuan-7B): inv_freq is a persistent buffer, loaded from the

checkpoint as bfloat16, but cos_cached/sin_cached are non-persistent and materialize as garbage under meta-init.

  • v2 (Baichuan2-*): inv_freq, cos_cached, sin_cached are all plain attributes (no register_buffer). v5’s meta-init materializes them on meta, and nothing in the checkpoint overwrites them.

Both cases are resolved by computing inv_freq + caches from scratch at float32 using config-derived head_dim and base=10000. Recomputing v1 at float32 is also an upgrade over its bfloat16 checkpoint values.

Baichuan2 Chat also uses NormHead which row-normalizes lm_head during forward. We apply that once here so the bridge sees the normalized weights directly without needing NormHead’s forward path.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Split fused W_pack QKV and optionally fold layer norms.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Inject per-layer rotary embedding for component testing.

class transformer_lens.model_bridge.supported_architectures.BambaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for BambaForCausalLM models.

Both mixers are mapped optional — each present only on its layer type. The Mamba-2 mixer is wired under the canonical .mixer slot (HF path .mamba) so SSM analyses reach it as on GraniteMoeHybrid / NemotronH.

__init__(cfg: Any) → None

Initialize the Bamba architecture adapter.

create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: Any) → Any

Unified DynamicCache carrying KV entries and SSM conv/recurrent state.

class transformer_lens.model_bridge.supported_architectures.BartArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for BartForConditionalGeneration models.

Post-LN with layernorm_embedding; checkpoints ship scale_embedding=False, so the family default-on is disabled.

force_scale_embedding: bool = False
has_layernorm_embedding: bool = True
class transformer_lens.model_bridge.supported_architectures.BertArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for BERT models.

__init__(cfg: Any) → None

Initialize the BERT architecture adapter.

Parameters:

cfg – The configuration object.

prepare_model(hf_model: Any) → None

Adjust component mapping based on the actual HF model variant.

BertForMaskedLM has cls.predictions (MLM head). BertForNextSentencePrediction has cls.seq_relationship (NSP head) and no MLM-specific LayerNorm.

supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.BitNetArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for BitNetForCausalLM models.

__init__(cfg: Any) → None

Initialize the BitNet architecture adapter.

applicable_phases: list[int] = [1, 2, 4]
class transformer_lens.model_bridge.supported_architectures.BlenderbotArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for BlenderbotForConditionalGeneration models.

has_final_stack_norm: bool = True
n_layers_from: str = 'decoder'
require_symmetric_layers: bool = False
class transformer_lens.model_bridge.supported_architectures.BloomArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Bloom models.

__init__(cfg: Any) → None

Initialize the Bloom architecture adapter.

split_qkv_matrix(original_attention_component: Any) → tuple[Linear, Linear, Linear]

Split the QKV matrix into separate linear transformations. :param attention_component: The original attention layer component

Returns:

Tuple of nn.Linear modules for Q, K, and V transformations

class transformer_lens.model_bridge.supported_architectures.CodeGenArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for CodeGen models.

CodeGen uses a parallel attention+MLP block (attn and MLP share the same LayerNorm input and their outputs are summed). The attention layer uses a fused qkv_proj weight whose layout follows GPT-J’s mp_num=4 tensor-parallel partitioning: the rows are interleaved as [Q_part, V_part, K_part] within each of the 4 MP partitions.

Optional Parameters (may be absent in some CodeGen checkpoints):

  • No bias on qkv_proj (fused QKV has no bias)

  • No bias on out_proj

  • No bias on mlp.fc_in or mlp.fc_out

__init__(cfg: Any) → None

Initialize the CodeGen architecture adapter.

split_qkv_matrix(attn_component: Any) → tuple[Linear, Linear, Linear]

Split the fused QKV weight into separate Q, K, V linear modules.

CodeGen uses GPT-J-style tensor-parallel partitioning with mp_num=4 partitions. Within each partition the row order is [Q_part, V_part, K_part], i.e. not the conventional Q/K/V order.

The fused weight has shape [3 * n_embd, n_embd]. We reshape to [mp_num, 3, local_dim, n_embd], extract the three slices, then flatten back to [n_embd, n_embd] for each of Q, K, V.

Parameters:

attn_component – The original CodeGenAttention module.

Returns:

Tuple of (q_linear, k_linear, v_linear) — three nn.Linear modules with no bias and weight shape [n_embd, n_embd].

class transformer_lens.model_bridge.supported_architectures.Cohere2ArchitectureAdapter(cfg: Any)

Bases: CohereArchitectureAdapter

Architecture adapter for Cohere2 / Command-A models.

Cohere2 keeps Cohere v1’s parallel block, LayerNorm, GQA, gated MLP and logit_scale behaviour, but interleaves sliding-window RoPE layers with full-attention NoPE layers. HF represents that either as an explicit layer_types list or as a legacy sliding_window_pattern integer.

__init__(cfg: Any) → None

Initialize the Cohere2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.CohereArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Cohere models (CohereForCausalLM).

Architectural quirks vs. standard decoder-only models: - Single input_layernorm per block; NO post_attention_layernorm.

Attention and MLP both read the SAME normed hidden states (parallel).

  • CohereLayerNorm is true LayerNorm (mean-subtracting), NOT RMSNorm. It has a weight parameter but NO bias parameter.

  • Logit scale: CohereForCausalLM.forward multiplies logits by logit_scale (default 0.0625 = 1/16). Folded into unembed.weight via preprocess_weights.

  • Rotary embeddings use repeat_interleave instead of cat-split (delegated to HF).

Optional parameters (absent from state_dict by default): - blocks.{i}.attn.b_Q/b_K/b_V/b_O — no bias on projections (attention_bias=False) - blocks.{i}.mlp.b_gate/b_in/b_out — no bias on MLP projections - blocks.{i}.ln1.b — CohereLayerNorm has no bias - ln_final.b — CohereLayerNorm has no bias

__init__(cfg: Any) → None

Initialize the Cohere architecture adapter.

apply_output_logits_transform(logits: Tensor) → Tensor

Match Cohere’s lm_head -> logit_scale -> optional softcap path.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Fold logit_scale into unembed weights before ProcessWeights runs.

bridge.py lines 726-732 clone unembed.weight before calling this, so scaling does not affect the tied embed.weight. logit_scale=1.0 is a no-op (skipped for efficiency).

class transformer_lens.model_bridge.supported_architectures.DeepSeekV2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for DeepSeek V2 / V2-Lite / Coder-V2 models.

Uses RMSNorm, MLA with compressed Q/KV projections (or direct Q projection when q_lora_rank is None), partial RoPE, MoE on most layers (dense MLP on first few), and no biases.

class transformer_lens.model_bridge.supported_architectures.DeepSeekV3ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for DeepSeek V3 / R1 models.

Uses RMSNorm, MLA with compressed Q/KV projections, partial RoPE, MoE on most layers (dense MLP on first few), and no biases.

class transformer_lens.model_bridge.supported_architectures.DreamArchitectureAdapter(cfg: Any)

Bases: Qwen2ArchitectureAdapter

Architecture adapter for DreamModel diffusion LMs.

applicable_phases: list[int] = [1, 2, 3, 4]
native_sampler: str = 'diffusion_generate'
native_sampler_kwargs(max_new_tokens: int, prompt_len: int) → dict

Dream denoises a fixed-length canvas; one step per token is its default ratio.

prepare_loading(model_name: str, model_kwargs: dict) → None

Shim the remote code’s two transformers-v4 dependencies.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.Emu3ArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for Emu3ForConditionalGeneration models.

class transformer_lens.model_bridge.supported_architectures.Ernie4_5ArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for Ernie4_5ForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.Ernie4_5_MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Ernie4_5_MoeForCausalLM models.

__init__(cfg: Any) → None

Initialize the ERNIE 4.5 MoE architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Exaone4ArchitectureAdapter(cfg: Any)

Bases: Olmo2ArchitectureAdapter

Architecture adapter for Exaone4ForCausalLM models.

__init__(cfg: Any) → None

Initialize the EXAONE 4.0 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.ExaoneArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for ExaoneForCausalLM (EXAONE-3.x) models.

The remote modeling code follows current HF conventions (Cache API, position_embeddings tuples), so the standard bridges delegate cleanly. Naming quirks: attention projections live one level deeper than usual (attn.attention.q_proj), the gated MLP uses c_fc_0 (gate) / c_fc_1 (up) / c_proj (down), and rotary sits at transformer.rotary.

__init__(cfg: Any) → None

Initialize the EXAONE architecture adapter.

prepare_model(hf_model: Any) → Any

Shim the EXAONE-3.x remote module for transformers >= 5.13, which renamed create_causal_mask’s input_embeds kwarg to inputs_embeds.

class transformer_lens.model_bridge.supported_architectures.FalconArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Falcon models (FalconForCausalLM).

prepare_model(hf_model: Any) → None

Patch Falcon modules to avoid backward hook conflicts.

Two issues: 1. FalconLinear does input @ self.weight.T where .T is a view —

clone the transpose to break the view chain.

  1. FalconDecoderLayer does mlp_output += attention_output (inplace) — this modifies a tensor captured by mlp.hook_out’s backward hook. Patch to use non-inplace addition.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Wire the shared rotary onto attention bridges (ALiBi variants skip).

class transformer_lens.model_bridge.supported_architectures.FalconH1ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for FalconH1ForCausalLM.

Parallel hybrid: every block runs GQA attention and a Mamba-2 mixer side by side, then a SwiGLU MLP. Both branches are mapped on every block so each sub-path is independently hookable for ablation studies.

applicable_phases: list[int] = []
apply_output_logits_transform(logits: Tensor) → Tensor

Match Falcon-H1’s post-unembedding multiplier.

class transformer_lens.model_bridge.supported_architectures.FalconMambaArchitectureAdapter(cfg: Any)

Bases: MambaArchitectureAdapter

Architecture adapter for FalconMambaForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.FlexOlmoArchitectureAdapter(cfg: Any)

Bases: Olmo2ArchitectureAdapter

Architecture adapter for FlexOlmoForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.Florence2ArchitectureAdapter(cfg: Any)

Bases: BartArchitectureAdapter

Architecture adapter for Florence2ForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Florence-2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.GPT2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GPT2 models.

Optional Parameters (may not exist in state_dict):

GPT-2 models HAVE biases on ALL linear layers:

✓ blocks.{i}.attn.b_Q - Has bias (from combined c_attn.bias) ✓ blocks.{i}.attn.b_K - Has bias (from combined c_attn.bias) ✓ blocks.{i}.attn.b_V - Has bias (from combined c_attn.bias) ✓ blocks.{i}.attn.b_O - Has bias (c_proj.bias) ✓ blocks.{i}.mlp.b_in - Has bias (c_fc.bias) ✓ blocks.{i}.mlp.b_out - Has bias (c_proj.bias) ✓ blocks.{i}.ln1.b - LayerNorm has bias ✓ blocks.{i}.ln2.b - LayerNorm has bias ✓ ln_final.b - LayerNorm has bias

No optional parameters - all biases exist in GPT-2.

__init__(cfg: Any) → None

Initialize the GPT2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.GPTBigCodeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GPTBigCode models.

GPTBigCode is a GPT-2 variant using Multi-Query Attention (MQA): a single fused c_attn projection whose output splits asymmetrically into [embed_dim, head_dim, head_dim] for Q/K/V (rather than three equal thirds). All other structure (module paths, LayerNorm, learned pos embeddings, standard MLP) is identical to GPT-2.

All public models use multi_query=True (1 KV head). The adapter assumes MQA throughout.

All linear layers have biases (c_attn, c_proj, c_fc, mlp.c_proj). lm_head has no bias and its weight is tied to transformer.wte.weight.

Weight layout difference from GPT-2: GPTBigCode uses nn.Linear (weights stored [out, in]) rather than GPT-2’s Conv1D ([in, out]), so no unembed weight transpose is needed.

class transformer_lens.model_bridge.supported_architectures.GPTOSSArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GPT-OSS model.

__init__(cfg: Any) → None

Initialize the GPT-OSS architecture adapter.

setup_hook_compatibility(bridge_model: Any) → None

Setup hook compatibility transformations for GPT-OSS models.

This configures rotary embedding references for attention layers, which is needed for models using RoPE (Rotary Position Embeddings).

This is called during Bridge.__init__ and should always be run.

Parameters:

bridge_model – The TransformerBridge instance

setup_no_processing_hooks(bridge_model: Any) → None

Backward compatibility alias for setup_hook_compatibility.

class transformer_lens.model_bridge.supported_architectures.Gemma1ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Gemma1 models.

__init__(cfg: Any) → None

Initialize the Gemma1 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Gemma2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Gemma2 models.

__init__(cfg: Any) → None

Initialize the Gemma2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Gemma3ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Gemma3 models.

__init__(cfg: Any) → None

Initialize the Gemma3 architecture adapter.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Wire local RoPE + eager attention; q/k norms delegate to HF autograd.

Gemma-3 uses dual RoPE (global + local); component tests share the local instance across all layers (layers on global RoPE accept the tradeoff).

class transformer_lens.model_bridge.supported_architectures.Gemma3MultimodalArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Gemma3 multimodal models (Gemma3ForConditionalGeneration).

This adapter handles vision-language models like Gemma 3 4B/12B/27B and MedGemma. The model structure is: - model.vision_tower: SigLIP vision encoder - model.multi_modal_projector: Projects vision embeddings to language space - model.language_model: Gemma3TextModel (same as text-only Gemma 3) - lm_head: Output projection

The language model component follows the same patterns as Gemma3ArchitectureAdapter.

__init__(cfg: Any) → None

Initialize the Gemma3 multimodal architecture adapter.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Wire rotary + eager, then enable native autograd on the Q/K norms.

class transformer_lens.model_bridge.supported_architectures.Gemma3nArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Text-only adapter for Gemma 3n (Gemma3nForConditionalGeneration).

applicable_phases: list[int] = [1, 2, 4]
required_libraries: list[str] = ['timm']
required_libraries_group: str = 'multimodal'
class transformer_lens.model_bridge.supported_architectures.Gemma4ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for Gemma 4 (Gemma4ForConditionalGeneration — multimodal, or Gemma4UnifiedForConditionalGeneration — text-only 12B).

applicable_phases: list[int] = [1, 2, 4]
class transformer_lens.model_bridge.supported_architectures.Gemma4TextArchitectureAdapter(cfg: Any)

Bases: Gemma4ArchitectureAdapter

Architecture adapter for Gemma4ForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.GiddArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GiddForDiffusionLM models.

__init__(cfg: Any) → None

Initialize the Gidd architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]
native_sampler: str = 'generate'
native_sampler_kwargs(max_new_tokens: int, prompt_len: int) → dict

Gidd’s max_length counts generated tokens: its windows start at prompt_length and span max_length, so adding the prompt over-generates.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch the remote class before from_pretrained runs.

Like BD3LM, the remote code’s attribute handling raises on v5’s all_tied_weights_keys lookup (the checkpoint is untied anyway).

prepare_model(hf_model: Any) → None

Restore the rotary table lost to meta-device loading.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention reads the rotary buffer inside HF; nothing to wire.

supports_fold_ln = False
supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.Glm4ArchitectureAdapter(cfg: Any)

Bases: GlmArchitectureAdapter

Architecture adapter for Glm4ForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.Glm4MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GLM-4.5 / 4.6 / 4.7 MoE decoder models.

GLM-4x MoE families use RMSNorm, RoPE and sparse routing, with early dense-MLP layers in some checkpoints. The dense layers are represented by a present-but-slightly-thinner mlp sub-module where routing is absent.

__init__(cfg: Any) → None

Initialize the GLM-4 MoE architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Glm4MoeLiteArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

GLM-4.7-Flash (Glm4MoeLiteForCausalLM) adapter: DeepSeek-V2 MLA + GLM-4-MoE routing (dense/sparse per mlp_layer_types).

class transformer_lens.model_bridge.supported_architectures.Glm4vArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Glm4vForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the GLM-4V architecture adapter.

required_libraries: list[str] = ['torchvision']
required_libraries_group: str = 'multimodal'
class transformer_lens.model_bridge.supported_architectures.GlmArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GlmForCausalLM models.

__init__(cfg: Any) → None

Initialize the GLM architecture adapter.

class transformer_lens.model_bridge.supported_architectures.GlmAsrArchitectureAdapter(cfg: Any)

Bases: Qwen2AudioArchitectureAdapter

Architecture adapter for GlmAsrForConditionalGeneration models.

class transformer_lens.model_bridge.supported_architectures.GlmMoeDsaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Z.ai GLM-5 / GLM-5.1 DSA models.

GLM-MoE-DSA combines MLA-style latent attention, a learned sparse-attention indexer, dense early MLP layers, and sparse MoE later layers.

class transformer_lens.model_bridge.supported_architectures.Gpt2LmHeadCustomArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GPT-2 LM Head Custom models.

__init__(cfg: Any) → None

Initialize the GPT-2 LM Head Custom architecture adapter.

class transformer_lens.model_bridge.supported_architectures.GptjArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for GPTJ models.

__init__(cfg: Any) → None

Initialize the GPTJ architecture adapter.

class transformer_lens.model_bridge.supported_architectures.GraniteArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for IBM Granite models (dense).

Granite is a Llama-like architecture with RMSNorm, rotary position embeddings (RoPE), GQA, and a gated MLP (SiLU activation). Granite-specific scaling multipliers are handled by the HF model’s native forward pass.

Optional Parameters (may not exist in state_dict):

Granite models do NOT have biases on attention and MLP projections:

  • blocks.{i}.attn.b_Q/b_K/b_V/b_O - No bias on attention projections

  • blocks.{i}.mlp.b_in/b_gate/b_out - No bias on MLP projections

  • blocks.{i}.ln1.b, blocks.{i}.ln2.b, ln_final.b - RMSNorm has no bias

__init__(cfg: Any) → None

Initialize the Granite architecture adapter.

apply_output_logits_transform(logits: Tensor) → Tensor

Match Granite’s lm_head / logits_scaling output path.

class transformer_lens.model_bridge.supported_architectures.GraniteMoeArchitectureAdapter(cfg: Any)

Bases: GraniteArchitectureAdapter

Architecture adapter for IBM Granite MoE models.

Identical to dense Granite but replaces the gated MLP with a Sparse Mixture of Experts block (block_sparse_moe) using batched expert parameters and top-k routing.

class transformer_lens.model_bridge.supported_architectures.GraniteMoeHybridArchitectureAdapter(cfg: Any)

Bases: GraniteArchitectureAdapter

Hybrid Mamba2 + Attention with Sparse MoE.

Attention is optional (absent on Mamba layers). shared_mlp and MoE are universal. Inherits Granite config and attention bridge construction.

applicable_phases: list[int] = [1, 2, 3, 4]
class transformer_lens.model_bridge.supported_architectures.HrmTextArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for HRM-Text (Sapient Intelligence).

Exposes L_blocks (fast/low-level stack) and H_blocks (slow/high-level stack) as sibling block lists. The nested recurrence loop is owned by HF’s forward; hooks fire once per iteration through the physical layers.

__init__(cfg: Any) → None

Initialize the HRM-Text architecture adapter.

applicable_phases: list[int] = [1, 2, 3]
setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for HRM-Text component testing.

HRM-Text uses RoPE. We set the rotary_emb reference on all attention bridge instances so component-level isolation tests can run.

supports_center_writing_weights = False
supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.HubertArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for HuBERT audio models.

HubertForCTC nests HubertModel under a ‘hubert.’ prefix; prepare_model() detects this and adjusts component paths.

prepare_loading(model_name: str, model_kwargs: dict) → None

Propagate HuBERT-specific HF config attributes to bridge config.

Prevents silent-default bugs where adapter reads from bridge config but the attribute was never propagated from HF config.

prepare_model(hf_model: Any) → None

Detect HubertForCTC (has ‘hubert.’ prefix) and add CTC head.

supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.HunYuanDenseV1ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for HunYuanDenseV1 models.

__init__(cfg: Any) → None

Initialize the HunYuanDenseV1 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.HyenaDNAArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for HyenaDNAForCausalLM models.

__init__(cfg: Any) → None

Initialize the HyenaDNA architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]
supports_batched_generation: bool = False
supports_fold_ln: bool = False
supports_generation: bool = True
supports_kv_cache: bool = False
class transformer_lens.model_bridge.supported_architectures.Idefics3ArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for Idefics3ForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Idefics3 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.InternLM2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for InternLM2 models.

InternLM2 uses remote code (trust_remote_code=True) and differs from Llama in: - Fused interleaved GQA wqkv weight (not standard [Q|K|V] split) - Non-standard module names: tok_embeddings, output, attention, feed_forward,

wqkv/wo, w1(gate)/w3(up)/w2(down), attention_norm, ffn_norm

  • Per-layer rotary_emb (no model-level shared instance)

  • supports_fold_ln=False: fold_ln is done manually in preprocess_weights because the bridge state dict has the fused qkv key, not split q/k/v keys, so fold_layer_norm’s extract_attention_tensors_for_folding would silently skip attn.

Optional parameters (may not exist in state_dict): - blocks.{i}.attn.b_Q / b_K / b_V / b_O — config.bias=False on shipped models - blocks.{i}.mlp.b_gate / b_in / b_out — MLP always bias=False - blocks.{i}.ln1.b / ln2.b / ln_final.b — RMSNorm has no bias

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch transformers v5 incompatibilities before from_pretrained runs.

prepare_model(hf_model: Any) → None

Restore per-layer rotary inv_freq lost to meta-device loading – this remote code predates HF’s original_inv_freq auto-restore, so positions would otherwise rotate by random values.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Fold layer norms into QKV and MLP weights.

Standard fold_ln can’t reach split Q/K/V when wqkv is fused in the bridge state dict. We extract and fold here, then write split keys so RearrangeTensorConversion can follow. MLP projections (w1/w2/w3) are separate linears so they fold normally. Mirrors phi3.py.preprocess_weights, adapted for InternLM2’s layout.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Inject per-layer rotary embedding for component testing.

class transformer_lens.model_bridge.supported_architectures.Jais2ArchitectureAdapter(cfg: Any)

Bases: NemotronArchitectureAdapter

Architecture adapter for Jais2ForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.JambaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for JambaForCausalLM.

Interleaved attention + Mamba-1 layers with optional sparse MoE FFN. Attention and Mamba streams are separate optional slots so each can be ablated independently.

applicable_phases: list[int] = [1, 2, 3, 4]
class transformer_lens.model_bridge.supported_architectures.JetMoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for JetMoeForCausalLM models.

__init__(cfg: Any) → None

Initialize the JetMoE architecture adapter.

component_test_skip_suffixes: tuple = ('mlp.gate', 'attn.experts.router')
setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.LEDArchitectureAdapter(cfg: Any)

Bases: BartArchitectureAdapter

Architecture adapter for LEDForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the LED architecture adapter.

class transformer_lens.model_bridge.supported_architectures.LLaDA2MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LLaDA2MoeModelLM models.

__init__(cfg: Any) → None

Initialize the LLaDA 2.0 MoE architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]
native_sampler: str = 'generate'
native_sampler_kwargs(max_new_tokens: int, prompt_len: int) → dict

gen_length must cover whole blocks; block_length caps at the budget.

prepare_loading(model_name: str, model_kwargs: dict) → None

Restore the v4 ‘default’ rope init the remote code looks up (Dream shim).

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

setup_hook_compatibility(bridge: Any) → None

Guard the remote forward against auto-passed 2D padding masks.

supports_fold_ln = False
supports_generation: bool = False
class transformer_lens.model_bridge.supported_architectures.LLaDAArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for the dense LLaDAModelLM architecture.

Support is deliberately limited to the released dense LLaDA block contract: Llama-style blocks, RMSNorm, separate bias-free projections, RoPE, bidirectional attention, an untied LM head, and no KV cache. The external iterative denoising/remasking loop is not a TransformerBridge generation API. Loading the Hugging Face checkpoint requires the caller to opt in with trust_remote_code=True.

applicable_phases: list[int] = []
prepare_loading(model_name: str, model_kwargs: dict) → None

Disable remote branches incompatible with single-pass hook support.

prepare_model(hf_model: Any) → None

Keep the wrapper and underlying model on the no-cache path.

supports_causal_loss: bool = False
supports_generation: bool = False
supports_hf_output_attentions: bool = False
class transformer_lens.model_bridge.supported_architectures.LagunaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LagunaForCausalLM models.

__init__(cfg: Any) → None

Initialize the Laguna architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

User-register Laguna’s native conversion mapping so the per-expert->batched expert merge runs under remote code (transformers skips it for custom-code modules, leaving the batched experts at random init).

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.Lfm2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Lfm2 models.

__init__(cfg: Any) → None

Initialize the Lfm2 architecture adapter.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up model-specific references for component testing.

class transformer_lens.model_bridge.supported_architectures.Lfm2MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LiquidAI LFM2 MoE models.

LFM2 MoE is a hybrid decoder with both short-convolution and full-attention layers. The adapter delegates each decoder layer to HF and exposes residual hooks around the whole layer rather than pretending every layer has a homogeneous attention/MLP substructure.

__init__(cfg: Any) → None

Initialize the LFM2 MoE architecture adapter.

applicable_phases: list[int] = [4]
class transformer_lens.model_bridge.supported_architectures.Llama4ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Llama4ForCausalLM models.

__init__(cfg: Any) → None

Initialize the Llama 4 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Llama4MultimodalArchitectureAdapter(cfg: Any)

Bases: Llama4ArchitectureAdapter

Architecture adapter for Llama4ForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Llama 4 multimodal architecture adapter.

class transformer_lens.model_bridge.supported_architectures.LlamaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Llama models.

Optional Parameters (may not exist in state_dict):

LLaMA models do NOT have biases on attention and MLP projections:

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP input (up_proj)

  • blocks.{i}.mlp.b_gate - No bias on MLP gate projection

  • blocks.{i}.mlp.b_out - No bias on MLP output (down_proj)

  • blocks.{i}.ln1.b - RMSNorm has no bias

  • blocks.{i}.ln2.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

Weight processing must handle these missing biases gracefully using ProcessWeights._safe_get_tensor() or by checking for None values.

__init__(cfg: Any) → None

Initialize the Llama architecture adapter.

class transformer_lens.model_bridge.supported_architectures.LlavaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LLava multimodal models (LlavaForConditionalGeneration).

This adapter handles vision-language models like LLava 1.5. The model structure is: - model.vision_tower: CLIP vision encoder - model.multi_modal_projector: 2-layer MLP (Linear -> GELU -> Linear) - model.language_model: LlamaForCausalLM

  • model.language_model.model.embed_tokens

  • model.language_model.model.layers[]: LLaMA transformer blocks

  • model.language_model.model.norm

  • model.language_model.lm_head

The language model component follows the same patterns as LlamaArchitectureAdapter.

__init__(cfg: Any) → None

Initialize the LLava architecture adapter.

class transformer_lens.model_bridge.supported_architectures.LlavaNextArchitectureAdapter(cfg: Any)

Bases: LlavaArchitectureAdapter

Architecture adapter for LLaVA-NeXT (1.6) models.

class transformer_lens.model_bridge.supported_architectures.LlavaOnevisionArchitectureAdapter(cfg: Any)

Bases: LlavaArchitectureAdapter

Architecture adapter for LLaVA-OneVision models.

prepare_model(hf_model: Any) → None

Fix weight tying when text_config and top-level config disagree.

Some checkpoints have tie_word_embeddings=True in text_config but False at the top level, leaving lm_head randomly initialized.

class transformer_lens.model_bridge.supported_architectures.LongT5ArchitectureAdapter(cfg: Any)

Bases: T5ArchitectureAdapter

Architecture adapter for LongT5ForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the LongT5 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.M2M100ArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for M2M100ForConditionalGeneration (M2M100 / NLLB) models.

has_final_stack_norm: bool = True
class transformer_lens.model_bridge.supported_architectures.MBartArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for MBartForConditionalGeneration models.

has_final_stack_norm: bool = True
has_layernorm_embedding: bool = True
class transformer_lens.model_bridge.supported_architectures.MPTArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

MPT adapter: ALiBi bias; all layers bias-free (no b_Q/b_K/b_V/b_O/b_in/b_out/ln bias).

validate_output_logits_transform() → None

Reject ambiguous remote-code logit scaling not used by integrated HF MPT.

class transformer_lens.model_bridge.supported_architectures.Mamba2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Wraps HF’s Mamba2ForCausalLM.

Differs from Mamba-1 at the mixer level: fused in_proj (no x_proj/dt_proj), two-input inner norm, multi-head structure with num_heads/head_dim/ n_groups, and an [num_heads]-shaped dt_bias. Shares SSMBlockBridge, DepthwiseConv1DBridge, and the stateful generation loop with Mamba-1.

applicable_phases: list[int] = [1, 2, 3, 4]
create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: dtype) → Any

Build a cache for the stateful generation loop.

class transformer_lens.model_bridge.supported_architectures.MambaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Wraps HF’s MambaForCausalLM. No attention, no positional embeddings.

SSM config fields (state_size, conv_kernel, expand, time_step_rank, intermediate_size) are propagated from the HF config via _HF_PASSTHROUGH_ATTRS in sources/transformers.py.

applicable_phases: list[int] = [1, 2, 3, 4]
create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: dtype) → Any

Build a cache for the stateful generation loop.

class transformer_lens.model_bridge.supported_architectures.MarianArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for MarianMTModel models (Helsinki-NLP opus-mt family).

class transformer_lens.model_bridge.supported_architectures.MingptArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for MinGPT models.

__init__(cfg: Any) → None

Initialize the MinGPT architecture adapter.

Parameters:

cfg – The configuration object.

class transformer_lens.model_bridge.supported_architectures.MiniMaxM2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for MiniMaxM2ForCausalLM models – Qwen3-MoE-like, but with full-width (not per-head) Q/K norm and a sigmoid + e_score_correction_bias router.

__init__(cfg: Any) → None

Initialize the MiniMax-M2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Ministral3ArchitectureAdapter(cfg: Any)

Bases: MistralArchitectureAdapter

Architecture adapter for Ministral3ForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.Mistral3ArchitectureAdapter(cfg: Any)

Bases: LlavaArchitectureAdapter

Architecture adapter for Mistral3ForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Mistral 3 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.MistralArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Mistral models.

__init__(cfg: Any) → None

Initialize the Mistral architecture adapter.

class transformer_lens.model_bridge.supported_architectures.MixtralArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Mixtral models.

Mixtral uses a pre-norm architecture with RMSNorm, rotary position embeddings (RoPE), and a Sparse Mixture of Experts MLP. Key features:

  • Pre-norm: RMSNorm applied BEFORE attention and BEFORE MLP.

  • Rotary embeddings: stored at model.rotary_emb and passed per-forward-call.

  • Sparse MoE: batched expert parameters (gate_up_proj, down_proj as 3D tensors).

  • MixtralAttention.forward() requires position_embeddings and attention_mask args.

  • Optional GQA (n_key_value_heads may differ from n_heads).

__init__(cfg: Any) → None

Initialize the Mixtral architecture adapter.

class transformer_lens.model_bridge.supported_architectures.ModernBertDecoderArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for ModernBertDecoderForCausalLM models.

__init__(cfg: Any) → None

Initialize the ModernBERT Decoder architecture adapter.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

supports_center_writing_weights = False
supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.MusicFlamingoArchitectureAdapter(cfg: Any)

Bases: AudioFlamingo3ArchitectureAdapter

Architecture adapter for MusicFlamingoForConditionalGeneration models.

class transformer_lens.model_bridge.supported_architectures.NanoChatArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for NanoChatForCausalLM models.

__init__(cfg: Any) → None

Initialize the NanoChat architecture adapter.

supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.NanogptArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for NanoGPT models.

__init__(cfg: Any) → None

Initialize the NanoGPT architecture adapter.

Parameters:

cfg – The configuration object.

convert_weights(remote_module: Any) → dict[str, Tensor]
class transformer_lens.model_bridge.supported_architectures.NativeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for NativeModel — TL-native, split-QKV, pre-LN; feature set driven by cfg (gated MLP, RMS norm, rotary, GQA, soft-cap, attn_only).

prepare_model(model: Any) → None

Reject modules whose attribute names collide with bridge slots.

Bridge’s __getattr__ falls back to getattr(original_model, name) for unknown attrs, so a name match — submodule, buffer, plain tensor, or property — makes add_module raise mid-setup with an opaque message. Failing here points at the real cause. Reserved set is derived from component_mapping.keys() so adapter variants stay in sync.

class transformer_lens.model_bridge.supported_architectures.NeelSoluOldArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Neel’s SOLU models (old style).

__init__(cfg: Any) → None

Initialize the Neel SOLU old-style architecture adapter.

Parameters:

cfg – The configuration object.

class transformer_lens.model_bridge.supported_architectures.NemotronArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for NemotronForCausalLM models.

__init__(cfg: Any) → None

Initialize the Nemotron architecture adapter.

class transformer_lens.model_bridge.supported_architectures.NemotronHArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for NemotronHForCausalLM.

Hybrid Mamba-2 + Attention + MoE + dense MLP model. All layers share a single pre-norm and a single residual connection; the mixer type per layer is determined by config.layers_block_type[layer_idx].

applicable_phases: list[int] = [1, 2, 3, 4]
create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: Any) → Any

Build the unified DynamicCache for stateful generation.

Transformers ≥ 5.12 ships a unified DynamicCache that carries both KV-cache entries (attention layers) and SSM conv/recurrent states (Mamba layers) in a single object, using has_previous_state() to distinguish which state is available for a given layer index. The config is required so the cache knows each layer’s type — matching NemotronHModel’s own initialization.

class transformer_lens.model_bridge.supported_architectures.NeoArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Neo models.

__init__(cfg: Any) → None

Initialize the Neo architecture adapter.

class transformer_lens.model_bridge.supported_architectures.NeoxArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for NeoX models.

__init__(cfg: Any) → None

Initialize the NeoX architecture adapter.

Parameters:

cfg – The configuration object.

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for GPT-NeoX/StableLM component testing.

GPT-NeoX models use RoPE (Rotary Position Embeddings) which need to be set on all attention bridge instances for component testing.

Parameters:

  • hf_model – The HuggingFace GPT-NeoX model instance

  • bridge_model – The TransformerBridge model (if available, set rotary_emb on actual instances)

split_qkv_matrix(original_attention_component: Any) → tuple[Linear, Linear, Linear]

Split the QKV matrix into separate linear transformations.

GPT-NeoX/StableLM uses an interleaved QKV format where the weights are stored as [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, …] - i.e., Q, K, V are interleaved per head.

The weight shape is [n_heads * 3 * d_head, d_model] and the output is reshaped by HuggingFace as [batch, seq, n_heads, 3*d_head] then split on the last dim.

Parameters:

original_attention_component – The original attention layer component

Returns:

Tuple of nn.Linear modules for Q, K, and V transformations

class transformer_lens.model_bridge.supported_architectures.Olmo2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OLMo 2 models.

OLMo 2 uses a post-norm architecture with RMSNorm, Q/K normalization in attention, rotary position embeddings (RoPE), and gated MLP (SwiGLU). Key differences from pre-norm models like Llama:

  • Post-norm: RMSNorm is applied AFTER attention and AFTER MLP, not before. ln1 maps to post_attention_layernorm, ln2 maps to post_feedforward_layernorm.

  • Q/K normalization: Per-head RMSNorm applied to queries and keys after projection.

  • No biases on any projections.

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP up_proj

  • blocks.{i}.mlp.b_gate - No bias on MLP gate_proj

  • blocks.{i}.mlp.b_out - No bias on MLP down_proj

  • blocks.{i}.ln1.b - RMSNorm has no bias

  • blocks.{i}.ln2.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

__init__(cfg: Any) → None

Initialize the OLMo 2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Olmo3ArchitectureAdapter(cfg: Any)

Bases: Olmo2ArchitectureAdapter

Architecture adapter for OLMo 3 / OLMo 3.1 models.

OLMo 3 is architecturally identical to OLMo 2 at the weight and component level. The only difference is sliding window attention on some layers (configurable via layer_types), which is handled by the HF model’s forward pass (mask creation) and does not affect weight structure or component mapping.

class transformer_lens.model_bridge.supported_architectures.OlmoArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OLMo (v1) models.

OLMo v1 uses a pre-norm architecture with a custom non-learnable LayerNorm (fixed weight=1, bias=0), rotary position embeddings (RoPE), and gated MLP (SwiGLU). Key differences from later OLMo variants:

  • Pre-norm: LayerNorm is applied BEFORE attention and BEFORE MLP.

  • Non-learnable LayerNorm: Weight and bias are not trainable parameters. Delegating to HF’s native forward via NormalizationBridge handles this correctly.

  • No Q/K normalization in attention.

  • Optional QKV clipping (handled by HF’s native attention forward).

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP up_proj

  • blocks.{i}.mlp.b_gate - No bias on MLP gate_proj

  • blocks.{i}.mlp.b_out - No bias on MLP down_proj

__init__(cfg: Any) → None

Initialize the OLMo architecture adapter.

prepare_model(hf_model: Any) → None

Patch OLMo’s in-place clamp_ to avoid backward hook conflicts.

OLMo v1 uses query_states.clamp_() when config.clip_qkv is set. In-place ops on tensors that pass through register_full_backward_hook trigger PyTorch’s “view modified inplace” error. This patch disables the in-place clamp branch during attention forward passes.

Note: clip_qkv clamping is skipped in the patched forward. In practice clip_qkv values (typically 100+) rarely activate. If exact clamping is needed, add out-of-place clamp hooks on hook_q/hook_k/hook_v.

class transformer_lens.model_bridge.supported_architectures.OlmoHybridArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OlmoHybridForCausalLM models.

__init__(cfg: Any) → None

Initialize the OLMo Hybrid architecture adapter.

create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: Any) → Any

OLMo Hybrid keeps per-layer q/k/v conv states in its own cache class.

supports_fold_ln = False
class transformer_lens.model_bridge.supported_architectures.OlmoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OLMoE (Mixture of Experts) models.

OLMoE uses a pre-norm architecture with RMSNorm, Q/K normalization in attention, rotary position embeddings (RoPE), and sparse Mixture of Experts MLP. Key features:

  • Pre-norm: RMSNorm applied BEFORE attention and BEFORE MLP.

  • Q/K normalization: RMSNorm applied to queries and keys after projection.

  • Sparse MoE: 64 experts with top-8 routing (configurable).

  • Batched expert parameters: gate_up_proj [num_experts, 2*d_mlp, d_model] and down_proj [num_experts, d_model, d_mlp] as single tensors, not a ModuleList.

  • Optional QKV clipping (handled by HF’s native attention forward).

  • No biases on any projections.

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.ln1.b - RMSNorm has no bias

  • blocks.{i}.ln2.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

__init__(cfg: Any) → None

Initialize the OLMoE architecture adapter.

prepare_model(hf_model: Any) → None

Patch OLMoE’s in-place clamp_ to avoid backward hook conflicts.

Same issue as OLMo v1 — see OlmoArchitectureAdapter.prepare_model.

class transformer_lens.model_bridge.supported_architectures.OpenAIGPTArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OpenAIGPTLMHeadModel (GPT-1) models.

__init__(cfg: Any) → None

Initialize the OpenAI GPT architecture adapter.

class transformer_lens.model_bridge.supported_architectures.OpenElmArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Apple OpenELM models.

OpenELM uses a unique architecture with per-layer varying head counts and FFN dimensions. Key characteristics:

  • Combined QKV projection (qkv_proj) with per-layer varying Q/KV head counts

  • Gated MLP with combined gate+up projection (proj_1) and per-layer FFN sizes

  • RMSNorm normalization

  • Full rotary embeddings (per-layer, not shared)

  • Optional Q/K RMSNorm (normalize_qk_projections=True)

  • Weight tying (share_input_output_layers=True typically)

  • Model root is ‘transformer’ (not ‘model’)

  • Requires trust_remote_code=True (custom HF code)

The native HF attention handles all per-layer dimension variations, RoPE, GQA group repeat, and Q/K normalization internally. The bridge delegates to the native forward for correct computation.

Note: Individual Q/K/V hooks are not available since the model uses a combined QKV projection. Attention-level hooks (hook_attn_in, hook_attn_out) are provided.

__init__(cfg: Any) → None

Initialize the OpenELM architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch OpenELM for compatibility with transformers v5.

Two patches are needed: 1. RotaryEmbedding: Custom _compute_sin_cos_embeddings fails on meta device

because it calls .cos() on meta tensors. We wrap it to catch NotImplementedError.

  1. Weight re-initialization: OpenELM’s _init_weights re-randomizes ALL weights after they’ve been loaded from safetensors because transformers v5’s _finalize_load_state_dict calls initialize_weights() on modules lacking the _is_hf_initialized flag. We patch _init_weights to skip real (non-meta) tensors.

Parameters:

  • model_name – The HuggingFace model name/path

  • model_kwargs – The kwargs dict for from_pretrained()

prepare_model(hf_model: Any) → None

Post-load fixes for non-persistent buffers zeroed during meta materialization.

Transformers v5 creates models on meta device then materializes weights from checkpoint. Non-persistent buffers (registered with persistent=False) are NOT in the checkpoint, so they materialize as zeros. OpenELM has two critical non-persistent buffers that must be recomputed:

  1. RoPE inv_freq — zeroed inv_freq produces cos=1, sin=0 for all positions, destroying positional information entirely.

  2. causal_mask — zeroed mask means no causal masking, allowing all positions to attend to future tokens. Single forward passes appear correct (no future tokens to leak) but autoregressive generation degenerates immediately.

We also create a synthetic lm_head for weight-tied models.

Note: We intentionally do NOT restore the original _compute_sin_cos_embeddings. The safe_compute wrapper is functionally equivalent for real (non-meta) tensors, and keeping it avoids issues when multiple models are loaded in the same process (e.g., benchmark suite loading both HF reference and bridge models).

Parameters:

hf_model – The loaded HuggingFace OpenELM model

class transformer_lens.model_bridge.supported_architectures.OptArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OPT models.

__init__(cfg: Any) → None

Initialize the OPT architecture adapter.

class transformer_lens.model_bridge.supported_architectures.OuroArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for ByteDance Ouro (LoopLM) models.

Ouro is a looped-depth (“Universal Transformer”) decoder: the remote-code OuroModel.forward applies the same num_hidden_layers-deep stack total_ut_steps times (4 for the released checkpoints) within a single forward pass, applying model.norm after every pass. The loop lives entirely inside the HF forward, which the bridge delegates to, so logits and generation are correct with no loop handling here. n_layers counts the physical layers; each block’s hooks fire once per loop step, and a cache records the final step’s value. The same holds for ln_final (model.norm): it runs after EVERY UT pass, so its hooks fire total_ut_steps times per forward and run_with_cache keeps only the last pass.

The backbone is Qwen2/Llama-shaped (RoPE, no-bias q/k/v/o projections, SwiGLU gate/up/down MLP, untied lm_head) with one twist: sandwich normalization. Each decoder layer has FOUR RMSNorms; the extra two (input_layernorm_2, post_attention_layernorm_2) apply to the sublayer outputs before the residual add, exactly like Gemma2’s ln1_post/ln2_post but without Gemma’s +1.0 RMSNorm offset.

Deliberately not mapped by this adapter:

  • per-loop-step hooks (a cache holds the final UT step only)

  • model.early_exit_gate, the adaptive-exit halting head

  • the UniversalTransformerCache slot layout (step * n_layers + layer)

Loading requires trust_remote_code=True (auto_map to modeling_ouro).

Optional Parameters (may not exist in state_dict):

Ouro models do NOT have biases on any mapped linear layers:

  • blocks.{i}.attn.b_Q / b_K / b_V / b_O - no attention biases

  • blocks.{i}.mlp.b_gate / b_in / b_out - no MLP biases

  • blocks.{i}.ln1.b / ln1_post.b / ln2.b / ln2_post.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

Weight processing must handle these missing biases gracefully using ProcessWeights._safe_get_tensor() or by checking for None values.

__init__(cfg: Any) → None

Initialize the Ouro architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch Ouro’s remote code for compatibility with transformers v5.

Ouro’s modeling code was written against transformers 4.55, where standard RoPE lived in ROPE_INIT_FUNCTIONS[“default”]. Transformers v5 removed that key and instead expects each *RotaryEmbedding class to carry a compute_default_rope_parameters static method. Two call sites break, so two patches:

  1. OuroRotaryEmbedding.__init__ does ROPE_INIT_FUNCTIONS[“default”] (KeyError). Rebind the module-level name inside the imported modeling_ouro module(s) to a copy with “default” restored; the shared transformers dict is left untouched.

  2. v5’s PreTrainedModel._init_weights re-initializes RotaryEmbedding buffers via module.compute_default_rope_parameters(config) (AttributeError). Attach the same function as a static method.

Parameters:

  • model_name – The HuggingFace model name/path

  • model_kwargs – The kwargs dict for from_pretrained()

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for Ouro component testing.

Ouro uses RoPE (Rotary Position Embeddings) with a single shared model.rotary_emb. We set the rotary_emb reference on all attention bridge instances for component testing.

Parameters:

  • hf_model – The HuggingFace Ouro model instance

  • bridge_model – The TransformerBridge model (if available, set rotary_emb on actual instances)

class transformer_lens.model_bridge.supported_architectures.PegasusArchitectureAdapter(cfg: Any)

Bases: BartFamilyArchitectureAdapter

Architecture adapter for PegasusForConditionalGeneration models.

has_final_stack_norm: bool = True
class transformer_lens.model_bridge.supported_architectures.Phi3ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Phi-3 models.

__init__(cfg: Any) → None

Initialize the Phi-3 architecture adapter.

Parameters:

cfg – The configuration object.

prepare_loading(model_name: str, model_kwargs: dict) → None

Patch cached Phi-3 remote code for transformers v5 compatibility.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Fold layer norms into joint QKV/gate_up projections.

Standard fold_ln can’t handle joint projections (shape mismatch on round-trip), so we scale the full joint weights directly.

class transformer_lens.model_bridge.supported_architectures.PhiArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Phi models.

__init__(cfg: Any) → None

Initialize the Phi architecture adapter.

Parameters:

cfg – The configuration object.

default_cfg: dict[str, Any] = {'use_fast': False}
class transformer_lens.model_bridge.supported_architectures.PhiMoEArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Microsoft PhiMoE models.

PhiMoE is a Phi-style decoder with LayerNorm, split Q/K/V attention, and a sparse MoE block. This adapter targets the native Transformers implementation (trust_remote_code=False); the archived remote implementation is not compatible with modern Transformers generation/cache semantics.

__init__(cfg: Any) → None

Initialize the PhiMoE architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

Disable remote code; base hook forces eager attention.

prepare_model(hf_model: Any) → None

Also force eager on the inner model module (PhiMoE re-derives it there).

class transformer_lens.model_bridge.supported_architectures.PretrainArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for a decoder-only transformer using RoPE, RMSNorm, gated SwiGLU MLPs, and optional sparse MoE feed-forward layers.

Uses an opaque NativeForwardAttentionBridge with no attention projection submodules, not JointQKVAttentionBridge/ PositionEmbeddingsAttentionBridge: those reimplement RoPE via HF’s rotate-half convention, wrong for a source model using the adjacent-pair convention. The opaque bridge delegates unchanged to Attention.forward, so RoPE runs as written – at the cost of no per-head hooks, only block-level resid_pre/resid_mid/resid_post.

Blocks use DelegatedAttentionBlockBridge rather than plain BlockBridge: that existing abstraction already exists for architectures where attention is delegated wholesale and the split-qkv-fork block-level aliases (hook_attn_in/hook_q_input/ hook_k_input/hook_v_input) don’t apply. It complements NativeForwardAttentionBridge.supports_split_qkv_fork = False (which prevents the split-QKV-fork machinery and its associated HookPoints from being exposed for this attention component) by also removing the now-dangling block-level aliases that would otherwise point at them. hook_attn_out is untouched by either change, since the attention component still fires its own hook_out normally.

self.cfg is mutated in place, not copied (matches nanogpt.py’s convention) – callers holding another reference to the same config will see these fields change.

Bridges built through build_pretrain_bridge are given a mode-propagating subclass so .train()/.eval() reach the wrapped source model (see that function’s docstring) – this adapter class itself has no lifecycle behavior of its own.

class transformer_lens.model_bridge.supported_architectures.Qwen2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Qwen2 models.

Qwen2 hardcodes q/k/v biases (o_proj, MLP, and norms are bias-free); the include_biases conversions keep GQA K/V biases in the per-head (n_kv_heads, d_head) layout weight processing expects.

__init__(cfg: Any) → None

Initialize the Qwen2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Qwen2AudioArchitectureAdapter(cfg: Any)

Bases: Qwen2ArchitectureAdapter

Architecture adapter for Qwen2AudioForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Qwen2-Audio architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Qwen2MoeArchitectureAdapter(cfg: Any)

Bases: Qwen2ArchitectureAdapter

Architecture adapter for Qwen2-MoE models.

Qwen2-MoE uses the Qwen2 attention stack plus a sparse MoE MLP with an always-on shared expert path.

__init__(cfg: Any) → None

Initialize the Qwen2-MoE architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Qwen2_5_VLArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Qwen2_5_VLForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Qwen2.5-VL architecture adapter.

required_libraries: list[str] = ['torchvision']
required_libraries_group: str = 'multimodal'
class transformer_lens.model_bridge.supported_architectures.Qwen3ArchitectureAdapter(cfg: Any, *, hybrid: bool = False, lm_prefix: str = 'model')

Bases: ArchitectureAdapter

Architecture adapter for Qwen3 dense models.

RMSNorm, RoPE, GQA, Q/K head norms, gated MLP. No biases. Serves as base class for Qwen3.5 and Qwen3Next hybrid variants.

class transformer_lens.model_bridge.supported_architectures.Qwen3MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Qwen3MoE (Mixture of Experts) models.

Qwen3MoE is a sparse MoE decoder-only Transformer, structurally close to OLMoE. Key features:

  • Pre-norm: RMSNorm applied BEFORE attention and BEFORE MLP.

  • Q/K normalization: RMSNorm applied to queries and keys after projection.

  • Sparse MoE: 128 experts with top-8 routing (public 30B-A3B checkpoints).

  • Batched expert parameters: gate_up_proj and down_proj as single 3D tensors, not a ModuleList.

  • final_rms=True (Qwen3-style; OLMoE uses False).

  • No biases on any projections.

  • GQA: n_key_value_heads < n_heads in all public checkpoints.

Only the all-MoE configuration is supported (decoder_sparse_step=1, mlp_only_layers=[]). Models with dense fallback layers cannot be wrapped because MoEBridge does not handle the dense Qwen3MoeMLP path.

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.ln1.b - RMSNorm has no bias

  • blocks.{i}.ln2.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

__init__(cfg: Any) → None

Initialize the Qwen3MoE architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Qwen3NextArchitectureAdapter(cfg: Any)

Bases: Qwen3ArchitectureAdapter

Hybrid linear-attention + full-attention with sparse MoE MLP.

Same hybrid design as Qwen3.5 but with MoE instead of dense MLP.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Slice query half from gated q_proj.weight for weight-space analysis.

class transformer_lens.model_bridge.supported_architectures.Qwen3VLArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Qwen3VLForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Qwen3-VL architecture adapter.

required_libraries: list[str] = ['torchvision']
required_libraries_group: str = 'multimodal'
class transformer_lens.model_bridge.supported_architectures.Qwen3VLMoeArchitectureAdapter(cfg: Any)

Bases: Qwen3VLArchitectureAdapter

Architecture adapter for Qwen3VLMoeForConditionalGeneration models.

class transformer_lens.model_bridge.supported_architectures.Qwen3_5ArchitectureAdapter(cfg: Any)

Bases: Qwen3ArchitectureAdapter

Hybrid linear-attention + full-attention with dense gated MLP.

Inherits Qwen3 config/attention/MLP structure. Differences: - Attention + linear_attn are optional (per-layer type) - Gated q_proj (2x wide) sliced by preprocess_weights for weight analysis

prepare_loading(model_name: str, model_kwargs: dict) → None

Swap multimodal Qwen3_5Config for text-only Qwen3_5TextConfig.

Published checkpoints carry architectures=[‘Qwen3_5ForConditionalGeneration’]. We replace config with text_config so AutoModelForCausalLM loads the text-only Qwen3_5ForCausalLM.

prepare_model(hf_model: Any) → None

Reject full multimodal Qwen3.5 models on this text-only adapter.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Slice query half from gated q_proj.weight for weight-space analysis.

In processed mode, W_Q is the pure query projection (for composition scores, logit lens). Gate signal available in unprocessed mode on full-attention layers via blocks.N.attn.hook_q_gate.

class transformer_lens.model_bridge.supported_architectures.Qwen3_5MoeArchitectureAdapter(cfg: Any)

Bases: Qwen3ArchitectureAdapter

Text-only Qwen3.5-MoE: hybrid GatedDeltaNet + full attention, sparse MoE MLP.

prepare_loading(model_name: str, model_kwargs: dict) → None

Swap to text_config so AutoModelForCausalLM loads the text-only model (checkpoints ship the ConditionalGeneration architecture).

prepare_model(hf_model: Any) → None

Reject full multimodal Qwen3.5-MoE models on this text-only adapter.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Slice query half from gated q_proj.weight for weight-space analysis.

class transformer_lens.model_bridge.supported_architectures.Qwen3_5MoeMultimodalArchitectureAdapter(cfg: Any)

Bases: Qwen3_5MultimodalArchitectureAdapter

Vision-language adapter for Qwen3_5MoeForConditionalGeneration.

Reuses the Qwen3.5 multimodal wiring (language model under model.language_model + vision tower) with the MLP swapped for sparse MoE.

class transformer_lens.model_bridge.supported_architectures.Qwen3_5MultimodalArchitectureAdapter(cfg: Any)

Bases: Qwen3ArchitectureAdapter

Full vision-language adapter for Qwen3_5ForConditionalGeneration.

preprocess_weights(state_dict: dict[str, Tensor]) → dict[str, Tensor]

Slice query half from gated q_proj.weight (matcher is path-prefix-agnostic).

required_libraries: list[str] = ['torchvision']
required_libraries_group: str = 'multimodal'
class transformer_lens.model_bridge.supported_architectures.QwenArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Qwen models.

__init__(cfg: Any) → None

Initialize the Qwen architecture adapter.

class transformer_lens.model_bridge.supported_architectures.RWKV7ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RWKV7ForCausalLM (RWKV-7 “Goose”).

Attention-free recurrent decoder: a flat stack of pre-norm blocks, each a generalized-delta-rule time-mixing sublayer plus a token-shifted squared-ReLU channel-mixing sublayer, wrapped by standard biased LayerNorm. The recurrence and the cross-block v_first threading live inside the fla remote-code forward, which the bridge delegates to; see the module docstring for the full set of adapter decisions.

__init__(cfg: Any) → None

Initialize the RWKV-7 architecture adapter.

applicable_phases: list[int] = []
prepare_loading(model_name: str, model_kwargs: dict) → None

Patch fla’s RWKV-7 remote code for transformers v5 compatibility.

Two defensive patches, mirroring raven:

  1. Tied-weights format. RWKV7ForCausalLM._tied_weights_keys is a list (["lm_head.weight"], the 4.x form). v5’s tie_weights -> get_expanded_tied_weights_keys calls .keys() on the mapping, which raises AttributeError on a list. Rewrite it to the v5 dict form {"lm_head.weight": "model.embeddings.weight"}. RWKV-7 defaults tie_word_embeddings=False (so the list path is usually short- circuited before .keys()), but the rewrite is harmless when untied and prevents the crash on any tied checkpoint.

  2. Weight re-init. Under v5’s meta-device load-then-materialise flow, PreTrainedModel._init_weights is invoked on modules that already hold checkpoint weights, re-randomising them. Guard it to skip modules whose parameters are already on a real (non-meta) device — the same defensive patch openelm.py / raven.py apply.

Parameters:

  • model_name – The HuggingFace model name/path.

  • model_kwargs – The kwargs dict for from_pretrained().

class transformer_lens.model_bridge.supported_architectures.RavenArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RavenForCausalLM (Huginn depth-recurrent decoder).

Prelude / weight-tied recurrent core / coda phases over a shared residual width. The recurrence and prelude re-injection live inside the remote-code HF forward, which the bridge delegates to; see the module docstring for the full set of adapter decisions.

__init__(cfg: Any) → None

Initialize the Raven / Huginn architecture adapter.

applicable_phases: list[int] = []
prepare_loading(model_name: str, model_kwargs: dict) → None

Patch Huginn’s remote code for transformers v5 compatibility.

Huginn’s modeling code targets transformers 4.44; two things break under v5 (5.8.1), so two patches:

  1. Tied-weights format. RavenForCausalLM._tied_weights_keys is a list (["lm_head.weight"], the 4.x format), but v5’s tie_weights -> get_expanded_tied_weights_keys calls .keys() on it and raises AttributeError. The model does not even construct. Rewrite it to the v5 dict form {"lm_head.weight": "transformer.wte.weight"} (Huginn ties lm_head to transformer.wte).

  2. Weight re-init. Under v5’s meta-device load-then-materialise flow, PreTrainedModel._init_weights is invoked on modules that already hold checkpoint weights, re-randomising them. Guard it to skip modules whose parameters are already on a real (non-meta) device — the same defensive patch openelm.py applies.

Parameters:

  • model_name – The HuggingFace model name/path.

  • model_kwargs – The kwargs dict for from_pretrained().

class transformer_lens.model_bridge.supported_architectures.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]
class transformer_lens.model_bridge.supported_architectures.RwkvArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RwkvForCausalLM models.

__init__(cfg: Any) → None

Initialize the RWKV architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]
prepare_loading(model_name: str, model_kwargs: dict) → None

Force use_cache off: per-layer in-place state writes break autograd under backward hooks, and only recurrent generation consumes them.

prepare_model(hf_model: Any) → None

Re-assert use_cache=False – prepare_loading only fires on the boot path, so directly-wrapped modules keep the default and leak state tuples into tensor-only hooks.

supports_batched_generation: bool = False
supports_fold_ln = False
supports_generation: bool = True
supports_kv_cache: bool = False
class transformer_lens.model_bridge.supported_architectures.SeedOssArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for SeedOssForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.SmolLM3ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for SmolLM3 models.

SmolLM3 is a pre-norm decoder with RMSNorm, grouped-query attention (GQA), a SwiGLU gated MLP, rotary position embeddings (RoPE), tied input and output embeddings, and no biases on any projection. The block shape matches Llama and Qwen2 exactly, so the component mapping and weight conversions mirror qwen2.py.

NoPE (No Positional Encoding): SmolLM3 disables RoPE on every no_rope_layer_interval-th layer (default every 4th) via config.no_rope_layers. That per-layer toggle lives inside HF’s SmolLM3Attention.forward, but the bridge reimplements attention and would otherwise rotate Q and K on those layers. The _SmolLM3AttentionBridge subclass handles it by suppressing position embeddings on NoPE layers, so the reimplemented attention matches HF.

No Q/K normalization: unlike Qwen3, SmolLM3 has no per-head Q or K RMSNorm, so the attention block uses the plain q/k/v/o submodules.

Optional Parameters (may not exist in state_dict):

SmolLM3 models do NOT have biases on any linear layers:

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP input (up_proj)

  • blocks.{i}.mlp.b_gate - No bias on MLP gate projection

  • blocks.{i}.mlp.b_out - No bias on MLP output (down_proj)

  • blocks.{i}.ln1.b - RMSNorm has no bias

  • blocks.{i}.ln2.b - RMSNorm has no bias

  • ln_final.b - RMSNorm has no bias

Weight processing must handle these missing biases gracefully using ProcessWeights._safe_get_tensor() or by checking for None values.

__init__(cfg: Any) → None

Initialize the SmolLM3 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.StableLmArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for StableLM models.

StableLM uses a Llama-like architecture with separate Q/K/V projections and gated MLP, but differs in using standard LayerNorm (not RMSNorm) and partial rotary embeddings (25% of head dimensions by default).

Supports optional features: - Grouped Query Attention (num_key_value_heads != num_attention_heads) - QKV bias (use_qkv_bias=True on some models like stable-code-3b) - Parallel residual connections (use_parallel_residual=True) - Per-head QK LayerNorm (qk_layernorm=True)

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - Only present when use_qkv_bias=True

  • blocks.{i}.attn.b_K - Only present when use_qkv_bias=True

  • blocks.{i}.attn.b_V - Only present when use_qkv_bias=True

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP up_proj

  • blocks.{i}.mlp.b_gate - No bias on MLP gate_proj

  • blocks.{i}.mlp.b_out - No bias on MLP down_proj

__init__(cfg: Any) → None

Initialize the StableLM architecture adapter.

class transformer_lens.model_bridge.supported_architectures.Starcoder2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Starcoder2ForCausalLM models.

__init__(cfg: Any) → None

Initialize the Starcoder2 architecture adapter.

class transformer_lens.model_bridge.supported_architectures.SwitchTransformersArchitectureAdapter(cfg: Any)

Bases: T5ArchitectureAdapter

Architecture adapter for SwitchTransformersForConditionalGeneration models.

__init__(cfg: Any) → None

Initialize the Switch Transformers architecture adapter.

prepare_loading(model_name: str, model_kwargs: dict) → None

The google/switch-base-* repos ship pytorch_model.bin only; skip v5’s Hub-side safetensors auto-conversion (it needs a conversion PR).

class transformer_lens.model_bridge.supported_architectures.T5ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for T5 models.

T5 is an encoder-decoder model with: - Shared embeddings - Encoder stack (self-attention + FFN) - Decoder stack (self-attention + cross-attention + FFN) - Language modeling head

Supports both standard T5 (DenseReluDense with wi/wo) and gated variants like Flan-T5 (T5DenseGatedActDense with wi_0/wi_1/wo).

__init__(cfg: Any) → None

Initialize the T5 architecture adapter.

Parameters:

cfg – The configuration object.

class transformer_lens.model_bridge.supported_architectures.T5Gemma2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for T5Gemma2ForConditionalGeneration (text-only).

Encoder: BlockBridge over model.encoder.text_model.layers (Gemma-style, QK-norm, no cross-attn) Decoder: T5Gemma2DecoderBlockBridge over model.decoder.layers (merged self+cross attention)

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for T5Gemma2 component testing.

Both the encoder text stack and the decoder carry their own rotary_emb. We set the reference on all PositionEmbeddingsAttentionBridge instances so that component-level forward calls can compute RoPE correctly, force eager attention (so patterns are hookable), and enable native layernorm autograd on QK-norm so the manual encoder path matches HF exactly.

class transformer_lens.model_bridge.supported_architectures.T5GemmaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for T5GemmaForConditionalGeneration.

Encoder: BlockBridge over model.encoder.layers (Gemma-style, no cross-attn) Decoder: T5GemmaDecoderBlockBridge over model.decoder.layers (adds cross-attn hooks)

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for T5Gemma component testing.

Both the encoder and decoder carry their own rotary_emb. We set the reference on all PositionEmbeddingsAttentionBridge instances so that component-level forward calls can compute RoPE correctly.

class transformer_lens.model_bridge.supported_architectures.VaultGemmaArchitectureAdapter(cfg: Any)

Bases: Gemma2ArchitectureAdapter

Architecture adapter for VaultGemmaForCausalLM models.

__init__(cfg: Any) → None

Initialize the VaultGemma architecture adapter.

applicable_phases: list[int] = [1, 2, 4]
supports_compatibility_mode: bool = False
class transformer_lens.model_bridge.supported_architectures.XGLMArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for XGLM models.

XGLM uses pre-norm LayerNorm, sinusoidal positional embeddings (no learnable weights), standard MHA with separate q/k/v/out_proj, and a 2-layer MLP (fc1/fc2) that lives directly on the decoder block rather than inside an mlp sub-module.

All attention projections and fc1/fc2 carry biases. lm_head has no bias. Embeddings are scaled by sqrt(d_model) at runtime in XGLMScaledWordEmbedding.

Optional Parameters (may not exist in state_dict):

None — all published XGLM checkpoints include all parameters listed above.

__init__(cfg: Any) → None

Initialize the XGLM architecture adapter.

class transformer_lens.model_bridge.supported_architectures.YoutuArchitectureAdapter(cfg: Any)

Bases: DeepSeekV2ArchitectureAdapter

Architecture adapter for YoutuForCausalLM models.

class transformer_lens.model_bridge.supported_architectures.Zamba2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for Zamba2ForCausalLM.

Hybrid Mamba-2 + shared global-attention model. Most layers are pure Mamba-2 SSM ("mamba"); a recurring subset are hybrid layers ("hybrid") that route through a shared attention block before the Mamba-2 step.

applicable_phases: list[int] = [1, 2, 3, 4]