transformer_lens.model_bridge.supported_architectures.pretrain module

Architecture adapter for a lightweight decoder-only pretraining model.

Maps a decoder-only transformer using RoPE, RMSNorm, gated SwiGLU MLPs, and optional sparse mixture-of-experts feed-forward layers into TransformerBridge, by wrapping the source module and delegating to its own forward rather than translating parameters into a second implementation.

Usage: build_pretrain_bridge(model, cfg) – the public entry point. PretrainModelContainer and direct build_bridge_from_module use are internal/advanced details (see PretrainModelContainer’s docstring).

Scope: maps a live module into TransformerBridge. Does not load checkpoints, merge tensor-parallel shards, or depend on a training framework.

Required module protocol – “lightweight decoder-only pretraining models” describes intent, not a generality guarantee. The wrapped model must expose:

model.embed (embedding lookup) model.blocks[i].norm1 (pre-attention norm) model.blocks[i].attn (called as attn(x, …)) model.blocks[i].norm2 (pre-MLP norm) model.blocks[i].mlp (gate/up/down, or router/experts) model.norm_f (final norm) model.lm_head (unembedding)

gate/up/down and router/experts name the supported protocol. DenseOrMoEFeedForwardBridge checks these structurally – attribute presence plus basic type (each is a module, experts is a registered module collection) – and raises clearly on a mismatch, but that is structural validation only: it does not and cannot validate that a module satisfying the shape actually implements matching forward semantics. Blocks must take more than the bare hidden state (this target passes cos/sin) – see PretrainModelContainer.

class transformer_lens.model_bridge.supported_architectures.pretrain.DenseOrMoEFeedForwardBridge(name: str, config: Any)

Bases: GeneralizedComponent

Wraps either a dense SwiGLU MLP or a sparse MoE layer behind a common interface. Dispatch is determined by structural inspection (router/experts vs gate/up/down) rather than configuration, so dense, MoE, and mixed architectures all use the same component mapping. This identifies the supported protocol – it does not validate that a module merely sharing those attribute names actually implements the matching forward behavior.

hasattr alone would let a module with, say, both router/experts and gate/up/down (or router/experts of the wrong types) win the MoE branch by attribute-name coincidence and fail later with a confusing error from deep inside MoEBridge, or not fail until forward time. set_original_component therefore also checks the basic shape of whichever protocol wins: router/gate/up/down must themselves be modules, and experts must be a registered module collection (nn.ModuleList/nn.ModuleDict) – a plain Python list/tuple of nn.Module experts is rejected even though every element is itself a valid module, because modules held in an ordinary list aren’t registered as children and would silently drop out of parameters()/state_dict()/.to(…)/train()/eval(), contradicting this adapter’s lifecycle guarantees.

forward(*args: Any, **kwargs: Any) → Any

Generic forward pass for bridge components with input/output hooks.

set_original_component(component: Module) → None

Set the original component that this bridge wraps.

Parameters:

original_component – The original transformer component to wrap

class transformer_lens.model_bridge.supported_architectures.pretrain.NativeForwardAttentionBridge(name: str | None, config: Any, submodules: Dict[str, GeneralizedComponent] | None = None, conversion_rule: BaseTensorConversion | None = None, pattern_conversion_rule: BaseTensorConversion | None = None, maintain_native_attention: bool = False, requires_position_embeddings: bool = False, requires_attention_mask: bool = False, attention_mask_4d: bool = False, requires_relative_position_bias: bool = False, is_cross_attention: bool = False, is_causal: bool = True, optional: bool = False)

Bases: AttentionBridge

Opaque attention bridge that delegates to the source attention.

This adapter intentionally exposes only input/output attention hooks. It has no mapped Q/K/V/O projection components, so the standard per-head aliases and weight aliases do not apply.

hook_aliases: Dict[str, str | List[str]] = {}

property_aliases: Dict[str, str] = {}

supports_split_qkv_fork: bool = False

class transformer_lens.model_bridge.supported_architectures.pretrain.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.pretrain.PretrainModelContainer(model: Module)

Bases: Module

Internal detail – build_pretrain_bridge applies this automatically. Three responsibilities, all in this container’s own forward:

1. Avoids a TransformerBridge.__getattr__ collision: a source model’s own self.embed/self.blocks clashes with identically named component_mapping keys. Wrapping one level deeper (container.inner.embed) fixes this without touching the source model.

2. Normalizes the return value to the .logits contract TransformerBridge expects: a plain “logits” dict (the target architecture’s actual shape) is wrapped in _LogitsAttrDict; a bare tensor, tensor-first tuple, or object already exposing .logits passes through after validating the tensor is present and is a tensor; anything else raises immediately with a clear message.

3. Strips _BRIDGE_COMPAT_KWARGS from kwargs before calling the wrapped model (see that constant’s comment).

Also sidesteps a BlockBridge convention where a bare-tensor block output gets wrapped in a 1-tuple for “standalone hidden_states calls”: since this target’s blocks take cos/sin too, that path never triggers, so the source forward loop needs no changes to be bridged.

self.inner is a regular registered submodule, so container.train()/.eval() already recurse into it via the normal nn.Module traversal – no override needed here. The propagation gap lives one level up, at TransformerBridge itself (see build_pretrain_bridge), whose .train()/.eval() do not walk down to original_model.

forward(*args: Any, **kwargs: Any) → Any

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

transformer_lens.model_bridge.supported_architectures.pretrain.build_pretrain_bridge(model: Module, cfg: TransformerBridgeConfig, *, device: Any = None, dtype: dtype | None = None, model_name: str | None = None) → TransformerBridge

Public entry point: wraps model in PretrainModelContainer and builds a TransformerBridge around it. Prefer this over calling build_bridge_from_module directly – the container is easy to forget.

device/dtype/model_name forward to build_bridge_from_module only when explicitly given.

bridge.train()/.eval() propagate to model via TransformerBridge.train() itself, which sets mode on original_model in addition to the registered module tree (original_model is deliberately not a registered submodule, so nn.Module.train()’s own recursion never reaches it). This adapter needs nothing extra for mode propagation.

Setting mode on model directly still works too and stays in sync.