Add TransformerBridge adapter for lightweight decoder-only pretraining models

[Cacapice]

Add TransformerBridge adapter for lightweight decoder-only pretraining models

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Summary

Adds an interoperability adapter for live model modules loaded from
lightweight decoder-only pretraining checkpoints (RoPE, RMSNorm, SwiGLU,
optional MoE), enabling mechanistic interpretability experiments on
models trained from scratch.

Rather than converting checkpoints into a separate TransformerLens
implementation, the adapter wraps the existing module and delegates
execution to its native forward, preserving the source model's semantics
(including RoPE conventions) while exposing the standard TransformerBridge
interface.

Motivation

Interpretability research on small from-scratch models (circuit discovery,
induction heads, toy superposition models, controlled probing) needs
precise control over architecture, data, and checkpoints. This PR adds
that interoperability path: train a small decoder-only transformer, wrap
it with build_pretrain_bridge, confirm the wrapping is transparent (the
bridge delegates to and matches the source model's own forward exactly,
rather than running a separately reimplemented TransformerLens forward),
then use TransformerLens's caching, hooks, and patching. The goal is
architecture interoperability, not a general-purpose training framework.

What changed

• PretrainArchitectureAdapter
  (transformer_lens/model_bridge/supported_architectures/pretrain.py)
  maps token embeddings, decoder blocks (RoPE attention, RMSNorm, gated
  SwiGLU MLP), final norm, and tied/untied unembedding.
• DenseOrMoEFeedForwardBridge dispatches each block's MLP
  structurally (router + experts versus gate + up + down), so
  dense, MoE, and mixed models use the same component mapping.
• The selected delegate is excluded from normal module registration to
  prevent duplicate named hook points. Structural validation requires
  projection/router fields to be modules and experts to be a registered
  ModuleList or ModuleDict, producing clear construction-time errors
  for unsupported module shapes.
• build_pretrain_bridge(model, cfg): the public entry point. Wraps
  the source model and normalizes its output contract. Direct
  build_bridge_from_module use still works for advanced cases.
• TransformerBridge.train() fix (bug fix, affects all bridges):
  original_model is intentionally stored outside the registered module
  tree (assigned via self.__dict__ rather than nn.Module.__setattr__);
  consequently, inherited nn.Module.train() -- which only recurses into
  registered submodules -- never reached the wrapped source model, so
  bridge.train()/bridge.eval() silently left mode-dependent layers
  (dropout etc.) in whatever mode the model was in at wrap time.
  TransformerBridge.train() now
  also sets mode on original_model; eval() needs no override since
  nn.Module.eval() calls self.train(False). Note this propagates to
  every bridge, including HF-backed ones -- consistent with standard
  nn.Module semantics and with the bridge's delegation of other
  lifecycle operations (to/cuda/cpu/mps, state_dict,
  parameters) to original_model -- but flagged here so the
  blast radius is reviewed consciously. Authoritative regression coverage
  lives at the bridge level in
  tests/unit/model_bridge/test_bridge_train_mode_propagation.py
  (TestTrainEvalModePropagation),
  against a minimal architecture-independent stub bridge (propagation both
  directions, train()/eval() returning self per nn.Module chaining
  convention, direct-on-model mode setting staying in sync); one
  end-to-end smoke test through build_pretrain_bridge remains in the
  container test module. The bug is reproducible on released
  transformer-lens 3.5.1 (the propagation tests fail there without the
  override), so this is a live-release fix, not one only reachable
  through the new adapter. No HookedTransformer counterpart change is
  needed: it has no wrapped original_model -- its components are
  ordinary registered submodules, which inherited nn.Module.train()
  already reaches.
• Registered "TransformerLensPretrain" in SUPPORTED_ARCHITECTURES and
  added it to the model-registry synchronization test's internal-only
  exclusion set, since no Hugging Face Hub architecture or checkpoint
  backs this key.
• Config mutation: like nanogpt.py, the adapter mutates the passed-in
  cfg in place rather than copying it.
