TransformerBridge Compatibility Mode¶

TransformerBridge.boot_transformers(...) returns a bridge whose default numerics match HuggingFace — raw weights, no folding, no centering. Calling bridge.enable_compatibility_mode() afterwards puts the bridge into HookedTransformer-equivalent numerics — weights folded and centered. The legacy hook names already resolve on every bridge (see Model Structure); compatibility mode changes the numerics, not the names.

Most research code that was written against HookedTransformer.from_pretrained(...) assumes compatibility mode. Most new code that needs HF-faithful logits does not.

Source: transformer_lens/model_bridge/transformer_bridge.py:enable_compatibility_mode.

When to enable it¶

Use case	Compatibility mode?	Why
Logit lens / direct logit attribution	Yes	These analyses reason in the post-fold-LN coordinate system; raw HF weights produce different (wrong) attributions.
Residual-stream norm analysis	Yes	Centered weights give the residual a meaningful zero.
Circuit analysis using HT-style hook names (`blocks.{i}.attn.hook_q`, `hook_resid_pre`, etc.)	Yes	Legacy hook names resolve on every bridge (see model structure); compat mode makes the cached values match HookedTransformer’s processed-weight numerics.
Logit parity against HuggingFace	No	Folding changes weights; logits will not match HF.
Generation / inference vs HF baseline	No	Same reason.
Verifying a new adapter’s forward pass	No (initially)	Use `enable_compatibility_mode(no_processing=True)` to exercise the compatibility-mode forward without weight processing — isolates forward-pass bugs from weight-processing bugs. Hook aliases resolve on every bridge and do not need it.

What each flag does¶

bridge.enable_compatibility_mode(
    disable_warnings: bool = False,
    no_processing: bool = False,
    fold_ln: bool = True,
    center_writing_weights: bool = True,
    center_unembed: bool = True,
    fold_value_biases: bool = True,
    refactor_factored_attn_matrices: bool = False,
)

Flag	Default	Effect
`no_processing`	`False`	If `True`, overrides all other processing flags to False — leaves weights raw while still switching the bridge and every component into compatibility mode (HookedTransformer-style attention masking and hook-name bookkeeping). The “HT forward semantics on HF weights” mode; hook aliases resolve on every bridge and do not need it.
`fold_ln`	`True`	Folds LayerNorm scale + bias into the subsequent linear weights so the LayerNorm modules become pure normalization. Changes weights; mathematically equivalent.
`center_writing_weights`	`True`	Subtracts the mean from each “writing” weight (`W_out` in attention, MLP-down). Makes residual contributions sum to zero per layer, which makes residual-stream norms interpretable.
`center_unembed`	`True`	Subtracts the mean from the unembedding matrix. Logits become mean-zero — affects logit-lens output but not argmax.
`fold_value_biases`	`True`	Folds attention value biases into the output bias. Same numerics, fewer parameters.
`refactor_factored_attn_matrices`	`False`	Refactors `W_Q @ W_K.T` and `W_V @ W_O` for analysis. Off by default because it’s slow and only matters for specific factored-matrix research.
`disable_warnings`	`False`	Suppresses warnings emitted by legacy component aliases when accessed.

After processing, the bridge also:

Re-initializes the hook registry.
Re-runs _setup_hook_compatibility() (it already ran at boot) so per-attention hook conversions such as the hook_z reshape to [batch, seq, n_heads, d_head] are rebuilt on the fresh registry.
Re-registers aliases recursively so the property aliases (W_Q, b_O, …) point at the processed weights. The hook aliases already exist from boot.

compatibility_mode is then True on the bridge and on every component, so subsequent operations behave as if the bridge were loaded by HookedTransformer.from_pretrained().

Hook semantic parity¶

With their config flag set, these HT hook names fire on the pre-norm residual in either mode (matching HookedTransformer semantics):

blocks.{i}.attn.hook_q_input, hook_k_input, hook_v_input
blocks.{i}.hook_attn_in
blocks.{i}.hook_mlp_in (gated on cfg.use_hook_mlp_in; toggle via bridge.set_use_hook_mlp_in(True))

Carve-outs (issue #1317):

Post-norm architectures (OLMo 2, BERT-style) read the post-attention residual instead, because the norm semantically lives elsewhere in the block.
MLA blocks (DeepSeek V2 / V3 / R1) do not expose the split-qkv aliases — MLA’s compressed K/V doesn’t have a clean split.

On post-norm architectures (Gemma 2/3’s ln1_post/ln2_post, OLMo 2/3’s ln1/ln2), blocks.{i}.hook_attn_out / hook_mlp_out fire after the post-sublayer norm, so that they capture the tensor added to the residual stream and the identities resid_pre + attn_out == resid_mid and resid_mid + mlp_out == resid_post hold (issue #1648). One consequence for head-level direct logit attribution: per-head contributions from attn.hook_result no longer sum to hook_attn_out, because the norm sits between them. This is inherent to post-norm — decompose heads on the pre-norm side (attn.hook_out for the raw module output) or attribute through the norm explicitly. The Granite family has the same raw-vs-contribution split for a different reason: HF scales each sublayer output by residual_multiplier before the residual add, so hook_attn_out / hook_mlp_out fire on the scaled contribution while attn.hook_out / mlp.hook_out stay raw — attn.hook_result sums to the raw output, off from the contribution by the multiplier.

An adapter author for a new post-norm or MLA-style architecture must handle these carve-outs in setup_hook_compatibility. The Gemma1/Gemma2 adapters are exemplars of when not to override setup_hook_compatibility — GemmaTextScaledWordEmbedding already scales internally, so any added hook_conversion would double-scale embed.hook_out.

The four-quadrant test matrix¶

The shared conftest at tests/conftest.py provides three bridge variants plus frozen reference goldens:

Variant	`compatibility_mode`	`no_processing`	Tests…
`gpt2_bridge`	off	n/a	HF-faithful numerics
`gpt2_bridge_compat`	on	`False`	HT-equivalent numerics
`gpt2_bridge_compat_no_processing`	on	`True`	Compatibility-mode forward without weight processing — used to bisect numerical bugs
(reference) `gpt2_goldens_processed`, `gpt2_goldens_unprocessed`	n/a	n/a	Frozen HookedTransformer outputs with/without weight processing

New integration tests should use the variant that matches the property they’re testing. Tests of HF parity → gpt2_bridge. Tests of HT-API behaviour → gpt2_bridge_compat. Tests of hook semantics regardless of weights → gpt2_bridge_compat_no_processing.

Cost¶

enable_compatibility_mode() mutates the bridge’s weights in-place. It is:

One-shot: calling it twice re-runs the centering subtractions. Don’t.
Not reversible from within the bridge — re-boot for raw weights.
_setup_hook_compatibility is idempotent; only process_weights mutates weights.
Incompatible with a CPU/disk-offloaded device_map unless no_processing=True. Weight processing (fold_ln etc.) reads and rewrites parameters directly across many components at once, not through a single component’s own forward() call, so it isn’t covered by GeneralizedComponent’s per-call materialization and raises immediately rather than crashing mid-fold on a raw meta tensor. The default (non-compat-mode) forward pass works normally under offload — this restriction is compat mode’s weight processing specifically.