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

[jlarson4]

Thanks for tackling this @Chinmayrawat15, just a couple small change requests. Let me know if you have any questions!

[jlarson4]

TransformerBridgeConfig has no use_parallel_residual field, so on the build_bridge_from_module(..., tl_config=...) path this fallback fires and a config with parallel_attn_mlp=False still gets ParallelBlockBridge with no hook_resid_mid. Can you make the fallback getattr(cfg, "use_parallel_residual", getattr(cfg, "parallel_attn_mlp", True)) and add a unit case asserting parallel_attn_mlp=False leads to a BlockBridge being created? A plain parallel_attn_mlp read won't work due to its constructor default being False rather than HF's True, as your PR description notes.

[jlarson4]

manual_seed(42)/manual_seed(0) runs only when cache misses and leaks RNG state to later unseeded tests on the same worker. boot_native wraps its seeding in torch.random.fork_rng(devices=[]) for exactly this reason (bridge.py:360), the same two-line wrapper works here.

[Chinmayrawat15]

Thanks @jlarson4 — both were real, fixed in a4e0ba6 and pinned in e06b394.

Fallback on the tl_config= path. Applied as suggested. One consequence worth your call: with use_parallel_residual absent, the chain reduces to a plain parallel_attn_mlp read, and _make_cfg in test_neox_adapter.py never set that flag — so it defaulted to False and test_bridge_types started failing. Rather than special-case the adapter I made the fixture state parallel_attn_mlp=True explicitly, since a caller-supplied NeoX config is otherwise indistinguishable from one asking for sequential.

That does change behaviour for hand-built configs: previously any TransformerBridgeConfig got ParallelBlockBridge regardless, now an unset (i.e. False) parallel_attn_mlp yields BlockBridge. HF-booted models are unaffected. Happy to invert it if you'd rather the adapter keep defaulting to parallel when the flag was never touched.

Added both requested cases — test_sequential_config_builds_plain_block_bridge and test_hf_use_parallel_residual_overrides_config_default.

RNG leak. Good catch, wrapped in torch.random.fork_rng(devices=[]) per bridge.py:360. Verified rather than assumed: RNG is preserved across _run() on both the cache-miss and cache-hit paths, and is not preserved with the wrapper removed.

One thing that fell out of verifying this. Re-checking the guards against upstream/dev, the fallback makes the _HF_PASSTHROUGH_ATTRS entry redundant — reverting it left every test green, because sources/transformers.py:236 already derives parallel_attn_mlp and the fallback picks it up. That only holds while TransformerBridgeConfig keeps defaulting parallel_attn_mlp to False; if that default ever flips, sequential NeoX breaks silently. I kept the passthrough and added test_hf_use_parallel_residual_reaches_the_bridge_config so it is no longer resting on two defaults agreeing — but say the word if you'd prefer the smaller diff and I'll drop it.

Guard checked against upstream/dev rather than a stash this time: reverting neox.py fails 3 tests, reverting the passthrough fails 2.