• Unit tests split by concern: test_pretrain_adapter.py (mapping,
  dispatch, tied embeddings) and test_pretrain_model_container.py
  (container, kwarg filtering, output normalization, hooks, train/eval,
  dtype, persistence -- see Scope boundaries), sharing mocks from
  _pretrain_mocks.py. Integration tests use a
  real reference model (tests/mocks/tiny_pretrain_model.py) with genuine
  adjacent-pair RoPE/RMSNorm math for logit and hook-placement parity
  (see the integration test module docstring for the caveat on what
  logit parity does and doesn't prove).
• MoEBridge defensive validation: added validation for tuple
  outputs -- rejecting empty tuples and non-tensor first elements with a
  clear TypeError (instead of a lower-level HookPoint type-check
  error), and guarding hook_router_scores so it is only invoked for
  tensor-valued router scores. Covered directly by new MoEBridge tests
  (empty tuple, non-tensor first element, non-tensor metadata preserved
  without firing the router hook, tensor router scores still firing it),
  not just indirectly through the pretrain adapter's tests. This is a
  small compatibility and error-reporting improvement discovered while
  integrating the pretrain adapter and does not change the behavior of
  standard (hidden_states, router_scores) MoE outputs.
• NativeForwardAttentionBridge: a thin AttentionBridge subclass
  used for this adapter's opaque attention wrap. Plain AttentionBridge
  declares class-level hook_aliases (hook_q/hook_k/hook_v/
  hook_z) and property_aliases (W_Q/W_K/W_V/W_O) that assume
  mapped Q/K/V/O submodules; this adapter's wrap deliberately has none
  (see Hook coverage), so those aliases could never resolve and would
  either emit spurious "did not resolve" warnings or misrepresent the
  adapter's actual supported interface. The subclass clears both alias
  dicts and sets supports_split_qkv_fork = False so it doesn't
  advertise per-head hooks, weight aliases, or split-QKV-fork machinery
  it can't back.
• Blocks use DelegatedAttentionBlockBridge, not plain BlockBridge:
  reuses an existing abstraction rather than adding another special-case
  block, since it was written precisely for architectures where attention
  is delegated wholesale. It complements
  NativeForwardAttentionBridge.supports_split_qkv_fork = False (which
  prevents the split-QKV-fork machinery and its associated
  hook_attn_in/hook_q_input/hook_k_input/hook_v_input HookPoints
  from being exposed for this attention component) by also removing the
  block-level aliases that would otherwise dangle, pointing at hooks that
  aren't exposed. hook_attn_out is untouched by either change -- the
  attention component still fires its own hook_out normally, and the
  block-level alias to it still resolves.

Supported behavior

build_pretrain_bridge() wraps the supplied model in place. Callers that
need a raw source-model checkpoint should capture its state_dict()
before bridge construction.

Requires a specific module protocol (documented in the adapter file's
docstring) rather than any decoder-only architecture generally. Within
that protocol:

• preserves the source model's rotary convention for native-forward
  execution, since the adapter delegates rather than reimplementing it —
  the tradeoff is that this adapter does not expose per-head attention
  hooks (hook_q/hook_k/hook_v/hook_pattern);
• RMS norm + gated SwiGLU MLP, dense/MoE/mixed, dispatched structurally;
• tied or untied embeddings; forward with or without a **kwargs
  catch-all (only output_attentions is stripped — anything else,
  including a caller typo, passes through and raises naturally);
• output as a "logits" dict, tensor, tensor-first tuple, or
  .logits-exposing object, all validated, with clear errors on
  malformed shapes.

Container wrapping

PretrainModelContainer resolves attribute-name collisions with
TransformerBridge and normalizes output to the .logits contract
(raising immediately on missing/malformed output). Full mechanics are in
the class docstring. Kept local rather than generalizing
TransformerBridge's output handling, to keep this PR architecture-scoped
— a candidate for later centralization, not addressed here.

The dense/MoE delegate is excluded from the public module tree to avoid
duplicate hook points; forward execution, gradients, dtype/device
conversion, state_dict-based reconstruction, hook lifecycle, and
train/eval mode (via the TransformerBridge.train() fix above) all still
reach the wrapped source model.

Hook coverage

Block-level hooks only (resid_pre/resid_mid/resid_post, MLP
hook_in/hook_out) — no per-head hook_q/hook_k/hook_v/
hook_pattern. TransformerLens's HF-oriented attention bridges use
rotate-half RoPE, which does not match this adapter's adjacent-pair
convention and would break numerical equivalence, so it uses an opaque
NativeForwardAttentionBridge (a thin AttentionBridge subclass with no
Q/K/V/O submodules, and with hook_aliases/property_aliases cleared so
it doesn't advertise per-head hooks or weight aliases it can't back)
delegating to the source model's own attention instead. Per-head hooks
would need a new attention bridge built for that convention.

Scope boundaries

No training loop, optimizers/schedules, dataset/tokenizer infrastructure,
experiment tracking, distributed training, or checkpoint-shard
reconstruction — those stay with the source framework. This PR only adds
the TransformerBridge mapping and its tests.

Train/eval mode propagation is fixed upstream in
TransformerBridge.train() (see What changed) rather than worked around
in the adapter -- an earlier draft reassigned bridge.__class__ to a
generated subclass; that machinery is gone, and build_pretrain_bridge
now returns the bridge from build_bridge_from_module unchanged.

Persistence is tested by saving the source model's pre-bridge
state_dict, reconstructing the source model and bridge from
configuration, and verifying output and mode-propagation equivalence.
Whole-object pickling is not asserted because it depends on unrelated
TransformerBridge internals, including locally defined attention
hook-conversion classes. A standardized TransformerBridge checkpoint
interface could provide a stable, versioned save/load contract,
improving usability and scalability without requiring whole-object
serialization.

Testing

python -m pytest tests/unit/model_bridge/test_bridge_train_mode_propagation.py -v
python -m pytest tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py -v
python -m pytest tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py -v
python -m pytest tests/integration/model_bridge/test_pretrain_adapter.py -v
python -m pytest tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py -v
python -m pytest tests/unit/tools/test_model_registry.py -k TestRegistrySyncedWithFactory -v

Locally: 82 tests passed (4 bridge-level unit tests
(TransformerBridge.train() mode propagation, stub bridge), 30
adapter unit, 31 container unit, 9 integration, 4
generalized-component unit tests (MoEBridge tuple-output
validation), and 4 registry-sync tests). Covers: component
mapping and MLP dispatch (dense/MoE/mixed, checking actual
run_with_cache keys, not just isinstance, plus basic
type validation when an attribute name matches the MoE/gated-MLP protocol
but the value doesn't); logit and residual-stream parity; that
integration fixtures supply a cfg truthfully describing the model they
build (build_pretrain_bridge itself performs no cfg-vs-model
validation, so a mismatch silently reports wrong numbers on bridge.cfg
-- see TestIntegrationFixturesSupplyATruthfulConfig); activation
caching and hook interventions; kwarg filtering and output-normalization
edge cases; hook registration, gradients, dtype/device, state_dict-based
reconstruction; train/eval mode propagation (authoritative
architecture-independent coverage at the bridge level in
test_bridge_train_mode_propagation.py, plus one end-to-end smoke through
build_pretrain_bridge); malformed-architecture errors;
MoEBridge's own tuple-output validation and router-score hook gating
(empty tuple, non-tensor first element, non-tensor metadata preserved
without firing the hook, tensor router scores still firing it);
NativeForwardAttentionBridge's cleared hook_aliases/property_aliases/
supports_split_qkv_fork, and that DelegatedAttentionBlockBridge
correctly removes the now-unexposed split-QKV-fork block-level aliases
while leaving hook_attn_out (both as the block-level alias and the
attention component's own hook_out) resolvable.

Exact logit parity mainly demonstrates that wrapping does not alter the
delegated forward; it does not independently establish that the source
RoPE implementation is correct. The residual-stream decomposition and
hook-intervention tests provide the stronger evidence on hook placement.

Numerical parity is tested at float32 on CPU; float64 is tested for
preservation through construction only, not full parity at that dtype.

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue --
  TransformerBridge.train()/eval() not reaching the wrapped
  original_model)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility