fix(bridge): respect use_parallel_residual in the GPTNeoX adapter

tests/integration/model_bridge/test_parallel_residual_identities.py

@@ -0,0 +1,197 @@
+"""Residual-stream decomposition identity for parallel-residual architectures.
+
+Companion to ``test_residual_decomposition_identities.py``, which covers the
+sequential wiring and scopes these architectures out. A parallel block has no
+``hook_resid_mid`` — attention and MLP both read ``resid_pre`` and both write
+into ``resid_post``, so the identity ``ParallelBlockBridge`` documents is:
+
+    resid_post == resid_pre + attn_out + mlp_out
+
+#1639 showed how quietly an alias can end up pointing at a residual-added state
+instead of an additive contribution. That identity held for every
+``ParallelBlockBridge`` adapter but nothing asserted it.
+
+GPTNeoX, StableLM and Falcon pick their block class from the config
+(``use_parallel_residual`` / ``parallel_attn``), so each is exercised in both
+wirings — the sequential cases are the regression guard for a NeoX adapter that
+hardcoded the parallel block and dropped ``hook_resid_mid`` on checkpoints that
+genuinely have one.
+
+Fixtures are seeded tinies built from local HF configs, so no hub access is
+needed and both wirings are reachable from one checkpoint shape. Bridge-vs-HF
+logit parity is deliberately not asserted here — that contract belongs to the
+per-adapter tests (e.g. ``test_cohere_adapter.py``).
+"""
+
+import copy
+from typing import Any
+
+import pytest
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    CodeGenConfig,
+    CohereConfig,
+    FalconConfig,
+    GPTJConfig,
+    GPTNeoXConfig,
+    PhiConfig,
+    StableLmConfig,
+)
+
+from transformer_lens.model_bridge.sources import build_bridge_from_module
+
+_SHARED = dict(vocab_size=128, pad_token_id=0, bos_token_id=1, eos_token_id=2)
+# Multi-head (no GQA); GPTNeoX splits a fused QKV and needs n_kv == n_heads.
+_MHA = dict(
+    hidden_size=64,
+    intermediate_size=128,
+    num_hidden_layers=2,
+    num_attention_heads=4,
+    max_position_embeddings=128,
+    **_SHARED,
+)
+_GQA = dict(_MHA, num_key_value_heads=2)
+_GPTJ = dict(n_embd=64, n_layer=2, n_head=4, n_positions=128, rotary_dim=16, **_SHARED)
+_FALCON = dict(
+    hidden_size=64,
+    num_hidden_layers=2,
+    num_attention_heads=4,
+    max_position_embeddings=128,
+    new_decoder_architecture=False,
+    **_SHARED,
+)
+
+PARALLEL_CASES = [
+    pytest.param("GPTJForCausalLM", GPTJConfig, _GPTJ, id="gptj"),
+    pytest.param("CodeGenForCausalLM", CodeGenConfig, _GPTJ, id="codegen"),
+    pytest.param("PhiForCausalLM", PhiConfig, _GQA, id="phi"),
+    pytest.param("CohereForCausalLM", CohereConfig, _GQA, id="cohere"),
+    pytest.param(
+        "GPTNeoXForCausalLM", GPTNeoXConfig, dict(_MHA, use_parallel_residual=True), id="neox"
+    ),
+    pytest.param(
+        "StableLmForCausalLM", StableLmConfig, dict(_GQA, use_parallel_residual=True), id="stablelm"
+    ),
+    pytest.param("FalconForCausalLM", FalconConfig, dict(_FALCON, parallel_attn=True), id="falcon"),
+]
+
+SEQUENTIAL_CASES = [
+    pytest.param(
+        "GPTNeoXForCausalLM", GPTNeoXConfig, dict(_MHA, use_parallel_residual=False), id="neox"
+    ),
+    pytest.param(
+        "StableLmForCausalLM",
+        StableLmConfig,
+        dict(_GQA, use_parallel_residual=False),
+        id="stablelm",
+    ),
+    pytest.param(
+        "FalconForCausalLM", FalconConfig, dict(_FALCON, parallel_attn=False), id="falcon"
+    ),
+]
+
+# Built once per (architecture, wiring); the tests only read from them.
+_CACHE: dict[tuple[str, str], tuple[Any, dict[str, torch.Tensor]]] = {}
+
+
+def _run(hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]):
+    """Boot a seeded tiny bridge for this architecture and cache one forward."""
+    key = (hf_architecture, repr(sorted(config_kwargs.items())))
+    if key not in _CACHE:
+        config = config_cls(**config_kwargs)
+        config._attn_implementation = "eager"
+        # Fork the RNG so seeding here stays deterministic on a cache miss
+        # without leaking state into later unseeded tests on this worker.
+        with torch.random.fork_rng(devices=[]):
+            torch.manual_seed(42)
+            hf = AutoModelForCausalLM.from_config(config).eval()
+            bridge = build_bridge_from_module(
+                hf, hf_architecture, hf_config=copy.deepcopy(config), tokenizer=None, device="cpu"
+            ).eval()
+
+            torch.manual_seed(0)
+            tokens = torch.randint(3, config_kwargs["vocab_size"], (1, 10))
+            with torch.no_grad():
+                logits, cache = bridge.run_with_cache(tokens)
+        assert torch.isfinite(logits).all(), f"{hf_architecture} produced non-finite logits"
+        _CACHE[key] = (bridge, cache)
+    return _CACHE[key]
+
+
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", PARALLEL_CASES)
+def test_parallel_block_omits_resid_mid(
+    hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]
+) -> None:
+    """A parallel block has no post-attention residual, so the alias must be absent."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+
+    assert type(bridge.blocks[0]).__name__ == "ParallelBlockBridge"
+    assert bridge.cfg.parallel_attn_mlp is True
+    for layer in range(bridge.cfg.n_layers):
+        assert f"blocks.{layer}.hook_resid_mid" not in cache
+
+
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", PARALLEL_CASES)
+def test_parallel_residual_decomposition(
+    hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]
+) -> None:
+    """resid_post == resid_pre + attn_out + mlp_out on every layer."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+
+    for layer in range(bridge.cfg.n_layers):
+        torch.testing.assert_close(
+            cache[f"blocks.{layer}.hook_resid_post"],
+            cache[f"blocks.{layer}.hook_resid_pre"]
+            + cache[f"blocks.{layer}.hook_attn_out"]
+            + cache[f"blocks.{layer}.hook_mlp_out"],
+            msg=lambda got: f"{hf_architecture} layer {layer} parallel identity broken:\n{got}",
+        )
+
+
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", SEQUENTIAL_CASES)
+def test_sequential_variant_exposes_resid_mid(
+    hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]
+) -> None:
+    """The sequential variant of a switchable family keeps its post-attention residual."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+
+    assert type(bridge.blocks[0]).__name__ == "BlockBridge"
+    assert bridge.cfg.parallel_attn_mlp is False
+    for layer in range(bridge.cfg.n_layers):
+        assert f"blocks.{layer}.hook_resid_mid" in cache
+
+
+@pytest.mark.parametrize("use_parallel_residual", [True, False])
+def test_hf_use_parallel_residual_reaches_the_bridge_config(use_parallel_residual: bool) -> None:
+    """The adapter reads HF's own flag, so the passthrough must carry it.
+
+    Without it the adapter falls back to the derived ``parallel_attn_mlp``,
+    which is only correct while ``TransformerBridgeConfig`` keeps defaulting it
+    to ``False`` — the wiring should not rest on two defaults agreeing.
+    """
+    bridge, _ = _run(
+        "GPTNeoXForCausalLM", GPTNeoXConfig, dict(_MHA, use_parallel_residual=use_parallel_residual)
+    )
+
+    assert bridge.cfg.use_parallel_residual is use_parallel_residual
+
+
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", SEQUENTIAL_CASES)
+def test_sequential_variant_decomposition(
+    hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]
+) -> None:
+    """Both HookedTransformer identities hold once resid_mid is exposed."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+
+    for layer in range(bridge.cfg.n_layers):
+        torch.testing.assert_close(
+            cache[f"blocks.{layer}.hook_resid_mid"],
+            cache[f"blocks.{layer}.hook_resid_pre"] + cache[f"blocks.{layer}.hook_attn_out"],
+            msg=lambda got: f"{hf_architecture} layer {layer} attn identity broken:\n{got}",
+        )
+        torch.testing.assert_close(
+            cache[f"blocks.{layer}.hook_resid_post"],
+            cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"],
+            msg=lambda got: f"{hf_architecture} layer {layer} mlp identity broken:\n{got}",
+        )

