From e6d19bdedf4c725c71a51d5341b19186dbc8d99f Mon Sep 17 00:00:00 2001
From: Katherine OMBRELLARO <kombrellaro@gmail.com>
Date: Fri, 17 Jul 2026 03:23:27 +0200
Subject: [PATCH 1/2] Add TransformerLens pretrain bridge adapter

---
 .../model_bridge/test_pretrain_adapter.py     | 312 ++++++++++
 tests/mocks/tiny_pretrain_model.py            | 193 +++++++
 .../test_moe_bridge_tuple_output.py           |  88 +++
 .../_pretrain_mocks.py                        | 214 +++++++
 .../test_pretrain_adapter.py                  | 453 +++++++++++++++
 .../test_pretrain_model_container.py          | 537 ++++++++++++++++++
 .../test_bridge_train_mode_propagation.py     | 100 ++++
 tests/unit/tools/test_model_registry.py       |   1 +
 .../factories/architecture_adapter_factory.py |  26 +-
 transformer_lens/model_bridge/bridge.py       |  74 ++-
 .../generalized_components/moe.py             |  17 +-
 .../supported_architectures/__init__.py       |  56 +-
 .../supported_architectures/pretrain.py       | 442 ++++++++++++++
 13 files changed, 2400 insertions(+), 113 deletions(-)
 create mode 100644 tests/integration/model_bridge/test_pretrain_adapter.py
 create mode 100644 tests/mocks/tiny_pretrain_model.py
 create mode 100644 tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py
 create mode 100644 tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py
 create mode 100644 tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py
 create mode 100644 tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py
 create mode 100644 tests/unit/model_bridge/test_bridge_train_mode_propagation.py
 create mode 100644 transformer_lens/model_bridge/supported_architectures/pretrain.py

