[Bug Report] TransformerBridge: NeoX/Pythia unembed still resolves to embed_out instead of lm_head (transformers >= 5.13)

[janmenjayap]

---

Describe the bug

TransformerBridge.boot_transformers("EleutherAI/pythia-70m") (and every other GPT-NeoX-family checkpoint) raised AttributeError during component setup on the pinned transformers version. The NeoX architecture adapter mapped the model's unembedding to the HF module named embed_out (transformer_lens/model_bridge/supported_architectures/neox.py:190), but transformers renamed GPTNeoXForCausalLM.embed_out to lm_head in the 5.13 layout. The repository pinned transformers==5.13.0 in uv.lock and required transformers>=5.9.0 in pyproject.toml, so the locked CI environment could not boot any NeoX/Pythia model through the bridge.

Suggested labels: bug, model-bridge, architecture-adapter

---

Root cause

1. transformers GPT-NeoX exposes its output projection at the top level. In the ≥ 5.13 layout that module is lm_head; older releases exposed embed_out. Verified on 5.15.1: GPTNeoXForCausalLM top-level children are ['gpt_neox', 'lm_head']; hasattr(model, "embed_out") is False, hasattr(model, "lm_head") is True.
2. The NeoX adapter hard-coded the old name: "unembed": UnembeddingBridge(name="embed_out") (neox.py:190).
3. ArchitectureAdapter.get_remote_component resolves each dotted path segment with a plain getattr and has no fallback / alias mechanism (transformer_lens/model_bridge/architecture_adapter.py:388), so a stale name raised immediately rather than degrading.

Every other component the NeoX adapter references still resolves on transformers 5.15.1, so the fix was isolated to the unembedding mapping:

Adapter mapping (neox.py)	HF module	Resolves on 5.13+?
embed → gpt_neox.embed_in (L151)	gpt_neox.embed_in	✅ still present
block MLP in/out → dense_h_to_4h / dense_4h_to_h (L179–180)	present	✅
block attn qkv/o → query_key_value / dense (L172–173)	present	✅
ln_final → gpt_neox.final_layer_norm (L186)	gpt_neox.final_layer_norm	✅
unembed → embed_out (L190)	renamed to lm_head	❌ AttributeError

The sibling GPT-Neo adapter (transformer_lens/model_bridge/supported_architectures/neo.py:160)
already used UnembeddingBridge(name="lm_head") and booted cleanly. EleutherAI/gpt-neo-125M (a dense torch.nn.Linear / out_in MLP model, the same layout family as Pythia) was booted end-to-end and used by the Backward Lens capture with exact reconstruction (abs_err = 0.0, rel_err = 0.0 across layers 0/6/11). This
confirmed the correct name is lm_head and that the NeoX-family MLP capture would work once the adapter booted. An in-repo comment already documented the rename boundary: qwen2_audio.py refers to "lm_head (the transformers >= 5.13 layout)".

---

Code example

import torch
from transformer_lens.model_bridge import TransformerBridge

TransformerBridge.boot_transformers("EleutherAI/pythia-70m", device="cpu", dtype=torch.float32)

Observed traceback (abridged, pre-fix):

File ".../model_bridge/component_setup.py", line 372, in setup_components
    original_component = architecture_adapter.get_remote_component(...)
File ".../model_bridge/architecture_adapter.py", line 388, in get_remote_component
    current = getattr(current, part)
File ".../torch/nn/modules/module.py", line 1967, in __getattr__
    raise AttributeError(...)
AttributeError: 'GPTNeoXForCausalLM' object has no attribute 'embed_out'

---

System Info

• transformer_lens installed from source (repo checkout), editable dev environment.
• transformers pinned in uv.lock: 5.13.0 (also reproduced on 5.15.1).
• transformers floor in pyproject.toml: >= 5.9.0 (no upper bound; no CI version matrix — CI runs the locked version).
• Affects the TransformerBridge path only. (The HookedTransformer legacy path is a separate code path, not affected.)
• OS: Linux.

---

Additional context

Scope of affected models: every checkpoint routed to NeoxArchitectureAdapter, including the Pythia suite (EleutherAI/pythia-70m, -160m, -410m, …) and EleutherAI/gpt-neox-20b.

Impact: no GPT-NeoX / Pythia model could be loaded through TransformerBridge.boot_transformers on the locked/supported transformers versions. This directly blocked the Backward Lens dense-MLP generalization, whose acceptance criterion requires Pythia-70m as a required out_in integration family. The Backward Lens code itself is architecture-neutral and was already verified on a real out_in model (gpt-neo-125M) — only the NeoX adapter was broken.

Design question weighed in the fix: pyproject.toml allows transformers>=5.9.0, but the lockfile and CI use 5.13.0. If releases in the [5.9, 5.13) window still expose embed_out, a hard rename to lm_head would regress them. The implementation plan weighed a straight rename (matches the lock + the GPT-Neo adapter) against a version/hasattr-aware resolution that accepts either name. Resolution: the landed fix (11912677) uses the straight rename to lm_head, matching the locked transformers==5.13.0 and the GPT-Neo adapter precedent.

Acceptance criteria (tracked in the implementation plan):

• TransformerBridge.boot_transformers("EleutherAI/pythia-70m") boots on the locked transformers version and exposes working blocks, ln_final, and unembed.
• A regression test boots a small NeoX/Pythia checkpoint through the bridge and asserts the unembedding resolves and produces logits of shape [..., d_vocab] (using a cached small model only; no new large download in CI). — covered by 9665e075.
• No other NeoX adapter behavior changes; existing bridge tests stay green.

---

Checklist

• I have checked that there is no similar issue in the repo (required)

---