fix(bridge): respect use_parallel_residual in the GPTNeoX adapter

[Chinmayrawat15]

Description

Fixes #1644.

The GPTNeoX adapter hardcoded the parallel-residual wiring — neox.py set cfg.parallel_attn_mlp = True and built ParallelBlockBridge unconditionally, ignoring HF's use_parallel_residual. GPTNeoX ships both wirings, and on the sequential branch HF computes a genuine post-attention residual that post_attention_layernorm reads:

# transformers/models/gpt_neox/modeling_gpt_neox.py — GPTNeoXLayer.forward
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

ParallelBlockBridge pops the hook_resid_mid alias by design, so that hook was silently missing on exactly the checkpoints that have one, and cfg.parallel_attn_mlp reported True for a model that is not parallel. Logits stayed correct (the bridge delegates to HF's own block), so nothing raised — which is why it went unnoticed.

This is not hypothetical. All six RedPajama-INCITE checkpoints in supported_models.json are GPTNeoXForCausalLM with use_parallel_residual: false, so resid-mid patching, attribution and SAE work were silently unavailable on registered models:

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

Pythia is genuinely parallel and is unaffected.

The fix

StableLmArchitectureAdapter and FalconArchitectureAdapter already guard this exact case, so the change mirrors them rather than inventing a pattern:

• neox.py — block_cls = ParallelBlockBridge if use_parallel_residual else BlockBridge, matching stablelm.py:121 / falcon.py:148, and cfg.parallel_attn_mlp now follows the flag instead of being overwritten.
• _bridge_builder.py — added use_parallel_residual to _HF_PASSTHROUGH_ATTRS, which is how Falcon's parallel_attn already reaches its adapter.
• loading_from_pretrained.py — same hardcode existed on the HookedTransformer side; mirrored per AGENTS.md §2.

Reading the already-resolved cfg.parallel_attn_mlp instead looked simpler but is wrong: TransformerBridgeConfig.parallel_attn_mlp defaults to False while HF's GPTNeoX default is True, so a hand-built config (boot_native, and the existing test_neox_adapter.py) would silently flip to sequential. Using the HF field name with HF's default keeps the parallel path byte-identical — I confirmed the existing NeoX adapter unit tests pass unchanged.

Test coverage

ParallelBlockBridge documents its identity in its own docstring (block.py:407, output = resid_pre + attn_out + mlp_out), but nothing asserted it for any of the seven adapters that use it — the same class of gap as #1639, and what let this wiring bug stay invisible.

New tests/integration/model_bridge/test_parallel_residual_identities.py covers all seven (gptj, codegen, phi, cohere, neox, stablelm, falcon) and exercises the three config-switchable families in both wirings. Fixtures are seeded tinies built from local HF configs, so there is no hub access and both wirings are reachable from one checkpoint shape — 20 tests in ~8s.

It is a real regression guard, not a rubber stamp: reverting the neox.py + _bridge_builder.py change fails exactly test_sequential_variant_exposes_resid_mid[neox] and test_sequential_variant_decomposition[neox], with the other 18 still passing. Reverting the loading_from_pretrained.py change fails the new convert_hf_model_config case.

The new file is a sibling to test_residual_decomposition_identities.py (which explicitly scoped these architectures out); I updated that docstring to cross-reference rather than moving anyone's tests.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Validation

make test-pr clean on macOS arm64 / Python 3.12 / transformers 5.14.1:

Tier	Result
make unit-test	4550 passed, 36 skipped, 44 deselected, 10 xfailed
make docstring-test	18 passed, 25 skipped
make acceptance-test	150 passed, 96 skipped, 11 deselected
make integration-test	1182 passed, 23 skipped, 191 deselected, 1 xfailed

• uv run pytest tests/integration/model_bridge/test_parallel_residual_identities.py — 20 passed in 8s
• make check-format clean; uv run mypy . — Success, no issues in 385 source files

Noted but not included

transformer_lens/benchmarks/hook_registration.py still carries a BLOOM carve-out (lines ~246-255 and ~592-599) whose comment — "hook_attn_out and hook_mlp_out capture attn+residual instead of just attn" — was made false by #1640/#1642. It is unreachable today (both call sites in main_benchmark.py:401-402 pass no reference_activations, so the functions return before it), so it is dead code with a misleading comment rather than a live suppression. Left out to keep this PR to one story — happy to send it separately.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (behaviour is covered by tests; no doc pages describe this wiring)
• 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