diff --git a/tests/integration/model_bridge/test_pretrain_adapter.py b/tests/integration/model_bridge/test_pretrain_adapter.py
new file mode 100644
index 000000000..5878532e3
--- /dev/null
+++ b/tests/integration/model_bridge/test_pretrain_adapter.py
@@ -0,0 +1,312 @@
+"""Integration tests for PretrainArchitectureAdapter -- numerical parity
+against a real (not structural-mock) reference model.
+
+Uses `tests/mocks/tiny_pretrain_model.py`'s `TinyPretrainModel`, which
+implements genuine adjacent-pair RoPE and RMSNorm math (unlike
+`_pretrain_mocks.py`'s unit-test mocks, which stub RoPE with constant
+cos/sin and only need to get attribute names and call-signatures right).
+
+Covers: source/bridge logit parity (dense, full-MoE, mixed dense/MoE,
+untied embeddings); that this file's fixtures supply a `cfg` truthfully
+describing the wrapped model's head/vocab dimensions (a fixture property,
+not something `build_pretrain_bridge` itself enforces -- see
+`TestIntegrationFixturesSupplyATruthfulConfig`); residual-stream
+decomposition via `run_with_cache`; a forward-hook intervention
+cross-checked against an independent way of expressing the same edit; and
+float64 construction-time preservation.
+
+Note on scope: since the adapter delegates the whole forward to the
+source model rather than reimplementing attention/RoPE, exact logit
+parity mainly demonstrates that wrapping and output normalization don't
+alter an unhooked forward, not that the RoPE convention itself is
+correct (the same implementation runs on both sides of the comparison).
+`TestResidualStreamDecomposition` and `TestHookIntervention` are the
+stronger evidence for correct intervention points.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from tests.mocks.tiny_pretrain_model import TinyPretrainModel
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.supported_architectures.pretrain import (
+    ARCHITECTURE_NAME,
+    build_pretrain_bridge,
+)
+
+
+def _make_cfg(
+    *,
+    d_model: int,
+    n_layers: int,
+    n_heads: int = 2,
+    d_vocab: int = 64,
+) -> TransformerBridgeConfig:
+    """Build a bridge config matching this file's TinyPretrainModel
+    fixtures (n_heads=2, vocab_size=64 by default). See
+    `TestIntegrationFixturesSupplyATruthfulConfig` for why this matching
+    matters and what happens when it doesn't."""
+    return TransformerBridgeConfig(
+        d_model=d_model,
+        d_head=d_model // n_heads,
+        n_layers=n_layers,
+        n_ctx=128,
+        n_heads=n_heads,
+        d_vocab=d_vocab,
+        d_mlp=d_model * 2,
+        architecture=ARCHITECTURE_NAME,
+    )
+
+
+class TestSourceBridgeLogitParity:
+    """Bridge output must equal the source model's own `"logits"` output
+    exactly -- not within tolerance. The adapter never translates weights
+    into a second parameter layout, so there is no floating-point drift
+    to tolerate; any mismatch here is a real bug, not numerical noise."""
+
+    def test_dense_model_logit_parity_exact(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64)
+        model.eval()
+        tokens = torch.randint(0, 64, (1, 5))
+        with torch.no_grad():
+            source_logits = model(tokens)["logits"]
+
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2))
+        with torch.no_grad():
+            bridge_logits = bridge(tokens)
+
+        torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0)
+
+    def test_full_moe_model_logit_parity_exact(self) -> None:
+        """Every block's MLP is MoE (moe_layer_indices covers all layers)."""
+        torch.manual_seed(0)
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=2,
+            d_ff=32,
+            vocab_size=64,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({0, 1}),
+        )
+        model.eval()
+        tokens = torch.randint(0, 64, (1, 5))
+        with torch.no_grad():
+            source_logits = model(tokens)["logits"]
+
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2))
+        with torch.no_grad():
+            bridge_logits = bridge(tokens)
+
+        torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0)
+
+    def test_mixed_dense_moe_model_logit_parity_exact(self) -> None:
+        """One block dense, one block MoE -- the dispatcher must route each
+        block's MLP correctly within a single model, not just across
+        separately-constructed all-dense/all-MoE models."""
+        torch.manual_seed(0)
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=2,
+            d_ff=32,
+            vocab_size=64,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({1}),
+        )
+        model.eval()
+        tokens = torch.randint(0, 64, (1, 5))
+        with torch.no_grad():
+            source_logits = model(tokens)["logits"]
+
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2))
+        with torch.no_grad():
+            bridge_logits = bridge(tokens)
+
+        torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0)
+
+    def test_untied_embeddings_logit_parity_exact(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+            tie_embeddings=False,
+        )
+        model.eval()
+        tokens = torch.randint(0, 64, (1, 5))
+        with torch.no_grad():
+            source_logits = model(tokens)["logits"]
+
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=1))
+        with torch.no_grad():
+            bridge_logits = bridge(tokens)
+
+        torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0)
+
+
+class TestIntegrationFixturesSupplyATruthfulConfig:
+    """Integration fixtures must supply a config matching the source
+    model. `build_pretrain_bridge` does not itself validate this -- `cfg`
+    is caller-owned, per ordinary TransformerBridge convention -- so a
+    mismatch shows up only in `bridge.cfg`, silently, never in the
+    logits (see the second test)."""
+
+    def test_cfg_matches_model_head_and_vocab_dimensions(self) -> None:
+        torch.manual_seed(0)
+        n_heads = 2
+        vocab_size = 64
+        model = TinyPretrainModel(
+            d_model=16, n_heads=n_heads, n_layers=2, d_ff=32, vocab_size=vocab_size
+        )
+        model.eval()
+        cfg = _make_cfg(d_model=16, n_layers=2, n_heads=n_heads, d_vocab=vocab_size)
+        bridge = build_pretrain_bridge(model, cfg)
+
+        assert bridge.cfg.n_heads == n_heads
+        assert bridge.cfg.d_head == 16 // n_heads
+        assert bridge.cfg.d_vocab == model.embed.num_embeddings
+
+    def test_build_pretrain_bridge_does_not_validate_cfg_against_model(self) -> None:
+        """A cfg describing a different architecture than the wrapped
+        model still builds and still produces exact logit parity
+        (delegation never reads cfg) -- only bridge.cfg is wrong.
+
+        Current framework convention; not an adapter guarantee. This test
+        documents present behavior, not a requirement -- if
+        build_pretrain_bridge later gains cfg-vs-model validation, this
+        test should be updated (or removed) rather than treated as a
+        regression to preserve."""
+        torch.manual_seed(0)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        model.eval()
+        tokens = torch.randint(0, 64, (1, 5))
+        with torch.no_grad():
+            source_logits = model(tokens)["logits"]
+
+        wrong_cfg = _make_cfg(d_model=16, n_layers=1, n_heads=4, d_vocab=256)
+        bridge = build_pretrain_bridge(model, wrong_cfg)
+        with torch.no_grad():
+            bridge_logits = bridge(tokens)
+
+        torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0)
+        assert bridge.cfg.n_heads != model.blocks[0].attn.n_heads
+        assert bridge.cfg.d_vocab != model.embed.num_embeddings
+
+
+class TestResidualStreamDecomposition:
+    """`run_with_cache`'s resid hooks must decompose the way the residual
+    stream actually works, not just be present under the right names
+    (attribute/key presence is already covered by the structural unit
+    tests) -- resid_pre + attn_out == resid_mid, resid_mid + mlp_out ==
+    resid_post, and one block's resid_post is the next block's resid_pre."""
+
+    def test_resid_stream_decomposes_exactly_via_run_with_cache(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64)
+        model.eval()
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2))
+        tokens = torch.randint(0, 64, (1, 5))
+
+        with torch.no_grad():
+            _, cache = bridge.run_with_cache(tokens)
+
+        torch.testing.assert_close(
+            cache["blocks.0.hook_resid_pre"], cache["hook_embed"], atol=0, rtol=0
+        )
+        torch.testing.assert_close(
+            cache["blocks.0.hook_resid_mid"],
+            cache["blocks.0.hook_resid_pre"] + cache["blocks.0.hook_attn_out"],
+        )
+        torch.testing.assert_close(
+            cache["blocks.0.hook_resid_post"],
+            cache["blocks.0.hook_resid_mid"] + cache["blocks.0.hook_mlp_out"],
+        )
+        torch.testing.assert_close(
+            cache["blocks.1.hook_resid_pre"],
+            cache["blocks.0.hook_resid_post"],
+            atol=0,
+            rtol=0,
+        )
+        torch.testing.assert_close(
+            cache["blocks.1.hook_resid_post"],
+            cache["blocks.1.hook_resid_mid"] + cache["blocks.1.hook_mlp_out"],
+        )
+
+
+class TestHookIntervention:
+    """A forward-hook edit at `blocks.{i}.mlp.hook_out` must land at the
+    point the residual-stream decomposition says it should -- cross-
+    checked against an independently-expressed version of the same edit
+    (forcing resid_post to equal resid_mid), rather than only checking
+    that *some* change occurred."""
+
+    def test_mlp_hook_ablation_matches_independent_resid_post_override(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64)
+        model.eval()
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2))
+        tokens = torch.randint(0, 64, (1, 5))
+
+        with torch.no_grad():
+            baseline_logits, baseline_cache = bridge.run_with_cache(tokens)
+        resid_mid0 = baseline_cache["blocks.0.hook_resid_mid"].clone()
+
+        def zero_mlp_out(value: torch.Tensor, hook) -> torch.Tensor:
+            return torch.zeros_like(value)
+
+        def force_resid_post_to_resid_mid(value: torch.Tensor, hook) -> torch.Tensor:
+            return resid_mid0
+
+        with torch.no_grad():
+            via_mlp_ablation = bridge.run_with_hooks(
+                tokens, fwd_hooks=[("blocks.0.mlp.hook_out", zero_mlp_out)]
+            )
+            via_resid_override = bridge.run_with_hooks(
+                tokens,
+                fwd_hooks=[("blocks.0.hook_resid_post", force_resid_post_to_resid_mid)],
+            )
+
+        # Sanity: the intervention actually changed something.
+        assert not torch.equal(via_mlp_ablation, baseline_logits)
+        # The real check: two different-looking edits that are
+        # mathematically the same edit must produce identical output.
+        torch.testing.assert_close(via_mlp_ablation, via_resid_override, atol=0, rtol=0)
+
+
+class TestDtypePreservation:
+    """float64 is checked for preservation through construction only, not
+    full numerical parity at that dtype -- RoPE's cos/sin computation
+    involves trig functions whose float32-vs-float64 rounding can
+    legitimately differ from the bridge's own dtype-casting path in ways
+    unrelated to correctness, so exact-equality parity isn't a meaningful
+    check at non-default dtypes the way it is at float32."""
+
+    def test_float64_construction_preserves_dtype_and_tied_weights(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+            tie_embeddings=True,
+        ).to(torch.float64)
+        assert model.lm_head.weight is model.embed.weight
+
+        bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=1))
+
+        assert next(model.parameters()).dtype == torch.float64
+        assert model.lm_head.weight is model.embed.weight
+
+        tokens = torch.randint(0, 64, (1, 4))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.dtype == torch.float64
+        assert torch.isfinite(output).all()
diff --git a/tests/mocks/tiny_pretrain_model.py b/tests/mocks/tiny_pretrain_model.py
new file mode 100644
index 000000000..2c3d9434c
--- /dev/null
+++ b/tests/mocks/tiny_pretrain_model.py
@@ -0,0 +1,193 @@
+"""A tiny, self-contained decoder-only reference model (RoPE, RMSNorm,
+gated SwiGLU MLP, optional MoE) used only by the pretrain-adapter
+integration tests, to exercise real numerical parity rather than a purely
+structural mock.
+
+Deliberately uses the *adjacent-pair* RoPE convention
+(`[x0, x1] -> [-x1, x0]`), not HuggingFace's rotate-half (contiguous-half)
+convention -- this is the concrete case `PretrainArchitectureAdapter`'s
+docstring is about: an existing HF-oriented attention bridge would silently
+apply the wrong rotation to a model using this convention.
+
+Also deliberately computes RoPE `cos`/`sin` once per forward pass and
+passes them into every block call as extra positional args
+(`block(x, cos, sin)`), matching the target architecture's convention.
+This matters beyond efficiency: `BlockBridge` wraps a bare-tensor output
+in a 1-tuple for single-positional-argument "standalone hidden_states"
+calls (an HF-compatibility convention). A loop calling `block(x)` alone
+would need to unwrap `[0]` from each result -- modifying the source
+model's forward specifically because it's being bridged. Passing
+`cos`/`sin` positionally avoids that, so the forward loop needs no
+changes at all.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _rotate_adjacent_pairs(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., 0::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def _apply_rotary_adjacent_pairs(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    cos = cos.repeat_interleave(2, dim=-1)
+    sin = sin.repeat_interleave(2, dim=-1)
+    return x * cos + _rotate_adjacent_pairs(x) * sin
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d_model))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return x * self.weight
+
+
+class Attention(nn.Module):
+    """Receives precomputed `cos`/`sin` rather than computing them
+    internally -- matches the real target architecture, where rotary
+    tables are computed once per forward pass and shared across blocks."""
+
+    def __init__(self, d_model: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.o = nn.Linear(d_model, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        b, s, d = x.shape
+        qkv = self.qkv(x).view(b, s, 3, self.n_heads, self.d_head).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # each: [b, n_heads, s, d_head]
+
+        q = _apply_rotary_adjacent_pairs(q, cos, sin)
+        k = _apply_rotary_adjacent_pairs(k, cos, sin)
+
+        attn = torch.softmax(
+            (q @ k.transpose(-2, -1)) / (self.d_head**0.5)
+            + torch.triu(torch.full((s, s), float("-inf"), device=x.device), diagonal=1),
+            dim=-1,
+        )
+        out = (attn @ v).transpose(1, 2).reshape(b, s, d)
+        return self.o(out)
+
+
+class DenseMLP(nn.Module):
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.gate = nn.Linear(d_model, d_ff, bias=False)
+        self.up = nn.Linear(d_model, d_ff, bias=False)
+        self.down = nn.Linear(d_ff, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class MoEMLP(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, n_experts: int, top_k: int):
+        super().__init__()
+        self.top_k = top_k
+        self.router = nn.Linear(d_model, n_experts, bias=False)
+        self.experts = nn.ModuleList([DenseMLP(d_model, d_ff) for _ in range(n_experts)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        weights = torch.softmax(self.router(x), dim=-1)
+        topk_weights, topk_idx = weights.topk(self.top_k, dim=-1)
+        out = torch.zeros_like(x)
+        for e in range(len(self.experts)):
+            mask = topk_idx == e
+            if not mask.any():
+                continue
+            gate = (mask.float() * topk_weights).sum(dim=-1, keepdim=True)
+            out = out + gate * self.experts[e](x)
+        return out
+
+
+class Block(nn.Module):
+    """Takes `cos`/`sin` as extra positional args alongside the hidden
+    state -- see the module docstring for why this matters beyond RoPE
+    plumbing."""
+
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, mlp: nn.Module):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.attn = Attention(d_model, n_heads)
+        self.norm2 = RMSNorm(d_model)
+        self.mlp = mlp
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class TinyPretrainModel(nn.Module):
+    """Full tiny reference model for numerical-parity integration tests.
+    Matches `PretrainArchitectureAdapter`'s expected attribute paths:
+    `embed`, `blocks`, `norm_f`, `lm_head`.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 32,
+        n_heads: int = 4,
+        n_layers: int = 2,
+        d_ff: int = 64,
+        vocab_size: int = 256,
+        n_experts: int = 4,
+        top_k: int = 2,
+        moe_layer_indices: frozenset[int] = frozenset(),
+        tie_embeddings: bool = True,
+        rope_base: float = 10000.0,
+    ):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    d_model,
+                    n_heads,
+                    d_ff,
+                    mlp=(
+                        MoEMLP(d_model, d_ff, n_experts, top_k)
+                        if i in moe_layer_indices
+                        else DenseMLP(d_model, d_ff)
+                    ),
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.norm_f = RMSNorm(d_model)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        if tie_embeddings:
+            self.lm_head.weight = self.embed.weight
+
+        d_head = d_model // n_heads
+        inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_head, 2).float() / d_head))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, tokens: torch.Tensor) -> dict:
+        b, s = tokens.shape
+        x = self.embed(tokens)
+
+        # Computed once per forward pass, passed into every block -- the
+        # real target architecture's convention (see module docstring).
+        pos = torch.arange(s, device=tokens.device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(pos, self.inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        x = self.norm_f(x)
+        return {"logits": self.lm_head(x)}
diff --git a/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py b/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py
new file mode 100644
index 000000000..3b644497b
--- /dev/null
+++ b/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py
@@ -0,0 +1,88 @@
+"""Tests for MoEBridge tuple-output validation and router-score hooks.
+
+Unlike the pretrain-adapter's TestDenseOrMoEFeedForwardBridgeTupleOutput
+(which patches _delegate.forward and therefore exercises
+DenseOrMoEFeedForwardBridge.forward in isolation), these tests call
+MoEBridge.forward() directly via a stub original_component, so they cover
+MoEBridge's own tuple-validation contract rather than the dispatcher's.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from transformer_lens.model_bridge.generalized_components import MoEBridge
+
+
+class _StubMoEModule(nn.Module):
+    """Minimal nn.Module whose forward is swapped in per-test, standing in
+    for a real MoE layer wrapped by MoEBridge."""
+
+    def __init__(self, fake_forward):
+        super().__init__()
+        # Give the module at least one parameter so
+        # `next(self.parameters())` in MoEBridge.forward doesn't raise
+        # StopIteration and skip the dtype-cast branch entirely.
+        self.dummy = nn.Linear(4, 4)
+        self._fake_forward = fake_forward
+
+    def forward(self, *args, **kwargs):
+        return self._fake_forward(*args, **kwargs)
+
+
+def _bridge_with_stub(fake_forward) -> MoEBridge:
+    bridge = MoEBridge(name="mlp")
+    bridge.set_original_component(_StubMoEModule(fake_forward))
+    return bridge
+
+
+class TestMoEBridgeTupleOutput:
+    def test_empty_tuple_raises_clear_type_error(self) -> None:
+        bridge = _bridge_with_stub(lambda *a, **kw: ())
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            bridge(torch.ones(1, 3, 4))
+
+    def test_non_tensor_metadata_preserved_without_router_hook(self) -> None:
+        bridge = _bridge_with_stub(lambda *a, **kw: (torch.zeros(1, 3, 4), "not_router_scores"))
+        fired = {"called": False}
+
+        def mark_fired(value, hook):
+            fired["called"] = True
+
+        bridge.hook_router_scores.add_hook(mark_fired)
+
+        output = bridge(torch.ones(1, 3, 4))
+
+        assert isinstance(output, tuple)
+        assert output[1] == "not_router_scores"
+        torch.testing.assert_close(output[0], torch.zeros(1, 3, 4))
+        assert fired["called"] is False
+
+    def test_tensor_router_scores_still_fire_hook(self) -> None:
+        router_scores = torch.rand(1, 3, 8)
+        bridge = _bridge_with_stub(lambda *a, **kw: (torch.zeros(1, 3, 4), router_scores))
+        captured = {}
+
+        def capture_router_scores(value, hook):
+            captured["router_scores"] = value
+            return value
+
+        bridge.hook_router_scores.add_hook(capture_router_scores)
+
+        output = bridge(torch.ones(1, 3, 4))
+
+        assert isinstance(output, tuple)
+        torch.testing.assert_close(output[1], router_scores)
+        assert "router_scores" in captured
+        torch.testing.assert_close(captured["router_scores"], router_scores)
+
+    def test_non_tensor_first_element_raises_clear_type_error(self) -> None:
+        """Exercises the output[0]-validation added to MoEBridge itself
+        (not just the dispatcher) -- a malformed first element should
+        raise a clear TypeError here rather than failing inside
+        HookPoint's own type checking."""
+        bridge = _bridge_with_stub(lambda *a, **kw: ("not_a_tensor", torch.zeros(1, 3, 4)))
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            bridge(torch.ones(1, 3, 4))
diff --git a/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py b/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py
new file mode 100644
index 000000000..28c4f94f4
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py
@@ -0,0 +1,214 @@
+"""Shared mock modules and fixtures for PretrainArchitectureAdapter and
+PretrainModelContainer unit tests. Not a test file itself (no `test_`
+prefix, so pytest won't collect it) -- imported by both
+test_pretrain_adapter.py and test_pretrain_model_container.py.
+
+Fully self-contained: builds tiny mock `nn.Module`s that reproduce the
+relevant module tree (embed / blocks / norm1 / attn / norm2 / mlp / norm_f /
+lm_head) rather than depending on any external training-framework package.
+No weight loading, no network access.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig
+from transformer_lens.model_bridge.supported_architectures.pretrain import (
+    ARCHITECTURE_NAME,
+)
+
+# --------------------------------------------------------------------------
+# Tiny mock module tree, matching the real target architecture's
+# block-calling convention (cos/sin as extra positional args -- see
+# tiny_pretrain_model.py's module docstring). Real RoPE/RMSNorm math isn't
+# needed here (covered by integration parity tests); only correct
+# attribute names/shapes/call-signatures for component_mapping to wire up.
+# --------------------------------------------------------------------------
+
+
+class TinyRMSNorm(nn.Module):
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.weight
+
+
+class TinyAttention(nn.Module):
+    """Opaque attention stand-in: linear-in, linear-out, no real RoPE math.
+    Takes (and ignores) cos/sin positionally, matching the real target
+    architecture's block-calling convention."""
+
+    def __init__(self, d_model: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.out = nn.Linear(d_model, d_model)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        b, s, d = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        attn = torch.softmax(q @ k.transpose(-2, -1) / (d**0.5), dim=-1)
+        return self.out(attn @ v)
+
+
+class TinyDenseMLP(nn.Module):
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.gate = nn.Linear(d_model, d_ff, bias=False)
+        self.up = nn.Linear(d_model, d_ff, bias=False)
+        self.down = nn.Linear(d_ff, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(torch.nn.functional.silu(self.gate(x)) * self.up(x))
+
+
+class TinyMoE(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, n_experts: int, top_k: int):
+        super().__init__()
+        self.top_k = top_k
+        self.router = nn.Linear(d_model, n_experts, bias=False)
+        self.experts = nn.ModuleList([TinyDenseMLP(d_model, d_ff) for _ in range(n_experts)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        weights = torch.softmax(self.router(x), dim=-1)
+        topk_weights, topk_idx = weights.topk(self.top_k, dim=-1)
+        out = torch.zeros_like(x)
+        for expert_idx in range(len(self.experts)):
+            mask = (topk_idx == expert_idx).any(dim=-1, keepdim=True)
+            if mask.any():
+                out = out + mask * self.experts[expert_idx](x)
+        return out
+
+
+class TinyBlock(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, mlp: nn.Module):
+        super().__init__()
+        self.norm1 = TinyRMSNorm(d_model)
+        self.attn = TinyAttention(d_model, n_heads)
+        self.norm2 = TinyRMSNorm(d_model)
+        self.mlp = mlp
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class TinyPretrainModel(nn.Module):
+    """Minimal decoder-only model matching the adapter's expected paths:
+    embed / blocks / norm_f / lm_head. `moe_layer_indices` (possibly empty)
+    controls which blocks get a `TinyMoE` mlp instead of a dense one.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 32,
+        n_heads: int = 4,
+        n_layers: int = 2,
+        d_ff: int = 64,
+        vocab_size: int = 256,
+        n_experts: int = 4,
+        top_k: int = 2,
+        moe_layer_indices: frozenset[int] = frozenset(),
+        tie_embeddings: bool = True,
+    ):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.blocks = nn.ModuleList(
+            [
+                TinyBlock(
+                    d_model,
+                    n_heads,
+                    d_ff,
+                    mlp=(
+                        TinyMoE(d_model, d_ff, n_experts, top_k)
+                        if i in moe_layer_indices
+                        else TinyDenseMLP(d_model, d_ff)
+                    ),
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.norm_f = TinyRMSNorm(d_model)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        if tie_embeddings:
+            self.lm_head.weight = self.embed.weight
+        self.d_head = d_model // n_heads
+
+    def _forward_impl(self, tokens: torch.Tensor) -> dict:
+        b, s = tokens.shape
+        x = self.embed(tokens)
+        # Dummy cos/sin (no real RoPE math needed for these structural
+        # tests) but real tensors, passed positionally into every block --
+        # matches the real target architecture's convention.
+        cos = torch.ones(s, self.d_head // 2)
+        sin = torch.zeros(s, self.d_head // 2)
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        x = self.norm_f(x)
+        return {"logits": self.lm_head(x)}
+
+    def forward(self, tokens: torch.Tensor, **kwargs) -> dict:
+        # **kwargs declared so pytest fixtures can pass through
+        # run_with_cache's force-injected kwargs -- the actual kwarg-
+        # forwarding contracts (strict vs. **kwargs-accepting source
+        # forwards) are exercised directly against ForwardStrict/
+        # ForwardVarKwargs below instead.
+        return self._forward_impl(tokens)
+
+
+class ForwardStrict(nn.Module):
+    """Declares no **kwargs catch-all -- exercises the "must filter"
+    branch of PretrainModelContainer's kwarg handling."""
+
+    def __init__(self):
+        super().__init__()
+        self.seen_kwargs: dict = {}
+
+    def forward(self, tokens: torch.Tensor, targets: torch.Tensor | None = None) -> dict:
+        self.seen_kwargs = {"targets": targets}
+        return {"logits": tokens.float()}
+
+
+class ForwardVarKwargs(nn.Module):
+    """Declares **kwargs -- exercises the "pass everything through" branch
+    of PretrainModelContainer's kwarg handling."""
+
+    def __init__(self):
+        super().__init__()
+        self.seen_kwargs: dict = {}
+
+    def forward(self, tokens: torch.Tensor, **kwargs) -> dict:
+        self.seen_kwargs = kwargs
+        return {"logits": tokens.float()}
+
+
+class MalformedMLP(nn.Module):
+    """Has neither a dense MLP's (gate, up, down) nor an MoE layer's
+    (router, experts) attributes -- used to test the adapter's failure
+    path."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.some_other_layer = nn.Linear(d_model, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.some_other_layer(x)
+
+
+def make_cfg(d_model: int = 32, n_layers: int = 2) -> TransformerBridgeConfig:
+    return TransformerBridgeConfig(
+        d_model=d_model,
+        d_head=d_model // 4,
+        n_layers=n_layers,
+        n_ctx=128,
+        n_heads=4,
+        d_vocab=256,
+        d_mlp=d_model * 2,
+        architecture=ARCHITECTURE_NAME,
+    )
diff --git a/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py b/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py
new file mode 100644
index 000000000..d2fa03f24
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py
@@ -0,0 +1,453 @@
+"""Unit tests for PretrainArchitectureAdapter and DenseOrMoEFeedForwardBridge
+-- component mapping, structural dispatch, and bridge-construction-level
+behavior. Container, builder, and wrapped-model lifecycle behavior (kwarg
+filtering, output-contract normalization, hook registration for the hidden
+MLP delegate, train/eval, dtype/device) lives in
+test_pretrain_model_container.py.
+
+Fully self-contained (see _pretrain_mocks.py): tiny mock `nn.Module`s, no
+external package dependency, no network access.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from transformer_lens.model_bridge.generalized_components import (
+    DelegatedAttentionBlockBridge,
+)
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_from_module,
+)
+from transformer_lens.model_bridge.supported_architectures.pretrain import (
+    ARCHITECTURE_NAME,
+    DenseOrMoEFeedForwardBridge,
+    NativeForwardAttentionBridge,
+    PretrainArchitectureAdapter,
+    PretrainModelContainer,
+    build_pretrain_bridge,
+)
+
+from ._pretrain_mocks import (
+    MalformedMLP,
+    TinyDenseMLP,
+    TinyMoE,
+    TinyPretrainModel,
+    make_cfg,
+)
+
+_make_cfg = make_cfg  # local alias, matches call sites below
+
+
+class TestPretrainAdapterConstruction:
+    def test_adapter_sets_required_config_flags(self) -> None:
+        adapter = PretrainArchitectureAdapter(_make_cfg())
+        assert adapter.cfg.normalization_type == "RMS"
+        assert adapter.cfg.positional_embedding_type == "rotary"
+        assert adapter.cfg.final_rms is True
+        assert adapter.cfg.gated_mlp is True
+        assert adapter.cfg.attn_only is False
+
+    def test_adapter_mutates_the_passed_in_cfg_object_in_place(self) -> None:
+        """Documented, intentional behavior (matches nanogpt.py's adapter
+        convention) -- not an accidental side effect. This test guards the
+        specific claim: the object passed in is the object mutated, not a
+        copy."""
+        cfg = _make_cfg()
+        adapter = PretrainArchitectureAdapter(cfg)
+        assert adapter.cfg is cfg
+        assert cfg.normalization_type == "RMS"
+
+    def test_component_mapping_has_expected_top_level_keys(self) -> None:
+        adapter = PretrainArchitectureAdapter(_make_cfg())
+        mapping = adapter.get_component_mapping()
+        assert set(mapping.keys()) == {"embed", "blocks", "ln_final", "unembed"}
+
+    def test_mlp_mapping_always_uses_the_dispatcher(self) -> None:
+        """Unconditional now -- no cfg.num_experts branch to depend on."""
+        adapter = PretrainArchitectureAdapter(_make_cfg())
+        mlp_component = adapter.component_mapping["blocks"].submodules["mlp"]
+        assert isinstance(mlp_component, DenseOrMoEFeedForwardBridge)
+
+
+class TestNativeForwardAttentionBridge:
+    def test_opaque_attention_bridge_does_not_advertise_per_head_aliases(self) -> None:
+        """The opaque attention wrap has no Q/K/V/O submodules, so it must
+        not advertise AttentionBridge's class-level hook_aliases/
+        property_aliases (which assume those submodules exist) or the
+        split-QKV-fork machinery -- see NativeForwardAttentionBridge's
+        docstring."""
+        adapter = PretrainArchitectureAdapter(_make_cfg())
+        attn = adapter.component_mapping["blocks"].submodules["attn"]
+
+        assert isinstance(attn, NativeForwardAttentionBridge)
+        assert attn.hook_aliases == {}
+        assert attn.property_aliases == {}
+        assert attn.supports_split_qkv_fork is False
+
+        # Also check the actual constructed runtime bridge, not just the
+        # adapter's own component_mapping -- proves bridge construction
+        # didn't reintroduce the aliases along the way.
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256)
+        bridge = build_pretrain_bridge(model, cfg)
+        assert bridge.blocks[0].attn.hook_aliases == {}
+
+    def test_block_level_attention_output_still_reaches_hook_out(self) -> None:
+        """Clearing the per-head aliases must not also silence the block's
+        own attn.hook_out -- that's the actual hook point this adapter's
+        block-level hook coverage relies on."""
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256)
+        bridge = build_pretrain_bridge(model, cfg)
+        tokens = torch.randint(0, 256, (1, 5))
+        _, cache = bridge.run_with_cache(tokens)
+        assert "blocks.0.attn.hook_out" in cache
+
+    def test_blocks_use_delegated_attention_block_bridge(self) -> None:
+        """The adapter reuses DelegatedAttentionBlockBridge rather than
+        plain BlockBridge, since NativeForwardAttentionBridge's
+        supports_split_qkv_fork = False means the block-level
+        hook_attn_in/hook_q_input/hook_k_input/hook_v_input aliases would
+        otherwise dangle (point at HookPoints that are never created)."""
+        adapter = PretrainArchitectureAdapter(_make_cfg())
+        assert isinstance(adapter.component_mapping["blocks"], DelegatedAttentionBlockBridge)
+
+    def test_split_qkv_hook_names_are_absent_but_attention_output_remains_available(
+        self,
+    ) -> None:
+        """hook_attn_in/hook_q_input/hook_k_input/hook_v_input must not
+        appear as resolvable hook names anywhere in the built bridge --
+        DelegatedAttentionBlockBridge strips the block-level aliases, and
+        NativeForwardAttentionBridge never creates the underlying
+        attention-level HookPoints in the first place. hook_attn_out is
+        untouched by either change and must still resolve, both as the
+        block-level alias and as the attention component's own hook_out."""
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256)
+        bridge = build_pretrain_bridge(model, cfg)
+        hook_dict = bridge.hook_dict
+        for dangling in (
+            "blocks.0.hook_attn_in",
+            "blocks.0.hook_q_input",
+            "blocks.0.hook_k_input",
+            "blocks.0.hook_v_input",
+            "blocks.0.attn.hook_attn_in",
+            "blocks.0.attn.hook_q_input",
+            "blocks.0.attn.hook_k_input",
+            "blocks.0.attn.hook_v_input",
+        ):
+            assert dangling not in hook_dict, f"{dangling} should not resolve"
+        assert "blocks.0.hook_attn_out" in hook_dict
+        assert "blocks.0.attn.hook_out" in hook_dict
+
+
+class TestDenseOrMoEFeedForwardBridgeDispatch:
+    def test_dispatches_to_moe_delegate_for_moe_module(self) -> None:
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        moe_module = TinyMoE(d_model=32, d_ff=64, n_experts=4, top_k=2)
+        bridge.set_original_component(moe_module)
+
+        from transformer_lens.model_bridge.generalized_components import MoEBridge
+
+        assert isinstance(bridge._delegate, MoEBridge)
+
+    def test_dispatches_to_dense_delegate_for_dense_module(self) -> None:
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        dense_module = TinyDenseMLP(d_model=32, d_ff=64)
+        bridge.set_original_component(dense_module)
+
+        from transformer_lens.model_bridge.generalized_components import GatedMLPBridge
+
+        assert isinstance(bridge._delegate, GatedMLPBridge)
+
+    def test_raises_clear_error_for_malformed_module(self) -> None:
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        malformed = MalformedMLP(d_model=32)
+
+        with pytest.raises(ValueError, match="doesn't know how to wrap it"):
+            bridge.set_original_component(malformed)
+
+
+class TestDenseOrMoEFeedForwardBridgeDispatchValidatesTypes:
+    """hasattr alone would let a module win a branch by attribute-name
+    coincidence even if the attributes are the wrong type. These tests
+    exercise the basic type checks added to set_original_component so
+    such a mismatch is caught here, with a clear message naming the
+    actual field and type, rather than failing later inside MoEBridge/
+    GatedMLPBridge or during forward."""
+
+    def test_router_of_wrong_type_raises_type_error(self) -> None:
+        class BadRouterMoE(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.router = "not_a_module"  # right name, wrong type
+                self.experts = torch.nn.ModuleList(
+                    [TinyDenseMLP(d_model=32, d_ff=64) for _ in range(4)]
+                )
+
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        with pytest.raises(TypeError, match="router"):
+            bridge.set_original_component(BadRouterMoE())
+
+    def test_experts_of_wrong_type_raises_type_error(self) -> None:
+        class BadExpertsMoE(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.router = torch.nn.Linear(32, 4)
+                self.experts = "not_a_module_collection"  # right name, wrong type
+
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        with pytest.raises(TypeError, match="experts"):
+            bridge.set_original_component(BadExpertsMoE())
+
+    def test_plain_list_of_experts_is_rejected(self) -> None:
+        """Modules held in an ordinary Python list are not registered as
+        children -- their parameters would silently drop out of
+        parameters()/state_dict()/.to(...)/train()/eval(), contradicting
+        this adapter's lifecycle guarantees. Rejected even though every
+        element is itself a valid nn.Module."""
+
+        class ListExpertsMoE(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.router = torch.nn.Linear(32, 4)
+                self.experts = [TinyDenseMLP(d_model=32, d_ff=64) for _ in range(4)]
+
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        with pytest.raises(TypeError, match="registered module collection"):
+            bridge.set_original_component(ListExpertsMoE())
+
+    def test_gate_of_wrong_type_raises_type_error(self) -> None:
+        class BadGateDense(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.gate = "not_a_module"  # right name, wrong type
+                self.up = torch.nn.Linear(32, 64, bias=False)
+                self.down = torch.nn.Linear(64, 32, bias=False)
+
+        cfg = _make_cfg()
+        bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg)
+        with pytest.raises(TypeError, match="gate"):
+            bridge.set_original_component(BadGateDense())
+
+
+class TestDenseOrMoEFeedForwardBridgeTupleOutput:
+    """Exercises the dispatcher's own tuple-output handling directly, by
+    patching an already-wired delegate's forward -- independent of
+    whether GatedMLPBridge/MoEBridge happen to return a tuple today. The
+    dispatcher explicitly supports and validates this shape, so it's
+    tested as its own contract rather than assumed from what the real
+    delegates currently do."""
+
+    def _dispatcher_with_patched_delegate(self, fake_forward):
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({0}),
+        )
+        bridge = build_pretrain_bridge(model, cfg)
+        dispatcher = bridge.blocks[0].mlp
+        dispatcher._delegate.forward = fake_forward
+        return dispatcher
+
+    def test_tensor_first_tuple_passes_through_with_hook_applied(self) -> None:
+        dispatcher = self._dispatcher_with_patched_delegate(
+            lambda *a, **kw: (torch.zeros(1, 3, 16), "aux")
+        )
+        output = dispatcher(torch.ones(1, 3, 16))
+        assert isinstance(output, tuple)
+        assert output[1] == "aux"
+        torch.testing.assert_close(output[0], torch.zeros(1, 3, 16))
+
+    def test_empty_tuple_raises_clear_type_error(self) -> None:
+        dispatcher = self._dispatcher_with_patched_delegate(lambda *a, **kw: ())
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            dispatcher(torch.ones(1, 3, 16))
+
+    def test_non_tensor_first_element_raises_clear_type_error(self) -> None:
+        dispatcher = self._dispatcher_with_patched_delegate(
+            lambda *a, **kw: ("not_a_tensor", torch.zeros(1))
+        )
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            dispatcher(torch.ones(1, 3, 16))
+
+
+class TestPretrainAdapterBridgeConstruction:
+    @pytest.fixture
+    def dense_bridge(self):
+        cfg = _make_cfg(n_layers=2)
+        model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=2, d_ff=64, vocab_size=256)
+        return build_pretrain_bridge(model, cfg)
+
+    @pytest.fixture
+    def moe_bridge(self):
+        cfg = _make_cfg(n_layers=2)
+        model = TinyPretrainModel(
+            d_model=32,
+            n_heads=4,
+            n_layers=2,
+            d_ff=64,
+            vocab_size=256,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({1}),
+        )
+        return build_pretrain_bridge(model, cfg)
+
+    def test_bridge_has_correct_block_count(self, dense_bridge) -> None:
+        assert len(dense_bridge.blocks) == 2
+
+    def test_bridge_has_embed_unembed_and_final_norm(self, dense_bridge) -> None:
+        assert hasattr(dense_bridge, "embed")
+        assert hasattr(dense_bridge, "unembed")
+        assert hasattr(dense_bridge, "ln_final")
+
+    def test_forward_returns_logits_of_expected_shape(self, dense_bridge) -> None:
+        tokens = torch.randint(0, 256, (1, 5))
+        with torch.no_grad():
+            output = dense_bridge(tokens)
+        assert output.shape == (1, 5, 256)
+
+    def test_run_with_cache_exposes_resid_hooks(self, dense_bridge) -> None:
+        tokens = torch.randint(0, 256, (1, 5))
+        _, cache = dense_bridge.run_with_cache(tokens)
+        assert any("resid_pre" in k for k in cache.keys())
+        assert any("resid_post" in k for k in cache.keys())
+
+    def test_no_per_head_attention_hooks(self, dense_bridge) -> None:
+        """This adapter deliberately does not expose hook_q/hook_k/hook_v --
+        see PretrainArchitectureAdapter's class docstring."""
+        tokens = torch.randint(0, 256, (1, 5))
+        _, cache = dense_bridge.run_with_cache(tokens)
+        assert not any("hook_q" in k for k in cache.keys())
+        assert not any("hook_pattern" in k for k in cache.keys())
+
+    def test_moe_bridge_forward_runs(self, moe_bridge) -> None:
+        tokens = torch.randint(0, 256, (1, 5))
+        with torch.no_grad():
+            output = moe_bridge(tokens)
+        assert output.shape == (1, 5, 256)
+        assert not torch.isnan(output).any()
+
+    def test_moe_block_uses_dispatcher(self, moe_bridge) -> None:
+        assert isinstance(moe_bridge.blocks[1].mlp, DenseOrMoEFeedForwardBridge)
+
+    def test_moe_layer_hooks_are_registered_at_the_dispatcher_path(self, moe_bridge) -> None:
+        """Verifies actual cache keys, not just isinstance -- the claim in
+        DenseOrMoEFeedForwardBridge's docstring that hooks fire at
+        blocks.{i}.mlp.hook_in/hook_out regardless of dense-vs-MoE."""
+        tokens = torch.randint(0, 256, (1, 5))
+        _, cache = moe_bridge.run_with_cache(tokens)
+        assert "blocks.1.mlp.hook_in" in cache
+        assert "blocks.1.mlp.hook_out" in cache
+        # Layer 0 is dense in this fixture -- same path convention applies.
+        assert "blocks.0.mlp.hook_in" in cache
+        assert "blocks.0.mlp.hook_out" in cache
+
+
+class TestMoEConfigFieldIndependence:
+    def test_bridge_behaves_identically_with_moe_config_fields_populated(self) -> None:
+        """MoEBridge itself (transformer_lens/model_bridge/generalized_
+        components/moe.py) never reads cfg.num_experts or
+        cfg.experts_per_token. Populating them with realistic values must
+        not change behavior.
+        Uses an actual MoE layer (moe_layer_indices={0}) -- both models
+        being dense would let this pass without ever constructing
+        MoEBridge, so the dispatch type is asserted directly too."""
+        from transformer_lens.model_bridge.generalized_components import MoEBridge
+
+        torch.manual_seed(0)
+        model_a = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({0}),
+        )
+        cfg_without = _make_cfg(n_layers=1)
+        bridge_without = build_pretrain_bridge(model_a, cfg_without)
+
+        torch.manual_seed(0)
+        model_b = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+            n_experts=4,
+            top_k=2,
+            moe_layer_indices=frozenset({0}),
+        )
+        cfg_with = _make_cfg(n_layers=1)
+        cfg_with.num_experts = 4
+        cfg_with.experts_per_token = 2
+        bridge_with = build_pretrain_bridge(model_b, cfg_with)
+
+        assert isinstance(bridge_without.blocks[0].mlp._delegate, MoEBridge)
+        assert isinstance(bridge_with.blocks[0].mlp._delegate, MoEBridge)
+
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            out_without = bridge_without(tokens)
+            out_with = bridge_with(tokens)
+        torch.testing.assert_close(out_without, out_with, atol=0, rtol=0)
+
+
+class TestPretrainAdapterTiedEmbeddings:
+    def test_tied_embedding_bridge_construction_succeeds(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(
+            d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=True
+        )
+        bridge = build_pretrain_bridge(model, cfg)
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.shape == (1, 3, 64)
+
+    def test_untied_embedding_bridge_construction_succeeds(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(
+            d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=False
+        )
+        bridge = build_pretrain_bridge(model, cfg)
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.shape == (1, 3, 64)
+
+
+class TestBuildBridgeFromModuleDirectlyStillWorks:
+    """Advanced/internal path -- build_pretrain_bridge is the intended
+    public entry point, but direct build_bridge_from_module use (with the
+    container applied manually) must keep working too."""
+
+    def test_direct_build_bridge_from_module_with_manual_container(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_bridge_from_module(
+            PretrainModelContainer(model),
+            architecture=ARCHITECTURE_NAME,
+            tl_config=cfg,
+        )
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.shape == (1, 3, 64)
diff --git a/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py b/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py
new file mode 100644
index 000000000..43c40901e
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py
@@ -0,0 +1,537 @@
+"""Container, builder, and wrapped-model lifecycle behavior for
+PretrainModelContainer and build_pretrain_bridge -- kwarg filtering,
+output-contract normalization, train/eval mode propagation, and
+dtype/tied-weight preservation. Hook registration tests for the hidden
+MLP delegate
+(implemented by DenseOrMoEFeedForwardBridge) live here too, since what
+they verify -- real parameters staying reachable through
+`original_model.inner` -- is a lifecycle concern, not a mapping one.
+Adapter-mapping-level behavior lives in test_pretrain_adapter.py.
+
+Fully self-contained (see _pretrain_mocks.py): tiny mock `nn.Module`s, no
+external package dependency, no network access.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from transformer_lens.model_bridge.supported_architectures.pretrain import (
+    PretrainModelContainer,
+    _LogitsAttrDict,
+    build_pretrain_bridge,
+)
+
+from ._pretrain_mocks import (
+    ForwardStrict,
+    ForwardVarKwargs,
+    TinyPretrainModel,
+    make_cfg,
+)
+
+_make_cfg = make_cfg  # local alias, matches call sites below
+
+
+class TestPretrainModelContainerKwargFiltering:
+    def test_output_attentions_is_stripped_for_strict_signature(self) -> None:
+        """The one TransformerBridge-injected kwarg (run_with_cache
+        forces output_attentions=True) must not reach a source forward with
+        no **kwargs catch-all."""
+        model = ForwardStrict()
+        container = PretrainModelContainer(model)
+        container(torch.tensor([[1, 2, 3]]), targets=None, output_attentions=True)
+        assert model.seen_kwargs == {"targets": None}
+
+    def test_output_attentions_is_stripped_even_for_var_kwargs_signature(self) -> None:
+        """Stripped unconditionally, regardless of whether the source would
+        have accepted it anyway -- keeps behavior uniform across source
+        forward signatures."""
+        model = ForwardVarKwargs()
+        container = PretrainModelContainer(model)
+        container(torch.tensor([[1, 2, 3]]), output_attentions=True, some_other_kwarg=42)
+        assert model.seen_kwargs == {"some_other_kwarg": 42}
+
+    def test_unrelated_kwarg_typo_is_not_silently_swallowed(self) -> None:
+        """A genuine caller mistake (e.g. `target=` instead of `targets=`)
+        must raise, not be silently discarded alongside the one kwarg this
+        container is actually responsible for stripping."""
+        model = ForwardStrict()
+        container = PretrainModelContainer(model)
+        with pytest.raises(TypeError):
+            container(torch.tensor([[1, 2, 3]]), target=torch.tensor([1]))
+
+
+class TestPretrainModelContainerHookRegistration:
+    """The delegate is hidden from nn.Module registration (see
+    DenseOrMoEFeedForwardBridge), so blocks.{i}.mlp.hook_in/out are the
+    only publicly registered MLP hook points. Forward, gradients, dtype
+    conversion, and state_dict() still reach the hidden delegate's real
+    weights through the raw wrapped model."""
+
+    def _build(self):
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        return build_pretrain_bridge(model, cfg)
+
+    def test_only_the_dispatcher_hook_out_is_registered(self) -> None:
+        bridge = self._build()
+        tokens = torch.randint(0, 64, (1, 4))
+        with torch.no_grad():
+            _, cache = bridge.run_with_cache(tokens)
+        assert "blocks.0.mlp.hook_out" in cache
+        assert "blocks.0.mlp._delegate.hook_out" not in cache
+
+    def test_broad_hook_selector_applies_intervention_exactly_once(self) -> None:
+        bridge = self._build()
+        tokens = torch.randint(0, 64, (1, 4))
+        calls: list[str] = []
+
+        def recording_add_one(value: torch.Tensor, hook) -> torch.Tensor:
+            calls.append(hook.name)
+            return value + 1.0
+
+        def broad_filter(name: str) -> bool:
+            return "mlp" in name and "hook_out" in name
+
+        with torch.no_grad():
+            bridge.run_with_hooks(tokens, fwd_hooks=[(broad_filter, recording_add_one)])
+
+        assert calls == ["blocks.0.mlp.hook_out"]
+
+    def test_gradients_still_reach_the_hidden_delegates_weights(self) -> None:
+        bridge = self._build()
+        tokens = torch.randint(0, 64, (1, 4))
+        bridge(tokens).sum().backward()
+        mlp = bridge.original_model.inner.blocks[0].mlp
+        assert mlp.gate.weight.grad is not None
+        assert mlp.gate.weight.grad.abs().sum().item() > 0
+
+    def test_bridge_to_dtype_reaches_the_hidden_delegates_weights(self) -> None:
+        bridge = self._build()
+        bridge.to(dtype=torch.float64)
+
+        assert bridge.original_model.inner.embed.weight.dtype == torch.float64
+        assert bridge.original_model.inner.blocks[0].mlp.gate.weight.dtype == torch.float64
+
+        tokens = torch.randint(0, 64, (1, 4))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.dtype == torch.float64
+
+    def test_state_dict_includes_the_hidden_delegates_weights(self) -> None:
+        bridge = self._build()
+        source_weight = bridge.original_model.inner.blocks[0].mlp.gate.weight
+
+        state = bridge.state_dict()
+        matching = [v for k, v in state.items() if k.endswith("blocks.0.mlp.gate.weight")]
+
+        assert len(matching) == 1
+        torch.testing.assert_close(matching[0], source_weight)
+
+    def test_repeated_hook_lifecycle_still_works(self) -> None:
+        """Guards against a future assumption that every internally-used
+        GeneralizedComponent must be registered -- the dispatcher's own
+        hooks are, and that's what reset_hooks()/run_with_cache rely on."""
+        bridge = self._build()
+        tokens = torch.randint(0, 64, (1, 4))
+        with torch.no_grad():
+            bridge.run_with_cache(tokens)
+            bridge.reset_hooks()
+            _, cache = bridge.run_with_cache(tokens)
+        assert "blocks.0.mlp.hook_out" in cache
+
+
+class TestTrainEvalModePropagation:
+    """bridge.train()/.eval() propagate to the wrapped model:
+    build_pretrain_bridge reassigns the returned bridge's class to a small
+    generated TransformerBridge subclass (see
+    _get_mode_propagating_bridge_class) whose overridden train() also sets
+    mode on original_model; inherited eval() calls train(False), so it
+    reaches the override too without needing its own override. This
+    closes a gap in TransformerBridge's own train()/eval(), which only
+    walk component_mapping and never reach original_model on their own."""
+
+    def test_bridge_eval_propagates_to_wrapped_model(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_pretrain_bridge(model, cfg)
+
+        assert model.training is True
+        bridge.eval()
+        assert model.training is False
+
+    def test_bridge_train_propagates_to_wrapped_model(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_pretrain_bridge(model, cfg)
+
+        model.eval()
+        assert model.training is False
+        bridge.train()
+        assert model.training is True
+
+    def test_caller_can_still_set_mode_directly_on_their_own_model_reference(self) -> None:
+        """Setting mode via the raw model reference still works and stays
+        in sync -- both paths write to the same underlying module."""
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        build_pretrain_bridge(model, cfg)
+
+        model.eval()
+        assert model.training is False
+        model.train()
+        assert model.training is True
+
+    def test_train_and_eval_return_the_bridge_itself(self) -> None:
+        """nn.Module convention: train()/eval() return self, so callers can
+        chain (`model.train().to(device)`, etc). The generated subclass
+        must preserve this, not just the mode-propagation side effect."""
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_pretrain_bridge(model, cfg)
+
+        assert bridge.train(False) is bridge
+        assert bridge.train(True) is bridge
+        assert bridge.eval() is bridge
+
+    def test_bridge_can_be_reconstructed_from_supported_state(self) -> None:
+        """The adapter's actual persistence contract is: rebuild the
+        bridge from the source model class + cfg + the source model's
+        *raw* state_dict, then confirm behavior (outputs, train/eval
+        propagation) survives the round trip.
+
+        The state_dict must be snapshotted BEFORE `build_pretrain_bridge`
+        is called. Building the bridge wraps the model's submodules in
+        place (each one gains an `_original_component` wrapper for
+        hooking), so `model.state_dict()` taken *after* wrapping has
+        keys like `embed._original_component.weight` instead of
+        `embed.weight` -- a fresh, unwrapped `TinyPretrainModel` can't
+        load that. Capturing the state_dict first, then wrapping a
+        second fresh model and loading the pre-wrap snapshot into it,
+        avoids that trap.
+
+        Only the source model's state_dict is used -- not
+        `bridge.state_dict()` -- since the bridge doesn't own an
+        independent set of learned weights; it exposes the wrapped
+        source model's parameters under bridge-shaped names.
+
+        This is deliberately weaker than whole-object `pickle.dumps`,
+        which would also require every nested `TransformerBridge`
+        implementation detail (e.g. `AttentionBridge`'s hook
+        conversions) to be picklable -- a guarantee unrelated to this
+        adapter and not part of what it needs to promise. If a stronger
+        whole-object-pickle guarantee is later required, it belongs in a
+        `TransformerBridge`-level test, not here.
+        """
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+
+        # Snapshot BEFORE wrapping -- see docstring.
+        raw_model_state = {key: value.detach().clone() for key, value in model.state_dict().items()}
+
+        bridge = build_pretrain_bridge(model, cfg)
+        bridge.eval()
+
+        restored_model = TinyPretrainModel(
+            d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64
+        )
+        restored_model.load_state_dict(raw_model_state)
+
+        restored_bridge = build_pretrain_bridge(restored_model, cfg)
+        restored_bridge.eval()
+
+        tokens = torch.randint(0, 64, (1, 4))
+        with torch.no_grad():
+            expected = bridge(tokens)
+            actual = restored_bridge(tokens)
+
+        expected_logits = expected.logits if hasattr(expected, "logits") else expected
+        actual_logits = actual.logits if hasattr(actual, "logits") else actual
+        torch.testing.assert_close(actual_logits, expected_logits, atol=0, rtol=0)
+
+        assert restored_bridge.training is False
+        assert restored_model.training is False
+
+        restored_bridge.train()
+        assert restored_model.training is True
+
+    def test_repeated_construction_reuses_the_same_generated_class(self) -> None:
+        """Guards the cache's stated purpose: build_pretrain_bridge should
+        not generate a new mode-propagating subclass on every call -- two
+        independently-built bridges from the same base bridge class must
+        share one generated class."""
+        cfg_a = _make_cfg(n_layers=1)
+        model_a = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge_a = build_pretrain_bridge(model_a, cfg_a)
+
+        cfg_b = _make_cfg(n_layers=1)
+        model_b = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge_b = build_pretrain_bridge(model_b, cfg_b)
+
+        assert type(bridge_a) is type(bridge_b)
+
+
+class TestDtypeAndTiedWeightPreservation:
+    def test_non_default_dtype_and_tied_weights_survive_construction(self) -> None:
+        torch.manual_seed(0)
+        model = TinyPretrainModel(
+            d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=True
+        ).to(torch.float64)
+        assert model.lm_head.weight is model.embed.weight
+
+        cfg = _make_cfg(n_layers=1)
+        bridge = build_pretrain_bridge(model, cfg)
+
+        assert next(model.parameters()).dtype == torch.float64
+        assert model.lm_head.weight is model.embed.weight
+
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.dtype == torch.float64
+
+
+class TestBuildPretrainBridgeErgonomics:
+    def test_explicit_model_name_is_forwarded(self) -> None:
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_pretrain_bridge(model, cfg, model_name="my-custom-name")
+        assert bridge.cfg.model_name == "my-custom-name"
+
+    def test_omitting_optional_kwargs_still_works(self) -> None:
+        """The default (no device/dtype/model_name passed) path, exercised
+        everywhere else in this file, must keep working unchanged."""
+        cfg = _make_cfg(n_layers=1)
+        model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64)
+        bridge = build_pretrain_bridge(model, cfg)
+        tokens = torch.randint(0, 64, (1, 3))
+        with torch.no_grad():
+            output = bridge(tokens)
+        assert output.shape == (1, 3, 64)
+
+
+class TestPretrainModelContainerOutputContract:
+    def test_container_normalizes_dict_output_to_logits_attr_dict(self) -> None:
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+        )
+        container = PretrainModelContainer(model)
+        tokens = torch.randint(0, 64, (1, 3))
+
+        with torch.no_grad():
+            output = container(tokens)
+
+        assert isinstance(output, _LogitsAttrDict)
+        assert output.logits is output["logits"]
+        assert output.logits.shape == (1, 3, 64)
+
+    def test_container_returns_a_fresh_wrapper_on_each_call(self) -> None:
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+        )
+        container = PretrainModelContainer(model)
+        tokens = torch.randint(0, 64, (1, 3))
+
+        with torch.no_grad():
+            out1 = container(tokens)
+            out2 = container(tokens)
+
+        assert isinstance(out1, _LogitsAttrDict)
+        assert isinstance(out2, _LogitsAttrDict)
+        assert out1 is not out2
+        assert out1["logits"] is not out2["logits"]
+        torch.testing.assert_close(out1["logits"], out2["logits"])
+
+    def test_container_does_not_reuse_or_mutate_output_mapping(self) -> None:
+        model = TinyPretrainModel(
+            d_model=16,
+            n_heads=2,
+            n_layers=1,
+            d_ff=32,
+            vocab_size=64,
+        ).eval()
+        container = PretrainModelContainer(model)
+        tokens = torch.randint(0, 64, (1, 3))
+
+        with torch.no_grad():
+            out1 = container(tokens)
+
+        out1["sentinel"] = True
+
+        with torch.no_grad():
+            out2 = container(tokens)
+
+        assert out1 is not out2
+        assert "sentinel" not in out2
+        torch.testing.assert_close(out1["logits"], out2["logits"])
+
+    def test_dict_missing_logits_key_raises_value_error(self) -> None:
+        class ForwardWrongKey(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> dict:
+                return {"scores": tokens.float()}
+
+        container = PretrainModelContainer(ForwardWrongKey())
+        with pytest.raises(ValueError, match="'logits'"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_dict_with_heterogeneous_keys_still_raises_value_error(self) -> None:
+        """sorted(output.keys()) would raise TypeError on non-comparable
+        mixed key types (e.g. str and int); the error message must use
+        list() instead, so an invalid model output produces the intended
+        ValueError, not an unrelated sorting error."""
+
+        class ForwardHeterogeneousKeys(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> dict:
+                return {"scores": tokens.float(), 1: "metadata"}
+
+        container = PretrainModelContainer(ForwardHeterogeneousKeys())
+        with pytest.raises(ValueError, match="'logits'"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_dict_with_non_tensor_logits_raises_type_error(self) -> None:
+        class ForwardBadLogitsType(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> dict:
+                return {"logits": "not a tensor"}
+
+        container = PretrainModelContainer(ForwardBadLogitsType())
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_already_logits_attr_dict_passes_through_by_identity(self) -> None:
+        """An already-wrapped output must not be re-wrapped."""
+        already_wrapped = _LogitsAttrDict({"logits": torch.zeros(1)})
+
+        class ForwardAlreadyWrapped(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> _LogitsAttrDict:
+                return already_wrapped
+
+        container = PretrainModelContainer(ForwardAlreadyWrapped())
+        output = container(torch.tensor([[1, 2, 3]]))
+        assert output is already_wrapped
+
+    def test_malformed_logits_attr_dict_raises_value_error_not_keyerror(self) -> None:
+        """Not a realistic target-model failure (_LogitsAttrDict is
+        private), but closes the contract: a malformed instance gets the
+        same clear ValueError an ordinary dict missing 'logits' gets,
+        rather than a raw KeyError leaking out."""
+        malformed = _LogitsAttrDict({"scores": torch.zeros(1)})
+
+        class ForwardMalformedWrapped(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> _LogitsAttrDict:
+                return malformed
+
+        container = PretrainModelContainer(ForwardMalformedWrapped())
+        with pytest.raises(ValueError, match="'logits'"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_hf_style_object_with_logits_attribute_passes_through_unchanged(self) -> None:
+        """A source model that already satisfies hasattr(output, 'logits')
+        on its own (an HF-ModelOutput-like object, not a dict) needs no
+        normalization at all."""
+
+        class FakeModelOutput:
+            def __init__(self, logits: torch.Tensor):
+                self.logits = logits
+
+        already_hf_style = FakeModelOutput(torch.zeros(1))
+
+        class ForwardHFStyle(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> FakeModelOutput:
+                return already_hf_style
+
+        container = PretrainModelContainer(ForwardHFStyle())
+        output = container(torch.tensor([[1, 2, 3]]))
+        assert output is already_hf_style
+
+    def test_property_backed_logits_is_only_evaluated_once(self) -> None:
+        """.logits must be read once and cached locally, not re-evaluated
+        on every access -- matters for a property-backed .logits, which
+        could otherwise do redundant (or side-effecting) work per access."""
+        access_count = 0
+
+        class PropertyBackedOutput:
+            @property
+            def logits(self) -> torch.Tensor:
+                nonlocal access_count
+                access_count += 1
+                return torch.zeros(1)
+
+        class ForwardPropertyBacked(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> PropertyBackedOutput:
+                return PropertyBackedOutput()
+
+        container = PretrainModelContainer(ForwardPropertyBacked())
+        container(torch.tensor([[1, 2, 3]]))
+        assert access_count == 1
+
+    def test_object_with_non_tensor_logits_attribute_raises_type_error(self) -> None:
+        class BadModelOutput:
+            logits = "not a tensor"
+
+        class ForwardBadModelOutput(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> BadModelOutput:
+                return BadModelOutput()
+
+        container = PretrainModelContainer(ForwardBadModelOutput())
+        with pytest.raises(TypeError, match=r"\.logits"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_bare_tensor_output_passes_through_unchanged(self) -> None:
+        class ForwardBareTensor(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+                return tokens.float()
+
+        container = PretrainModelContainer(ForwardBareTensor())
+        tokens = torch.tensor([[1, 2, 3]])
+        output = container(tokens)
+        assert torch.equal(output, tokens.float())
+
+    def test_tuple_with_tensor_first_element_passes_through_unchanged(self) -> None:
+        """TransformerBridge extracts output[0] as logits for tuple
+        returns -- no normalization needed."""
+
+        class ForwardTuple(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> tuple:
+                return (tokens.float(), "some_aux_value")
+
+        container = PretrainModelContainer(ForwardTuple())
+        output = container(torch.tensor([[1, 2, 3]]))
+        assert isinstance(output, tuple)
+        assert output[1] == "some_aux_value"
+
+    def test_tuple_with_non_tensor_first_element_raises_type_error(self) -> None:
+        class ForwardBadTuple(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> tuple:
+                return ("not a tensor", tokens)
+
+        container = PretrainModelContainer(ForwardBadTuple())
+        with pytest.raises(TypeError, match="torch.Tensor"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_list_output_raises_type_error(self) -> None:
+        class ForwardList(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> list:
+                return [tokens.float()]
+
+        container = PretrainModelContainer(ForwardList())
+        with pytest.raises(TypeError, match="must return"):
+            container(torch.tensor([[1, 2, 3]]))
+
+    def test_string_output_raises_type_error(self) -> None:
+        class ForwardString(nn.Module):
+            def forward(self, tokens: torch.Tensor) -> str:
+                return "wrong"
+
+        container = PretrainModelContainer(ForwardString())
+        with pytest.raises(TypeError, match="must return"):
+            container(torch.tensor([[1, 2, 3]]))
diff --git a/tests/unit/model_bridge/test_bridge_train_mode_propagation.py b/tests/unit/model_bridge/test_bridge_train_mode_propagation.py
new file mode 100644
index 000000000..4d08d0c71
--- /dev/null
+++ b/tests/unit/model_bridge/test_bridge_train_mode_propagation.py
@@ -0,0 +1,100 @@
+"""Regression tests for TransformerBridge train/eval mode propagation.
+
+`original_model` is stored outside the registered module tree, so
+inherited `nn.Module.train()` does not reach it. These tests verify that
+bridge mode changes propagate to the wrapped model, whose mode-dependent
+layers (Dropout) rely on that flag. Architecture-independent stub bridge;
+no HF Hub access, no weights.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components import LinearBridge
+
+
+class _StubModel(nn.Module):
+    """Minimal source model with a mode-dependent layer (Dropout), so
+    train/eval propagation has an observable behavioral consequence, not
+    just a flag."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.proj = nn.Linear(4, 4)
+        self.drop = nn.Dropout(0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(self.proj(x))
+
+
+class _StubAdapter(ArchitectureAdapter):
+    """Smallest adapter TransformerBridge will accept: a one-entry
+    component_mapping. The key is `stub_proj` (not `proj` or `embed`)
+    because TransformerBridge reserves several canonical attribute names
+    on itself; a colliding key fails at add_module time."""
+
+    def __init__(self, cfg: TransformerBridgeConfig) -> None:
+        super().__init__(cfg)
+        self.component_mapping = {"stub_proj": LinearBridge(name="proj")}
+
+
+def _make_stub_bridge() -> tuple[TransformerBridge, _StubModel]:
+    cfg = TransformerBridgeConfig(
+        d_model=4,
+        d_head=2,
+        n_layers=1,
+        n_ctx=8,
+        n_heads=2,
+        d_vocab=8,
+        architecture="StubForTest",
+    )
+    model = _StubModel()
+    bridge = TransformerBridge(model, _StubAdapter(cfg), tokenizer=None)
+    return bridge, model
+
+
+class TestTrainEvalModePropagation:
+    """bridge.train()/.eval() must reach original_model.
+
+    Without the TransformerBridge.train() override these fail with
+    original_model.training stuck at its wrap-time value (bridge.training
+    flips, the wrapped model's does not) -- i.e. bridge.eval() would leave
+    dropout active."""
+
+    def test_eval_propagates_to_original_model(self) -> None:
+        bridge, model = _make_stub_bridge()
+        assert model.training is True  # nn.Module default at wrap time
+        bridge.eval()
+        assert bridge.training is False
+        assert model.training is False
+
+    def test_train_propagates_to_original_model(self) -> None:
+        bridge, model = _make_stub_bridge()
+        bridge.eval()
+        assert model.training is False
+        bridge.train()
+        assert bridge.training is True
+        assert model.training is True
+
+    def test_train_and_eval_return_the_bridge_itself(self) -> None:
+        """nn.Module convention: train()/eval() return self so callers can
+        chain (`bridge.eval().to(device)`, etc.)."""
+        bridge, _ = _make_stub_bridge()
+        assert bridge.train(False) is bridge
+        assert bridge.train(True) is bridge
+        assert bridge.eval() is bridge
+
+    def test_direct_mode_on_original_model_stays_in_sync(self) -> None:
+        """Setting mode via the caller's own reference to the source model
+        still works -- both paths write the same underlying flags, so
+        neither clobbers the other."""
+        bridge, model = _make_stub_bridge()
+        bridge.eval()
+        model.train()
+        assert model.training is True
+        bridge.eval()
+        assert model.training is False
diff --git a/tests/unit/tools/test_model_registry.py b/tests/unit/tools/test_model_registry.py
index 8fbb12c01..53627c8fa 100644
--- a/tests/unit/tools/test_model_registry.py
+++ b/tests/unit/tools/test_model_registry.py
@@ -807,6 +807,7 @@ class TestRegistrySyncedWithFactory:
             "NeelSoluOldForCausalLM",
             "GPT2LMHeadCustomModel",
             "TransformerLensNative",
+            "TransformerLensPretrain",
             # Group 2: factory-internal alias casings (HF emits the canonical name).
             "Gemma1ForCausalLM",  # HF emits: GemmaForCausalLM
             "NeoForCausalLM",  # HF emits: GPTNeoForCausalLM
diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py
index bfe1b90ff..0e7eba9bb 100644
--- a/transformer_lens/factories/architecture_adapter_factory.py
+++ b/transformer_lens/factories/architecture_adapter_factory.py
@@ -11,20 +11,15 @@
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.supported_architectures import (
     ApertusArchitectureAdapter,
-    ArceeArchitectureAdapter,
     BaichuanArchitectureAdapter,
     BartArchitectureAdapter,
-    BD3LMArchitectureAdapter,
     BertArchitectureAdapter,
     BloomArchitectureAdapter,
     CodeGenArchitectureAdapter,
-    Cohere2ArchitectureAdapter,
     CohereArchitectureAdapter,
     DeepSeekV2ArchitectureAdapter,
     DeepSeekV3ArchitectureAdapter,
-    DeepSeekV4ArchitectureAdapter,
     FalconArchitectureAdapter,
-    FalconH1ArchitectureAdapter,
     Gemma1ArchitectureAdapter,
     Gemma2ArchitectureAdapter,
     Gemma3ArchitectureAdapter,
@@ -44,9 +39,7 @@
     HubertArchitectureAdapter,
     HunYuanDenseV1ArchitectureAdapter,
     InternLM2ArchitectureAdapter,
-    Lfm2ArchitectureAdapter,
     Lfm2MoeArchitectureAdapter,
-    LLaDAArchitectureAdapter,
     LlamaArchitectureAdapter,
     LlavaArchitectureAdapter,
     LlavaNextArchitectureAdapter,
@@ -69,46 +62,37 @@
     OlmoeArchitectureAdapter,
     OpenElmArchitectureAdapter,
     OptArchitectureAdapter,
-    OuroArchitectureAdapter,
     Phi3ArchitectureAdapter,
     PhiArchitectureAdapter,
     PhiMoEArchitectureAdapter,
+    PretrainArchitectureAdapter,
     Qwen2ArchitectureAdapter,
-    Qwen2MoeArchitectureAdapter,
     Qwen3_5ArchitectureAdapter,
     Qwen3_5MultimodalArchitectureAdapter,
     Qwen3ArchitectureAdapter,
     Qwen3MoeArchitectureAdapter,
     Qwen3NextArchitectureAdapter,
     QwenArchitectureAdapter,
-    RecurrentGemmaArchitectureAdapter,
     SmolLM3ArchitectureAdapter,
     StableLmArchitectureAdapter,
     T5ArchitectureAdapter,
-    T5Gemma2ArchitectureAdapter,
     T5GemmaArchitectureAdapter,
     XGLMArchitectureAdapter,
-    Zamba2ArchitectureAdapter,
 )
 
 # Export supported architectures
 SUPPORTED_ARCHITECTURES = {
     "ApertusForCausalLM": ApertusArchitectureAdapter,
-    "ArceeForCausalLM": ArceeArchitectureAdapter,
     "BaiChuanForCausalLM": BaichuanArchitectureAdapter,
     "BaichuanForCausalLM": BaichuanArchitectureAdapter,
     "BartForConditionalGeneration": BartArchitectureAdapter,
-    "BD3LM": BD3LMArchitectureAdapter,
     "BertForMaskedLM": BertArchitectureAdapter,
     "BloomForCausalLM": BloomArchitectureAdapter,
     "CodeGenForCausalLM": CodeGenArchitectureAdapter,
-    "Cohere2ForCausalLM": Cohere2ArchitectureAdapter,
     "CohereForCausalLM": CohereArchitectureAdapter,
     "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter,
     "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter,
-    "DeepseekV4ForCausalLM": DeepSeekV4ArchitectureAdapter,
     "FalconForCausalLM": FalconArchitectureAdapter,
-    "FalconH1ForCausalLM": FalconH1ArchitectureAdapter,
     "GemmaForCausalLM": Gemma1ArchitectureAdapter,  # Default to Gemma1 as it's the original version
     "Gemma1ForCausalLM": Gemma1ArchitectureAdapter,
     "Gemma2ForCausalLM": Gemma2ArchitectureAdapter,
@@ -134,12 +118,10 @@
     "HubertModel": HubertArchitectureAdapter,
     "HunYuanDenseV1ForCausalLM": HunYuanDenseV1ArchitectureAdapter,
     "InternLM2ForCausalLM": InternLM2ArchitectureAdapter,
-    "LLaDAModelLM": LLaDAArchitectureAdapter,
     "LlamaForCausalLM": LlamaArchitectureAdapter,
     "LlavaForConditionalGeneration": LlavaArchitectureAdapter,
     "LlavaNextForConditionalGeneration": LlavaNextArchitectureAdapter,
     "LlavaOnevisionForConditionalGeneration": LlavaOnevisionArchitectureAdapter,
-    "Lfm2ForCausalLM": Lfm2ArchitectureAdapter,
     "Lfm2MoeForCausalLM": Lfm2MoeArchitectureAdapter,
     "Mamba2ForCausalLM": Mamba2ArchitectureAdapter,
     "MambaForCausalLM": MambaArchitectureAdapter,
@@ -156,29 +138,25 @@
     "OlmoeForCausalLM": OlmoeArchitectureAdapter,
     "OpenELMForCausalLM": OpenElmArchitectureAdapter,
     "OPTForCausalLM": OptArchitectureAdapter,
-    "OuroForCausalLM": OuroArchitectureAdapter,
     "PhiForCausalLM": PhiArchitectureAdapter,
     "Phi3ForCausalLM": Phi3ArchitectureAdapter,
     "PhiMoEForCausalLM": PhiMoEArchitectureAdapter,
     "QwenForCausalLM": QwenArchitectureAdapter,
     "Qwen2ForCausalLM": Qwen2ArchitectureAdapter,
-    "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter,
     "Qwen3ForCausalLM": Qwen3ArchitectureAdapter,
     "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter,
     "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter,
     "Qwen3_5ForCausalLM": Qwen3_5ArchitectureAdapter,
     "Qwen3_5ForConditionalGeneration": Qwen3_5MultimodalArchitectureAdapter,
-    "RecurrentGemmaForCausalLM": RecurrentGemmaArchitectureAdapter,
     "SmolLM3ForCausalLM": SmolLM3ArchitectureAdapter,
     "StableLmForCausalLM": StableLmArchitectureAdapter,
     "T5ForConditionalGeneration": T5ArchitectureAdapter,
     "MT5ForConditionalGeneration": T5ArchitectureAdapter,
     "T5GemmaForConditionalGeneration": T5GemmaArchitectureAdapter,
-    "T5Gemma2ForConditionalGeneration": T5Gemma2ArchitectureAdapter,
     "XGLMForCausalLM": XGLMArchitectureAdapter,
-    "Zamba2ForCausalLM": Zamba2ArchitectureAdapter,
     "NanoGPTForCausalLM": NanogptArchitectureAdapter,
     "TransformerLensNative": NativeArchitectureAdapter,
+    "TransformerLensPretrain": PretrainArchitectureAdapter,
     "MinGPTForCausalLM": MingptArchitectureAdapter,
     "GPTNeoForCausalLM": NeoArchitectureAdapter,
     "GPTNeoXForCausalLM": NeoxArchitectureAdapter,
diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py
index 3c470cf72..239eaa363 100644
--- a/transformer_lens/model_bridge/bridge.py
+++ b/transformer_lens/model_bridge/bridge.py
@@ -50,7 +50,6 @@
     VARIANT_SUBMODULE_NAMES,
 )
 from transformer_lens.model_bridge.get_params_util import get_bridge_params
-from transformer_lens.utilities.activation_functions import softcap_enabled
 from transformer_lens.utilities.aliases import resolve_alias
 from transformer_lens.utilities.devices import move_to_and_update_config
 from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss
@@ -220,7 +219,7 @@ def boot_transformers(
         hf_model: Optional[Any] = None,
         device_map: Optional[Union[str, Dict[str, Union[str, int]]]] = None,
         n_devices: Optional[int] = None,
-        max_memory: Optional[Dict[Union[str, int], Union[str, int]]] = None,
+        max_memory: Optional[Dict[Union[str, int], str]] = None,
         n_ctx: Optional[int] = None,
         revision: Optional[str] = None,
         checkpoint_index: Optional[int] = None,
@@ -921,7 +920,7 @@ def process_weights(
             print(f"Processing weights for {self.cfg.model_name}...")
 
         # Soft capping (tanh) is not translation-invariant; centering would change output.
-        if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)):
+        if center_unembed and getattr(self.cfg, "output_logits_soft_cap", -1.0) > 0.0:
             import logging
 
             logging.warning(
@@ -1747,14 +1746,6 @@ def forward(
             Model output based on return_type
         """
 
-        if return_type in ("loss", "both") and not self.adapter.supports_causal_loss:
-            architecture = self.cfg.architecture or type(self.adapter).__name__
-            raise NotImplementedError(
-                f"{architecture} does not support TransformerBridge's shifted causal "
-                "loss. Request return_type='logits' and compute the architecture-specific "
-                "masked-token objective explicitly."
-            )
-
         if start_at_layer is not None:
             raise NotImplementedError(
                 "start_at_layer is not supported in TransformerBridge. "
@@ -2198,10 +2189,7 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
                 stacklevel=2,
             )
         try:
-            if (
-                "output_attentions" not in filtered_kwargs
-                and self.adapter.supports_hf_output_attentions
-            ):
+            if "output_attentions" not in filtered_kwargs:
                 filtered_kwargs["output_attentions"] = True
             if processed_args:
                 output = self.forward(processed_args[0], **filtered_kwargs)
@@ -2683,15 +2671,6 @@ def _generate_tokens(
             if all_finished:
                 return
 
-    def _ensure_generation_supported(self, api_name: str) -> None:
-        """Reject autoregressive generation for forward-only architectures."""
-        if not self.adapter.supports_generation:
-            architecture = self.cfg.architecture or type(self.adapter).__name__
-            raise NotImplementedError(
-                f"TransformerBridge.{api_name}() generation is not supported by "
-                f"the {architecture} architecture."
-            )
-
     def generate(
         self,
         input: Union[str, List[str], torch.Tensor] = "",
@@ -2811,7 +2790,6 @@ def generate(
             the tokens that were fed to the model, useful for verifying BOS handling with
             chat templates.
         """
-        self._ensure_generation_supported("generate")
         # padding_side is handled internally: for batched list inputs, left-padding
         # is forced to ensure correct generation. See _is_batched_list logic below.
 
@@ -3228,7 +3206,6 @@ def generate_stream(
             max_tokens_per_yield tokens between yields. First yield includes
             the input tokens; subsequent yields contain only new tokens.
         """
-        self._ensure_generation_supported("generate_stream")
         # --- Input parsing (mirrors generate()) ---
         _is_batched_list = isinstance(input, list) and len(input) > 1
 
@@ -3460,7 +3437,6 @@ def hf_generate(
             print(result.logits)  # Logits for each generation step
             print(result.attentions)  # Attention weights
         """
-        self._ensure_generation_supported("hf_generate")
         # Handle string input by tokenizing it
         if isinstance(input, str):
             inputs = self.tokenizer(input, return_tensors="pt", padding=False, truncation=False).to(
@@ -3703,6 +3679,33 @@ def mps(self) -> "TransformerBridge":
         """
         return self.to(torch.device("mps"))
 
+    def train(self, mode: bool = True) -> "TransformerBridge":
+        """Set training mode, propagating to the wrapped source model.
+
+        ``nn.Module.train()`` only recurses into *registered* submodules.
+        ``original_model`` is intentionally stored outside the registered
+        module tree (assigned via ``self.__dict__`` rather than
+        ``nn.Module.__setattr__``); consequently, inherited ``train()``
+        does not reach it -- without this override, ``bridge.eval()`` /
+        ``bridge.train()`` silently leave dropout and other mode-dependent
+        layers in whatever mode the source model was in at wrap time.
+
+        ``eval()`` needs no override of its own: ``nn.Module.eval()`` calls
+        ``self.train(False)``, which dispatches here.
+
+        Args:
+            mode:
+                ``True`` for training mode, ``False`` for evaluation mode.
+
+        Returns:
+            The bridge itself (``nn.Module`` chaining convention).
+        """
+        super().train(mode)
+        original = getattr(self, "original_model", None)
+        if isinstance(original, torch.nn.Module):
+            original.train(mode)
+        return self
+
     def add_hook(
         self,
         name: Union[str, Callable[[str], bool]],
@@ -3966,25 +3969,18 @@ def _validate_attention_fork_supported(self, flag_name: str) -> None:
                 continue
             total_with_attn += 1
             attn_classes.add(type(attn).__name__)
-            supports_flag = (
-                bool(getattr(attn, "supports_attn_result", False))
-                if flag_name == "use_attn_result"
-                else isinstance(attn, supported_classes)
-            )
-            if supports_flag:
+            if isinstance(attn, supported_classes):
                 supporting_layers.append(idx)
         if total_with_attn == 0:
             raise NotImplementedError(f"{flag_name}: no attention bridges found on self.blocks.")
         if not supporting_layers:
-            if flag_name == "use_attn_result":
-                capability_detail = "Per-head result computation is unavailable."
-            else:
-                capability_detail = "No hook point is available before the Q/K/V projections."
             raise NotImplementedError(
                 f"{flag_name}: none of this model's attention bridges support "
-                "the requested fine-grained attention hook. Found attention classes: "
+                "the fine-grained Q/K/V hook fork. Found attention classes: "
                 f"{sorted(attn_classes)}. Supported classes: "
-                f"{[c.__name__ for c in supported_classes]}. {capability_detail}"
+                f"{[c.__name__ for c in supported_classes]}. Plain "
+                "AttentionBridge delegates to HuggingFace and exposes no hook "
+                "point before the Q/K/V projection."
             )
         if len(supporting_layers) < total_with_attn:
             skipped = total_with_attn - len(supporting_layers)
diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py
index 18bbdac7a..67a677ffc 100644
--- a/transformer_lens/model_bridge/generalized_components/moe.py
+++ b/transformer_lens/model_bridge/generalized_components/moe.py
@@ -2,6 +2,7 @@
 
 This module contains the bridge component for Mixture of Experts layers.
 """
+
 from __future__ import annotations
 
 from typing import Any, Dict, Optional
@@ -110,10 +111,24 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
             kwargs = {**kwargs, "hidden_states": hooked}
         output = self.original_component(*args, **kwargs)
         if isinstance(output, tuple):
+            if not output:
+                raise TypeError(
+                    f"{self.name}: expected a non-empty tuple whose first "
+                    "element is a torch.Tensor from the wrapped MoE component, "
+                    "got an empty tuple."
+                )
+
             hidden_states = output[0]
+            if not isinstance(hidden_states, torch.Tensor):
+                raise TypeError(
+                    f"{self.name}: expected the first tuple element from the "
+                    f"wrapped MoE component to be a torch.Tensor, got "
+                    f"{type(hidden_states).__name__}."
+                )
             if len(output) > 1:
                 router_scores = output[1]
-                self.hook_router_scores(router_scores)
+                if isinstance(router_scores, torch.Tensor):
+                    self.hook_router_scores(router_scores)
             hidden_states = self.hook_out(hidden_states)
             return (hidden_states,) + output[1:]
         else:
diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py
index 0eac330e2..d886cadec 100644
--- a/transformer_lens/model_bridge/supported_architectures/__init__.py
+++ b/transformer_lens/model_bridge/supported_architectures/__init__.py
@@ -6,18 +6,12 @@
 from transformer_lens.model_bridge.supported_architectures.apertus import (
     ApertusArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.arcee import (
-    ArceeArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.baichuan import (
     BaichuanArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.bart import (
     BartArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.bd3lm import (
-    BD3LMArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.bert import (
     BertArchitectureAdapter,
 )
@@ -28,7 +22,6 @@
     CodeGenArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.cohere import (
-    Cohere2ArchitectureAdapter,
     CohereArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.deepseek_v2 import (
@@ -37,15 +30,9 @@
 from transformer_lens.model_bridge.supported_architectures.deepseek_v3 import (
     DeepSeekV3ArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import (
-    DeepSeekV4ArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.falcon import (
     FalconArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.falcon_h1 import (
-    FalconH1ArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.gemma1 import (
     Gemma1ArchitectureAdapter,
 )
@@ -67,9 +54,6 @@
 from transformer_lens.model_bridge.supported_architectures.glm4_moe import (
     Glm4MoeArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import (
-    GlmMoeDsaArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.gpt2 import (
     GPT2ArchitectureAdapter,
 )
@@ -85,6 +69,9 @@
 from transformer_lens.model_bridge.supported_architectures.gptj import (
     GptjArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import (
+    GlmMoeDsaArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.granite import (
     GraniteArchitectureAdapter,
 )
@@ -103,15 +90,9 @@
 from transformer_lens.model_bridge.supported_architectures.internlm2 import (
     InternLM2ArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.lfm2 import (
-    Lfm2ArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.lfm2_moe import (
     Lfm2MoeArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.llada import (
-    LLaDAArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.llama import (
     LlamaArchitectureAdapter,
 )
@@ -178,9 +159,6 @@
 from transformer_lens.model_bridge.supported_architectures.opt import (
     OptArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.ouro import (
-    OuroArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.phi import (
     PhiArchitectureAdapter,
 )
@@ -190,15 +168,15 @@
 from transformer_lens.model_bridge.supported_architectures.phimoe import (
     PhiMoEArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.pretrain import (
+    PretrainArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.qwen import (
     QwenArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.qwen2 import (
     Qwen2ArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.qwen2_moe import (
-    Qwen2MoeArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.qwen3 import (
     Qwen3ArchitectureAdapter,
 )
@@ -214,9 +192,6 @@
 from transformer_lens.model_bridge.supported_architectures.qwen3_next import (
     Qwen3NextArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.recurrent_gemma import (
-    RecurrentGemmaArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.smollm3 import (
     SmolLM3ArchitectureAdapter,
 )
@@ -229,31 +204,21 @@
 from transformer_lens.model_bridge.supported_architectures.t5gemma import (
     T5GemmaArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.t5gemma2 import (
-    T5Gemma2ArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.xglm import (
     XGLMArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.zamba2 import (
-    Zamba2ArchitectureAdapter,
-)
 
 __all__ = [
     "ApertusArchitectureAdapter",
-    "ArceeArchitectureAdapter",
-    "BD3LMArchitectureAdapter",
     "BaichuanArchitectureAdapter",
     "BartArchitectureAdapter",
     "BertArchitectureAdapter",
     "BloomArchitectureAdapter",
     "CodeGenArchitectureAdapter",
-    "Cohere2ArchitectureAdapter",
     "CohereArchitectureAdapter",
     "DeepSeekV2ArchitectureAdapter",
     "DeepSeekV3ArchitectureAdapter",
     "FalconArchitectureAdapter",
-    "FalconH1ArchitectureAdapter",
     "Gemma1ArchitectureAdapter",
     "Gemma2ArchitectureAdapter",
     "Gemma3ArchitectureAdapter",
@@ -273,12 +238,10 @@
     "HubertArchitectureAdapter",
     "HunYuanDenseV1ArchitectureAdapter",
     "InternLM2ArchitectureAdapter",
-    "LLaDAArchitectureAdapter",
     "LlamaArchitectureAdapter",
     "LlavaArchitectureAdapter",
     "LlavaNextArchitectureAdapter",
     "LlavaOnevisionArchitectureAdapter",
-    "Lfm2ArchitectureAdapter",
     "Lfm2MoeArchitectureAdapter",
     "MambaArchitectureAdapter",
     "Mamba2ArchitectureAdapter",
@@ -287,7 +250,6 @@
     "MixtralArchitectureAdapter",
     "MPTArchitectureAdapter",
     "NanogptArchitectureAdapter",
-    "NemotronHArchitectureAdapter",
     "NativeArchitectureAdapter",
     "NeelSoluOldArchitectureAdapter",
     "NemotronHArchitectureAdapter",
@@ -299,24 +261,20 @@
     "Olmo3ArchitectureAdapter",
     "OlmoeArchitectureAdapter",
     "OptArchitectureAdapter",
-    "OuroArchitectureAdapter",
     "PhiArchitectureAdapter",
     "Phi3ArchitectureAdapter",
     "PhiMoEArchitectureAdapter",
+    "PretrainArchitectureAdapter",
     "QwenArchitectureAdapter",
     "Qwen2ArchitectureAdapter",
-    "Qwen2MoeArchitectureAdapter",
     "Qwen3ArchitectureAdapter",
     "Qwen3MoeArchitectureAdapter",
     "Qwen3NextArchitectureAdapter",
     "Qwen3_5ArchitectureAdapter",
     "Qwen3_5MultimodalArchitectureAdapter",
-    "RecurrentGemmaArchitectureAdapter",
     "SmolLM3ArchitectureAdapter",
     "StableLmArchitectureAdapter",
     "T5ArchitectureAdapter",
     "T5GemmaArchitectureAdapter",
-    "T5Gemma2ArchitectureAdapter",
     "XGLMArchitectureAdapter",
-    "Zamba2ArchitectureAdapter",
 ]
diff --git a/transformer_lens/model_bridge/supported_architectures/pretrain.py b/transformer_lens/model_bridge/supported_architectures/pretrain.py
new file mode 100644
index 000000000..20091b65f
--- /dev/null
+++ b/transformer_lens/model_bridge/supported_architectures/pretrain.py
@@ -0,0 +1,442 @@
+"""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`.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components import (
+    AttentionBridge,
+    DelegatedAttentionBlockBridge,
+    EmbeddingBridge,
+    GatedMLPBridge,
+    LinearBridge,
+    MoEBridge,
+    RMSNormalizationBridge,
+    UnembeddingBridge,
+)
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+ARCHITECTURE_NAME = "TransformerLensPretrain"
+
+# Reserved bridge kwargs removed before forwarding to the wrapped model.
+# Only known bridge-compatibility kwargs are stripped so genuine caller
+# mistakes (e.g. `target=` vs `targets=`) still raise naturally, and no
+# signature introspection is needed to support an arbitrary forward.
+_BRIDGE_COMPAT_KWARGS = frozenset({"output_attentions"})
+
+
+class DenseOrMoEFeedForwardBridge(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.
+    """
+
+    def __init__(self, name: str, config: Any):
+        super().__init__(name, config=config, submodules={})
+        self._delegate: GeneralizedComponent | None = None
+
+    def set_original_component(self, component: torch.nn.Module) -> None:
+        super().set_original_component(component)
+        # This subclass's own __init__ takes `name: str` (non-optional), so
+        # self.name is always a str here -- but the base GeneralizedComponent
+        # attribute is typed `str | None`, which is all mypy sees without this
+        # narrowing. MoEBridge/GatedMLPBridge both require a plain `str` name.
+        assert self.name is not None
+        if hasattr(component, "router") and hasattr(component, "experts"):
+            if not isinstance(component.router, torch.nn.Module):
+                raise TypeError(
+                    f"{type(component).__name__}.router must be an nn.Module; "
+                    f"got {type(component.router).__name__}."
+                )
+            if not isinstance(component.experts, (torch.nn.ModuleList, torch.nn.ModuleDict)):
+                raise TypeError(
+                    f"{type(component).__name__}.experts must be a registered "
+                    "module collection (nn.ModuleList or nn.ModuleDict); got "
+                    f"{type(component.experts).__name__}."
+                )
+            delegate: GeneralizedComponent = MoEBridge(
+                name=self.name,
+                config=self.config,
+                submodules={"gate": LinearBridge(name="router")},
+            )
+        elif hasattr(component, "gate") and hasattr(component, "up") and hasattr(component, "down"):
+            for field in ("gate", "up", "down"):
+                value = getattr(component, field)
+                if not isinstance(value, torch.nn.Module):
+                    raise TypeError(
+                        f"{type(component).__name__}.{field} must be an "
+                        f"nn.Module; got {type(value).__name__}."
+                    )
+            delegate = GatedMLPBridge(
+                name=self.name,
+                config=self.config,
+                submodules={
+                    "gate": LinearBridge(name="gate"),
+                    "in": LinearBridge(name="up"),
+                    "out": LinearBridge(name="down"),
+                },
+            )
+        else:
+            raise ValueError(
+                f"Block.mlp is a {type(component).__name__} with neither "
+                "an MoE layer's (router, experts) nor a gated MLP's "
+                "(gate, up, down) attributes -- this adapter doesn't know "
+                "how to wrap it."
+            )
+        delegate.set_original_component(component)
+        # Not `self._delegate = delegate`: normal registration would
+        # duplicate hook_in/hook_out under a nested `._delegate.` path,
+        # risking a broad hook selector firing twice.
+        # `_delegate` is an execution helper, absent from named_modules()
+        # -- safe since parameters/state_dict/dtype are read from the raw
+        # wrapped model (see PretrainModelContainer), not this tree.
+        object.__setattr__(self, "_delegate", delegate)
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        assert self._delegate is not None, f"{self.name}: original component not set"
+        if args:
+            args = (self.hook_in(args[0]),) + args[1:]
+        elif "hidden_states" in kwargs:
+            kwargs = {**kwargs, "hidden_states": self.hook_in(kwargs["hidden_states"])}
+        output = self._delegate(*args, **kwargs)
+
+        if isinstance(output, tuple):
+            if len(output) == 0:
+                raise TypeError(
+                    "DenseOrMoEFeedForwardBridge expected a non-empty tuple "
+                    "whose first element is a torch.Tensor"
+                )
+
+            first = output[0]
+
+            if not isinstance(first, torch.Tensor):
+                raise TypeError(
+                    "DenseOrMoEFeedForwardBridge expected the first tuple element "
+                    f"to be a torch.Tensor, got {type(first).__name__}"
+                )
+
+            hooked_first = self.hook_out(first)
+
+            # Preserve every auxiliary element without sending it through HookPoint.
+            return (hooked_first, *output[1:])
+
+        if not isinstance(output, torch.Tensor):
+            raise TypeError(
+                "DenseOrMoEFeedForwardBridge expected a torch.Tensor or a tuple "
+                f"whose first element is a torch.Tensor, got {type(output).__name__}"
+            )
+
+        return self.hook_out(output)
+
+
+class _LogitsAttrDict(dict):
+    """Makes a plain dict's keys accessible as attributes (`d.logits` reads
+    `d["logits"]`), so a source model's plain-dict forward output satisfies
+    the `hasattr(output, "logits")` contract `TransformerBridge` expects.
+    Behaves as a plain dict everywhere else (indexing, `.get`, `in`, ...).
+    """
+
+    def __getattr__(self, key: str) -> Any:
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(key) from e
+
+
+class PretrainModelContainer(torch.nn.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`.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        super().__init__()
+        self.inner = model
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        filtered = {k: v for k, v in kwargs.items() if k not in _BRIDGE_COMPAT_KWARGS}
+        output = self.inner(*args, **filtered)
+
+        # Already-normalized or already-HF-style outputs pass through
+        # unchanged, but only after checking .logits is actually a tensor
+        # -- an object merely exposing the attribute isn't enough.
+        if isinstance(output, _LogitsAttrDict):
+            if "logits" not in output:
+                raise ValueError(
+                    f"{type(self.inner).__name__}.forward returned a "
+                    "_LogitsAttrDict without a 'logits' key."
+                )
+            if not isinstance(output["logits"], torch.Tensor):
+                raise TypeError(
+                    f"{type(self.inner).__name__}.forward returned a "
+                    f"_LogitsAttrDict with a non-tensor 'logits' value: "
+                    f"{type(output['logits']).__name__}."
+                )
+            return output
+
+        # try/except, not hasattr(): hasattr() would evaluate a
+        # property-backed .logits once, then a separate read would
+        # evaluate it again. This reads it exactly once.
+        try:
+            logits = output.logits
+        except AttributeError:
+            pass
+        else:
+            if not isinstance(logits, torch.Tensor):
+                raise TypeError(
+                    f"{type(self.inner).__name__}.forward returned a "
+                    f"{type(output).__name__} with a non-tensor .logits value: "
+                    f"{type(logits).__name__}."
+                )
+            return output
+
+        if isinstance(output, torch.Tensor):
+            return output
+
+        # TransformerBridge extracts output[0] as logits for tuple returns.
+        if isinstance(output, tuple):
+            if output and isinstance(output[0], torch.Tensor):
+                return output
+            raise TypeError(
+                f"{type(self.inner).__name__}.forward returned a tuple whose "
+                f"first element is a {type(output[0]).__name__ if output else 'empty tuple'}, "
+                "not a torch.Tensor -- TransformerBridge extracts output[0] as "
+                "logits for tuple returns, so it must be a tensor."
+            )
+
+        # The primary case: a plain dict, normalized into _LogitsAttrDict.
+        if isinstance(output, dict):
+            if "logits" not in output:
+                raise ValueError(
+                    f"{type(self.inner).__name__}.forward returned a dict with keys "
+                    # list(), not sorted(): sorted() raises TypeError on
+                    # heterogeneous keys, which would mask this error.
+                    f"{list(output.keys())}, but PretrainModelContainer requires a "
+                    "'logits' key -- without it, TransformerBridge's own "
+                    "hasattr(output, 'logits') check would silently fail the same "
+                    "way this container exists to prevent."
+                )
+            if not isinstance(output["logits"], torch.Tensor):
+                raise TypeError(
+                    f"{type(self.inner).__name__}.forward returned a dict whose "
+                    f"'logits' value is a {type(output['logits']).__name__}, not a "
+                    "torch.Tensor."
+                )
+            return _LogitsAttrDict(output)
+
+        raise TypeError(
+            f"{type(self.inner).__name__}.forward must return a torch.Tensor, a "
+            "tuple whose first element is a tensor, a dict containing a "
+            "tensor-valued 'logits' key, or an object with a tensor-valued "
+            f".logits attribute; got {type(output).__name__}."
+        )
+
+
+class NativeForwardAttentionBridge(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 = {}
+    property_aliases = {}
+    supports_split_qkv_fork = False
+
+
+class PretrainArchitectureAdapter(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.
+    """
+
+    def __init__(self, cfg: Any) -> None:
+        super().__init__(cfg)
+
+        self.cfg.normalization_type = "RMS"
+        self.cfg.positional_embedding_type = "rotary"
+        self.cfg.final_rms = True
+        self.cfg.gated_mlp = True
+        self.cfg.attn_only = False
+
+        self.component_mapping = {
+            # "inner." because this adapter expects the source model to
+            # arrive wrapped in `PretrainModelContainer` -- see that
+            # class's docstring for why.
+            "embed": EmbeddingBridge(name="inner.embed"),
+            "blocks": DelegatedAttentionBlockBridge(
+                name="inner.blocks",
+                config=self.cfg,
+                submodules={
+                    "ln1": RMSNormalizationBridge(name="norm1", config=self.cfg),
+                    "attn": NativeForwardAttentionBridge(
+                        name="attn",
+                        config=self.cfg,
+                        submodules={},  # opaque wrap -- see class docstring
+                    ),
+                    "ln2": RMSNormalizationBridge(name="norm2", config=self.cfg),
+                    "mlp": DenseOrMoEFeedForwardBridge(name="mlp", config=self.cfg),
+                },
+            ),
+            "ln_final": RMSNormalizationBridge(name="inner.norm_f", config=self.cfg),
+            "unembed": UnembeddingBridge(name="inner.lm_head"),
+        }
+
+
+def build_pretrain_bridge(
+    model: torch.nn.Module,
+    cfg: TransformerBridgeConfig,
+    *,
+    device: Any = None,
+    dtype: torch.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.
+    """
+    from transformer_lens.model_bridge.sources._bridge_builder import (
+        build_bridge_from_module,
+    )
+
+    kwargs: dict[str, Any] = {}
+    if device is not None:
+        kwargs["device"] = device
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    if model_name is not None:
+        kwargs["model_name"] = model_name
+
+    return build_bridge_from_module(
+        PretrainModelContainer(model),
+        architecture=ARCHITECTURE_NAME,
+        tl_config=cfg,
+        **kwargs,
+    )

From e8dd8031e88ca5690c68239eeb0a66875e3eb905 Mon Sep 17 00:00:00 2001
From: Katherine OMBRELLARO <kombrellaro@gmail.com>
Date: Fri, 17 Jul 2026 03:23:27 +0200
Subject: [PATCH 2/2] Add TransformerLens pretrain bridge adapter

---
 .../factories/architecture_adapter_factory.py | 24 +++++++
 transformer_lens/model_bridge/bridge.py       | 68 +++++++++++--------
 .../supported_architectures/__init__.py       | 52 +++++++++++++-
 3 files changed, 113 insertions(+), 31 deletions(-)

diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py
index 0e7eba9bb..4f1e099dd 100644
--- a/transformer_lens/factories/architecture_adapter_factory.py
+++ b/transformer_lens/factories/architecture_adapter_factory.py
@@ -11,15 +11,20 @@
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.supported_architectures import (
     ApertusArchitectureAdapter,
+    ArceeArchitectureAdapter,
     BaichuanArchitectureAdapter,
     BartArchitectureAdapter,
+    BD3LMArchitectureAdapter,
     BertArchitectureAdapter,
     BloomArchitectureAdapter,
     CodeGenArchitectureAdapter,
+    Cohere2ArchitectureAdapter,
     CohereArchitectureAdapter,
     DeepSeekV2ArchitectureAdapter,
     DeepSeekV3ArchitectureAdapter,
+    DeepSeekV4ArchitectureAdapter,
     FalconArchitectureAdapter,
+    FalconH1ArchitectureAdapter,
     Gemma1ArchitectureAdapter,
     Gemma2ArchitectureAdapter,
     Gemma3ArchitectureAdapter,
@@ -39,7 +44,9 @@
     HubertArchitectureAdapter,
     HunYuanDenseV1ArchitectureAdapter,
     InternLM2ArchitectureAdapter,
+    Lfm2ArchitectureAdapter,
     Lfm2MoeArchitectureAdapter,
+    LLaDAArchitectureAdapter,
     LlamaArchitectureAdapter,
     LlavaArchitectureAdapter,
     LlavaNextArchitectureAdapter,
@@ -62,37 +69,47 @@
     OlmoeArchitectureAdapter,
     OpenElmArchitectureAdapter,
     OptArchitectureAdapter,
+    OuroArchitectureAdapter,
     Phi3ArchitectureAdapter,
     PhiArchitectureAdapter,
     PhiMoEArchitectureAdapter,
     PretrainArchitectureAdapter,
     Qwen2ArchitectureAdapter,
+    Qwen2MoeArchitectureAdapter,
     Qwen3_5ArchitectureAdapter,
     Qwen3_5MultimodalArchitectureAdapter,
     Qwen3ArchitectureAdapter,
     Qwen3MoeArchitectureAdapter,
     Qwen3NextArchitectureAdapter,
     QwenArchitectureAdapter,
+    RecurrentGemmaArchitectureAdapter,
     SmolLM3ArchitectureAdapter,
     StableLmArchitectureAdapter,
     T5ArchitectureAdapter,
+    T5Gemma2ArchitectureAdapter,
     T5GemmaArchitectureAdapter,
     XGLMArchitectureAdapter,
+    Zamba2ArchitectureAdapter,
 )
 
 # Export supported architectures
 SUPPORTED_ARCHITECTURES = {
     "ApertusForCausalLM": ApertusArchitectureAdapter,
+    "ArceeForCausalLM": ArceeArchitectureAdapter,
     "BaiChuanForCausalLM": BaichuanArchitectureAdapter,
     "BaichuanForCausalLM": BaichuanArchitectureAdapter,
     "BartForConditionalGeneration": BartArchitectureAdapter,
+    "BD3LM": BD3LMArchitectureAdapter,
     "BertForMaskedLM": BertArchitectureAdapter,
     "BloomForCausalLM": BloomArchitectureAdapter,
     "CodeGenForCausalLM": CodeGenArchitectureAdapter,
+    "Cohere2ForCausalLM": Cohere2ArchitectureAdapter,
     "CohereForCausalLM": CohereArchitectureAdapter,
     "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter,
     "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter,
+    "DeepseekV4ForCausalLM": DeepSeekV4ArchitectureAdapter,
     "FalconForCausalLM": FalconArchitectureAdapter,
+    "FalconH1ForCausalLM": FalconH1ArchitectureAdapter,
     "GemmaForCausalLM": Gemma1ArchitectureAdapter,  # Default to Gemma1 as it's the original version
     "Gemma1ForCausalLM": Gemma1ArchitectureAdapter,
     "Gemma2ForCausalLM": Gemma2ArchitectureAdapter,
@@ -118,10 +135,12 @@
     "HubertModel": HubertArchitectureAdapter,
     "HunYuanDenseV1ForCausalLM": HunYuanDenseV1ArchitectureAdapter,
     "InternLM2ForCausalLM": InternLM2ArchitectureAdapter,
+    "LLaDAModelLM": LLaDAArchitectureAdapter,
     "LlamaForCausalLM": LlamaArchitectureAdapter,
     "LlavaForConditionalGeneration": LlavaArchitectureAdapter,
     "LlavaNextForConditionalGeneration": LlavaNextArchitectureAdapter,
     "LlavaOnevisionForConditionalGeneration": LlavaOnevisionArchitectureAdapter,
+    "Lfm2ForCausalLM": Lfm2ArchitectureAdapter,
     "Lfm2MoeForCausalLM": Lfm2MoeArchitectureAdapter,
     "Mamba2ForCausalLM": Mamba2ArchitectureAdapter,
     "MambaForCausalLM": MambaArchitectureAdapter,
@@ -138,22 +157,27 @@
     "OlmoeForCausalLM": OlmoeArchitectureAdapter,
     "OpenELMForCausalLM": OpenElmArchitectureAdapter,
     "OPTForCausalLM": OptArchitectureAdapter,
+    "OuroForCausalLM": OuroArchitectureAdapter,
     "PhiForCausalLM": PhiArchitectureAdapter,
     "Phi3ForCausalLM": Phi3ArchitectureAdapter,
     "PhiMoEForCausalLM": PhiMoEArchitectureAdapter,
     "QwenForCausalLM": QwenArchitectureAdapter,
     "Qwen2ForCausalLM": Qwen2ArchitectureAdapter,
+    "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter,
     "Qwen3ForCausalLM": Qwen3ArchitectureAdapter,
     "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter,
     "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter,
     "Qwen3_5ForCausalLM": Qwen3_5ArchitectureAdapter,
     "Qwen3_5ForConditionalGeneration": Qwen3_5MultimodalArchitectureAdapter,
+    "RecurrentGemmaForCausalLM": RecurrentGemmaArchitectureAdapter,
     "SmolLM3ForCausalLM": SmolLM3ArchitectureAdapter,
     "StableLmForCausalLM": StableLmArchitectureAdapter,
     "T5ForConditionalGeneration": T5ArchitectureAdapter,
     "MT5ForConditionalGeneration": T5ArchitectureAdapter,
     "T5GemmaForConditionalGeneration": T5GemmaArchitectureAdapter,
+    "T5Gemma2ForConditionalGeneration": T5Gemma2ArchitectureAdapter,
     "XGLMForCausalLM": XGLMArchitectureAdapter,
+    "Zamba2ForCausalLM": Zamba2ArchitectureAdapter,
     "NanoGPTForCausalLM": NanogptArchitectureAdapter,
     "TransformerLensNative": NativeArchitectureAdapter,
     "TransformerLensPretrain": PretrainArchitectureAdapter,
diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py
index 239eaa363..7b2239532 100644
--- a/transformer_lens/model_bridge/bridge.py
+++ b/transformer_lens/model_bridge/bridge.py
@@ -50,6 +50,7 @@
     VARIANT_SUBMODULE_NAMES,
 )
 from transformer_lens.model_bridge.get_params_util import get_bridge_params
+from transformer_lens.utilities.activation_functions import softcap_enabled
 from transformer_lens.utilities.aliases import resolve_alias
 from transformer_lens.utilities.devices import move_to_and_update_config
 from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss
@@ -219,7 +220,7 @@ def boot_transformers(
         hf_model: Optional[Any] = None,
         device_map: Optional[Union[str, Dict[str, Union[str, int]]]] = None,
         n_devices: Optional[int] = None,
-        max_memory: Optional[Dict[Union[str, int], str]] = None,
+        max_memory: Optional[Dict[Union[str, int], Union[str, int]]] = None,
         n_ctx: Optional[int] = None,
         revision: Optional[str] = None,
         checkpoint_index: Optional[int] = None,
@@ -920,7 +921,7 @@ def process_weights(
             print(f"Processing weights for {self.cfg.model_name}...")
 
         # Soft capping (tanh) is not translation-invariant; centering would change output.
-        if center_unembed and getattr(self.cfg, "output_logits_soft_cap", -1.0) > 0.0:
+        if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)):
             import logging
 
             logging.warning(
@@ -1746,6 +1747,14 @@ def forward(
             Model output based on return_type
         """
 
+        if return_type in ("loss", "both") and not self.adapter.supports_causal_loss:
+            architecture = self.cfg.architecture or type(self.adapter).__name__
+            raise NotImplementedError(
+                f"{architecture} does not support TransformerBridge's shifted causal "
+                "loss. Request return_type='logits' and compute the architecture-specific "
+                "masked-token objective explicitly."
+            )
+
         if start_at_layer is not None:
             raise NotImplementedError(
                 "start_at_layer is not supported in TransformerBridge. "
@@ -2189,7 +2198,10 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
                 stacklevel=2,
             )
         try:
-            if "output_attentions" not in filtered_kwargs:
+            if (
+                "output_attentions" not in filtered_kwargs
+                and self.adapter.supports_hf_output_attentions
+            ):
                 filtered_kwargs["output_attentions"] = True
             if processed_args:
                 output = self.forward(processed_args[0], **filtered_kwargs)
@@ -2671,6 +2683,15 @@ def _generate_tokens(
             if all_finished:
                 return
 
+    def _ensure_generation_supported(self, api_name: str) -> None:
+        """Reject autoregressive generation for forward-only architectures."""
+        if not self.adapter.supports_generation:
+            architecture = self.cfg.architecture or type(self.adapter).__name__
+            raise NotImplementedError(
+                f"TransformerBridge.{api_name}() generation is not supported by "
+                f"the {architecture} architecture."
+            )
+
     def generate(
         self,
         input: Union[str, List[str], torch.Tensor] = "",
@@ -2790,6 +2811,7 @@ def generate(
             the tokens that were fed to the model, useful for verifying BOS handling with
             chat templates.
         """
+        self._ensure_generation_supported("generate")
         # padding_side is handled internally: for batched list inputs, left-padding
         # is forced to ensure correct generation. See _is_batched_list logic below.
 
@@ -3206,6 +3228,7 @@ def generate_stream(
             max_tokens_per_yield tokens between yields. First yield includes
             the input tokens; subsequent yields contain only new tokens.
         """
+        self._ensure_generation_supported("generate_stream")
         # --- Input parsing (mirrors generate()) ---
         _is_batched_list = isinstance(input, list) and len(input) > 1
 
@@ -3437,6 +3460,7 @@ def hf_generate(
             print(result.logits)  # Logits for each generation step
             print(result.attentions)  # Attention weights
         """
+        self._ensure_generation_supported("hf_generate")
         # Handle string input by tokenizing it
         if isinstance(input, str):
             inputs = self.tokenizer(input, return_tensors="pt", padding=False, truncation=False).to(
@@ -3680,26 +3704,7 @@ def mps(self) -> "TransformerBridge":
         return self.to(torch.device("mps"))
 
     def train(self, mode: bool = True) -> "TransformerBridge":
-        """Set training mode, propagating to the wrapped source model.
-
-        ``nn.Module.train()`` only recurses into *registered* submodules.
-        ``original_model`` is intentionally stored outside the registered
-        module tree (assigned via ``self.__dict__`` rather than
-        ``nn.Module.__setattr__``); consequently, inherited ``train()``
-        does not reach it -- without this override, ``bridge.eval()`` /
-        ``bridge.train()`` silently leave dropout and other mode-dependent
-        layers in whatever mode the source model was in at wrap time.
-
-        ``eval()`` needs no override of its own: ``nn.Module.eval()`` calls
-        ``self.train(False)``, which dispatches here.
-
-        Args:
-            mode:
-                ``True`` for training mode, ``False`` for evaluation mode.
-
-        Returns:
-            The bridge itself (``nn.Module`` chaining convention).
-        """
+        """Set training mode, propagating to the wrapped source model."""
         super().train(mode)
         original = getattr(self, "original_model", None)
         if isinstance(original, torch.nn.Module):
@@ -3969,18 +3974,25 @@ def _validate_attention_fork_supported(self, flag_name: str) -> None:
                 continue
             total_with_attn += 1
             attn_classes.add(type(attn).__name__)
-            if isinstance(attn, supported_classes):
+            supports_flag = (
+                bool(getattr(attn, "supports_attn_result", False))
+                if flag_name == "use_attn_result"
+                else isinstance(attn, supported_classes)
+            )
+            if supports_flag:
                 supporting_layers.append(idx)
         if total_with_attn == 0:
             raise NotImplementedError(f"{flag_name}: no attention bridges found on self.blocks.")
         if not supporting_layers:
+            if flag_name == "use_attn_result":
+                capability_detail = "Per-head result computation is unavailable."
+            else:
+                capability_detail = "No hook point is available before the Q/K/V projections."
             raise NotImplementedError(
                 f"{flag_name}: none of this model's attention bridges support "
-                "the fine-grained Q/K/V hook fork. Found attention classes: "
+                "the requested fine-grained attention hook. Found attention classes: "
                 f"{sorted(attn_classes)}. Supported classes: "
-                f"{[c.__name__ for c in supported_classes]}. Plain "
-                "AttentionBridge delegates to HuggingFace and exposes no hook "
-                "point before the Q/K/V projection."
+                f"{[c.__name__ for c in supported_classes]}. {capability_detail}"
             )
         if len(supporting_layers) < total_with_attn:
             skipped = total_with_attn - len(supporting_layers)
diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py
index d886cadec..cc79dd117 100644
--- a/transformer_lens/model_bridge/supported_architectures/__init__.py
+++ b/transformer_lens/model_bridge/supported_architectures/__init__.py
@@ -6,12 +6,18 @@
 from transformer_lens.model_bridge.supported_architectures.apertus import (
     ApertusArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.arcee import (
+    ArceeArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.baichuan import (
     BaichuanArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.bart import (
     BartArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.bd3lm import (
+    BD3LMArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.bert import (
     BertArchitectureAdapter,
 )
@@ -22,6 +28,7 @@
     CodeGenArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.cohere import (
+    Cohere2ArchitectureAdapter,
     CohereArchitectureAdapter,
 )
 from transformer_lens.model_bridge.supported_architectures.deepseek_v2 import (
@@ -30,9 +37,15 @@
 from transformer_lens.model_bridge.supported_architectures.deepseek_v3 import (
     DeepSeekV3ArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import (
+    DeepSeekV4ArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.falcon import (
     FalconArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.falcon_h1 import (
+    FalconH1ArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.gemma1 import (
     Gemma1ArchitectureAdapter,
 )
@@ -54,6 +67,9 @@
 from transformer_lens.model_bridge.supported_architectures.glm4_moe import (
     Glm4MoeArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import (
+    GlmMoeDsaArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.gpt2 import (
     GPT2ArchitectureAdapter,
 )
@@ -69,9 +85,6 @@
 from transformer_lens.model_bridge.supported_architectures.gptj import (
     GptjArchitectureAdapter,
 )
-from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import (
-    GlmMoeDsaArchitectureAdapter,
-)
 from transformer_lens.model_bridge.supported_architectures.granite import (
     GraniteArchitectureAdapter,
 )
@@ -90,9 +103,15 @@
 from transformer_lens.model_bridge.supported_architectures.internlm2 import (
     InternLM2ArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.lfm2 import (
+    Lfm2ArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.lfm2_moe import (
     Lfm2MoeArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.llada import (
+    LLaDAArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.llama import (
     LlamaArchitectureAdapter,
 )
@@ -159,6 +178,9 @@
 from transformer_lens.model_bridge.supported_architectures.opt import (
     OptArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.ouro import (
+    OuroArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.phi import (
     PhiArchitectureAdapter,
 )
@@ -177,6 +199,9 @@
 from transformer_lens.model_bridge.supported_architectures.qwen2 import (
     Qwen2ArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.qwen2_moe import (
+    Qwen2MoeArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.qwen3 import (
     Qwen3ArchitectureAdapter,
 )
@@ -192,6 +217,9 @@
 from transformer_lens.model_bridge.supported_architectures.qwen3_next import (
     Qwen3NextArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.recurrent_gemma import (
+    RecurrentGemmaArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.smollm3 import (
     SmolLM3ArchitectureAdapter,
 )
@@ -204,21 +232,31 @@
 from transformer_lens.model_bridge.supported_architectures.t5gemma import (
     T5GemmaArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.t5gemma2 import (
+    T5Gemma2ArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.xglm import (
     XGLMArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.zamba2 import (
+    Zamba2ArchitectureAdapter,
+)
 
 __all__ = [
     "ApertusArchitectureAdapter",
+    "ArceeArchitectureAdapter",
+    "BD3LMArchitectureAdapter",
     "BaichuanArchitectureAdapter",
     "BartArchitectureAdapter",
     "BertArchitectureAdapter",
     "BloomArchitectureAdapter",
     "CodeGenArchitectureAdapter",
+    "Cohere2ArchitectureAdapter",
     "CohereArchitectureAdapter",
     "DeepSeekV2ArchitectureAdapter",
     "DeepSeekV3ArchitectureAdapter",
     "FalconArchitectureAdapter",
+    "FalconH1ArchitectureAdapter",
     "Gemma1ArchitectureAdapter",
     "Gemma2ArchitectureAdapter",
     "Gemma3ArchitectureAdapter",
@@ -238,10 +276,12 @@
     "HubertArchitectureAdapter",
     "HunYuanDenseV1ArchitectureAdapter",
     "InternLM2ArchitectureAdapter",
+    "LLaDAArchitectureAdapter",
     "LlamaArchitectureAdapter",
     "LlavaArchitectureAdapter",
     "LlavaNextArchitectureAdapter",
     "LlavaOnevisionArchitectureAdapter",
+    "Lfm2ArchitectureAdapter",
     "Lfm2MoeArchitectureAdapter",
     "MambaArchitectureAdapter",
     "Mamba2ArchitectureAdapter",
@@ -250,6 +290,7 @@
     "MixtralArchitectureAdapter",
     "MPTArchitectureAdapter",
     "NanogptArchitectureAdapter",
+    "NemotronHArchitectureAdapter",
     "NativeArchitectureAdapter",
     "NeelSoluOldArchitectureAdapter",
     "NemotronHArchitectureAdapter",
@@ -261,20 +302,25 @@
     "Olmo3ArchitectureAdapter",
     "OlmoeArchitectureAdapter",
     "OptArchitectureAdapter",
+    "OuroArchitectureAdapter",
     "PhiArchitectureAdapter",
     "Phi3ArchitectureAdapter",
     "PhiMoEArchitectureAdapter",
     "PretrainArchitectureAdapter",
     "QwenArchitectureAdapter",
     "Qwen2ArchitectureAdapter",
+    "Qwen2MoeArchitectureAdapter",
     "Qwen3ArchitectureAdapter",
     "Qwen3MoeArchitectureAdapter",
     "Qwen3NextArchitectureAdapter",
     "Qwen3_5ArchitectureAdapter",
     "Qwen3_5MultimodalArchitectureAdapter",
+    "RecurrentGemmaArchitectureAdapter",
     "SmolLM3ArchitectureAdapter",
     "StableLmArchitectureAdapter",
     "T5ArchitectureAdapter",
     "T5GemmaArchitectureAdapter",
+    "T5Gemma2ArchitectureAdapter",
     "XGLMArchitectureAdapter",
+    "Zamba2ArchitectureAdapter",
 ]