tests/integration/model_bridge/test_residual_decomposition_identities.py

@@ -14,8 +14,10 @@
 - mistral: Llama-style pre-RMSNorm block-level adds
 - bloom: residual added *inside* the HF attention/MLP modules (the #1639 case)
 
-Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope:
-they have no ``hook_resid_mid``.
+Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope
+here: they have no ``hook_resid_mid``. Their two-term identity
+(``resid_post == resid_pre + attn_out + mlp_out``) is covered by
+``test_parallel_residual_identities.py``.
 """
 
 import pytest

tests/unit/model_bridge/supported_architectures/test_neox_adapter.py

@@ -14,6 +14,7 @@
 
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.generalized_components import (
+    BlockBridge,
     EmbeddingBridge,
     JointQKVPositionEmbeddingsAttentionBridge,
     LinearBridge,
@@ -51,6 +52,9 @@ def _make_cfg(
         d_mlp=d_mlp,
         default_prepend_bos=False,
         architecture="GPTNeoXForCausalLM",
+        # Explicit: the adapter now honours this, and TransformerBridgeConfig
+        # defaults it to False where HF's GPTNeoX default is True.
+        parallel_attn_mlp=True,
     )
 
 
@@ -113,6 +117,35 @@ def test_bridge_types(self, adapter: NeoxArchitectureAdapter) -> None:
         assert isinstance(mapping["ln_final"], NormalizationBridge)
         assert isinstance(mapping["unembed"], UnembeddingBridge)
 
+    def test_sequential_config_builds_plain_block_bridge(self) -> None:
+        """A caller-supplied config with parallel_attn_mlp=False must not get the parallel block.
+
+        Only HF-booted configs carry ``use_parallel_residual``; the
+        ``build_bridge_from_module(..., tl_config=...)`` path has to fall back
+        to ``parallel_attn_mlp`` or the sequential wiring loses hook_resid_mid.
+        """
+        cfg = _make_cfg()
+        cfg.parallel_attn_mlp = False
+
+        adapter = NeoxArchitectureAdapter(cfg)
+        blocks = adapter.component_mapping["blocks"]
+
+        assert isinstance(blocks, BlockBridge)
+        assert not isinstance(blocks, ParallelBlockBridge)
+        assert adapter.cfg.parallel_attn_mlp is False
+        assert "hook_resid_mid" in blocks.hook_aliases
+
+    def test_hf_use_parallel_residual_overrides_config_default(self) -> None:
+        """HF-booted configs carry use_parallel_residual; it wins over the TL default."""
+        cfg = _make_cfg()
+        cfg.parallel_attn_mlp = False
+        cfg.use_parallel_residual = True
+
+        adapter = NeoxArchitectureAdapter(cfg)
+
+        assert isinstance(adapter.component_mapping["blocks"], ParallelBlockBridge)
+        assert adapter.cfg.parallel_attn_mlp is True
+
     def test_top_level_hf_paths(self, adapter: NeoxArchitectureAdapter) -> None:
         mapping = adapter.component_mapping
         assert mapping["embed"].name == "gpt_neox.embed_in"

tests/unit/test_loading_from_pretrained_utilities.py

@@ -278,3 +278,28 @@ def test_apertus_instruct_config(self):
             raise
         assert cfg.original_architecture == "ApertusForCausalLM"
         assert cfg.act_fn == "xielu"
+
+    @pytest.mark.parametrize("use_parallel_residual", [True, False])
+    def test_gpt_neox_parallel_residual_follows_hf_config(
+        self, tmp_path, use_parallel_residual: bool
+    ):
+        """GPTNeoX ships sequential variants; parallel_attn_mlp must not be hardcoded."""
+        from transformers import GPTNeoXConfig
+
+        from transformer_lens.loading_from_pretrained import convert_hf_model_config
+
+        hf_config = GPTNeoXConfig(
+            vocab_size=128,
+            hidden_size=64,
+            intermediate_size=128,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            max_position_embeddings=128,
+            use_parallel_residual=use_parallel_residual,
+        )
+        hf_config.architectures = ["GPTNeoXForCausalLM"]
+        hf_config.save_pretrained(tmp_path)
+
+        cfg_dict = convert_hf_model_config(str(tmp_path))
+
+        assert cfg_dict["parallel_attn_mlp"] is use_parallel_residual

transformer_lens/loading_from_pretrained.py

@@ -532,7 +532,8 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]:
             "use_attn_scale": True,
             "use_local_attn": False,
             "scale_attn_by_inverse_layer_idx": False,
-            "parallel_attn_mlp": True,
+            # GPTNeoX ships sequential variants too (use_parallel_residual=False).
+            "parallel_attn_mlp": getattr(hf_config, "use_parallel_residual", True),
             "positional_embedding_type": "rotary",
             "rotary_adjacent_pairs": False,
             "normalization_type": "LN",

transformer_lens/model_bridge/sources/_bridge_builder.py

@@ -40,6 +40,8 @@
     "new_decoder_architecture",
     "alibi",
     "num_ln_in_parallel_attn",
+    # GPTNeoX
+    "use_parallel_residual",
     # Mamba (SSM config)
     "state_size",
     "conv_kernel",

transformer_lens/model_bridge/supported_architectures/neox.py

@@ -16,6 +16,7 @@
 )
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
+    BlockBridge,
     EmbeddingBridge,
     JointQKVPositionEmbeddingsAttentionBridge,
     LinearBridge,
@@ -44,7 +45,18 @@ def __init__(self, cfg: Any) -> None:
         self.cfg.final_rms = False
         self.cfg.gated_mlp = False
         self.cfg.attn_only = False
-        self.cfg.parallel_attn_mlp = True
+
+        # GPTNeoX ships both parallel (Pythia, HF's default) and sequential
+        # variants. Hardcoding parallel drops hook_resid_mid on sequential
+        # checkpoints that genuinely have a post-attention residual.
+        # HF-booted configs carry use_parallel_residual; a caller-supplied
+        # TransformerBridgeConfig only has parallel_attn_mlp, so fall back to
+        # it before defaulting to HF's True.
+        use_parallel_residual = getattr(
+            cfg, "use_parallel_residual", getattr(cfg, "parallel_attn_mlp", True)
+        )
+        self.cfg.parallel_attn_mlp = use_parallel_residual
+        block_cls = ParallelBlockBridge if use_parallel_residual else BlockBridge
 
         # NeoX/Pythia models were not trained with BOS tokens
         self.cfg.default_prepend_bos = False
@@ -138,7 +150,7 @@ def __init__(self, cfg: Any) -> None:
         self.component_mapping = {
             "embed": EmbeddingBridge(name="gpt_neox.embed_in"),
             "rotary_emb": RotaryEmbeddingBridge(name="gpt_neox.rotary_emb"),
-            "blocks": ParallelBlockBridge(
+            "blocks": block_cls(
                 name="gpt_neox.layers",
                 submodules={
                     "ln1": NormalizationBridge(