[Bug Report] GPTNeoX bridge ignores `use_parallel_residual`, dropping `hook_resid_mid` on sequential checkpoints (all RedPajama-INCITE models)

[Chinmayrawat15]

Describe the bug

The GPTNeoX adapter hardcodes the parallel-residual wiring. neox.py sets self.cfg.parallel_attn_mlp = True unconditionally and builds ParallelBlockBridge unconditionally, ignoring HF's use_parallel_residual.

GPTNeoX ships both wirings. GPTNeoXConfig.use_parallel_residual defaults to True (Pythia), but False is a supported, real configuration, and HF branches on it:

# transformers/models/gpt_neox/modeling_gpt_neox.py — GPTNeoXLayer.forward
if self.use_parallel_residual:
    mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
    hidden_states = mlp_output + attn_output + hidden_states
else:
    attn_output = attn_output + hidden_states              # <- a real resid_mid
    mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
    hidden_states = mlp_output + attn_output

On the sequential branch there is a distinct post-attention residual, and post_attention_layernorm reads it — so ln2.hook_in is exactly resid_mid. But ParallelBlockBridge pops the hook_resid_mid alias by design, so the hook is silently unavailable on precisely the checkpoints that have one. cfg.parallel_attn_mlp also reports True for a model that is not parallel.

The failure is silent: logits are correct (the bridge delegates to HF's own block), so nothing raises. Only the hook surface and the config flag are wrong.

This affects registered models, not just hypothetical configs. Every RedPajama-INCITE checkpoint in supported_models.json is GPTNeoXForCausalLM with use_parallel_residual: false:

Model	use_parallel_residual
togethercomputer/RedPajama-INCITE-7B-Base	False
togethercomputer/RedPajama-INCITE-7B-Chat	False
togethercomputer/RedPajama-INCITE-7B-Instruct	False
togethercomputer/RedPajama-INCITE-Base-3B-v1	False
togethercomputer/RedPajama-INCITE-Chat-3B-v1	False
togethercomputer/RedPajama-INCITE-Instruct-3B-v1	False

(Read from each repo's config.json.) On all six, blocks.{i}.hook_resid_mid is missing today, so resid-mid activation patching, attribution and SAE work are silently unavailable on models the registry claims to support. Pythia is unaffected — it is genuinely parallel.

StableLmArchitectureAdapter and FalconArchitectureAdapter already guard exactly this — stablelm.py:121 and falcon.py:148 both select ParallelBlockBridge if <flag> else BlockBridge. GPTNeoX looks like an oversight rather than an intentional difference.

The same hardcode exists on the HookedTransformer side: loading_from_pretrained.py:535 sets "parallel_attn_mlp": True in the GPTNeoXForCausalLM branch.

Code example

Seeded tiny, no hub access:

import copy
import torch
from transformers import AutoModelForCausalLM, GPTNeoXConfig
from transformer_lens.model_bridge.sources import build_bridge_from_module

cfg = 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=False,          # sequential wiring
)
cfg._attn_implementation = "eager"
torch.manual_seed(42)
hf = AutoModelForCausalLM.from_config(cfg).eval()
bridge = build_bridge_from_module(
    hf, "GPTNeoXForCausalLM", hf_config=copy.deepcopy(cfg), tokenizer=None, device="cpu"
).eval()

print(type(bridge.blocks[0]).__name__)      # ParallelBlockBridge  (expected BlockBridge)
print(bridge.cfg.parallel_attn_mlp)         # True                 (expected False)

_, cache = bridge.run_with_cache(torch.randint(3, 128, (1, 10)))
print("blocks.0.hook_resid_mid" in cache)   # False                (expected True)

Logit parity against HF is exact (0.0) in both wirings, which is why this has gone unnoticed.

System Info

• Installed from source, branch dev
• macOS (Darwin arm64), Python 3.12, transformers 5.14.1, torch CPU, float32

Additional context

The root cause of the flag not reaching the adapter is that use_parallel_residual is not in _HF_PASSTHROUGH_ATTRS (model_bridge/sources/_bridge_builder.py), which is how Falcon's parallel_attn reaches its adapter. sources/transformers.py:236 does derive tl_config.parallel_attn_mlp from it, but neox.py then overwrites that value.

Reading the resolved parallel_attn_mlp instead is not a safe substitute: TransformerBridgeConfig.parallel_attn_mlp defaults to False, whereas HF's GPTNeoX default is True, so a hand-constructed config (boot_native, and tests/unit/model_bridge/supported_architectures/test_neox_adapter.py) would silently flip to sequential. The HF field name with HF's default is the correct source, mirroring Falcon.

Separately, no test asserted the identity ParallelBlockBridge documents in its own docstring (block.py:407, output = resid_pre + attn_out + mlp_out) for any of the seven adapters that use it — which is the same class of gap #1639 was. I verified the identity does hold for all seven; the missing coverage is what let this wiring bug stay invisible.

Expected behaviour & fix pointers

• neox.py selects the block class from use_parallel_residual, mirroring stablelm.py:121 / falcon.py:148, and stops overwriting cfg.parallel_attn_mlp.
• Add use_parallel_residual to _HF_PASSTHROUGH_ATTRS so the adapter can see it.
• Mirror on the HT side at loading_from_pretrained.py:535 (AGENTS.md §2).

Acceptance:

• A sequential GPTNeoX bridge builds BlockBridge, reports cfg.parallel_attn_mlp is False, and exposes hook_resid_mid
• resid_pre + attn_out == resid_mid and resid_mid + mlp_out == resid_post on the sequential variant
• Parallel (Pythia/default) behaviour unchanged, including with a hand-built TransformerBridgeConfig
• convert_hf_model_config returns parallel_attn_mlp matching use_parallel_residual
• Regression coverage for the ParallelBlockBridge identity across the seven adapters that use it
• make unit-test, make integration-test and uv run mypy . clean

Checklist

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

[Chinmayrawat15]

Opened #1649 with the fix.

It takes the approach sketched above: select the block class from use_parallel_residual (mirroring the guards already in stablelm.py:121 and falcon.py:148), add the flag to _HF_PASSTHROUGH_ATTRS so the adapter can see it, and mirror the same hardcode on the HookedTransformer side per AGENTS.md §2.

One thing worth flagging for review: I initially read the already-resolved cfg.parallel_attn_mlp instead, which broke test_neox_adapter.py::test_bridge_types. That surfaced a real trap — TransformerBridgeConfig.parallel_attn_mlp defaults to False while HF's GPTNeoX default is True, so hand-built configs (boot_native) would have silently flipped to sequential. Reading the HF field name with HF's default keeps the parallel path byte-identical.

The PR also adds coverage for the identity ParallelBlockBridge documents in its own docstring, across all seven adapters that use it — that gap is what let this stay invisible. make test-pr is clean.

Happy to change direction if you'd rather scope it differently.