From fc548244fa9fe1fe56677a04655c8442563518ca Mon Sep 17 00:00:00 2001 From: Chinmayrawat15 Date: Tue, 11 Aug 2026 14:46:26 -0700 Subject: [PATCH 1/3] fix(bridge): respect use_parallel_residual in the GPTNeoX adapter The NeoX adapter hardcoded the parallel-residual wiring, ignoring HF's use_parallel_residual. On the sequential branch HF computes a genuine post-attention residual that post_attention_layernorm reads, but ParallelBlockBridge pops the hook_resid_mid alias by design -- so the 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 because the bridge delegates to HF's own block, so nothing raised. All six RedPajama-INCITE checkpoints in the registry are GPTNeoXForCausalLM with use_parallel_residual=false, so resid-mid patching, attribution and SAE work were unavailable on registered models. Pythia is genuinely parallel and is unaffected. Select the block class from the flag, mirroring the guards already in stablelm.py:121 and falcon.py:148, and add use_parallel_residual to _HF_PASSTHROUGH_ATTRS so the adapter can see it -- reading the resolved cfg.parallel_attn_mlp instead would flip hand-built configs to sequential, since it defaults to False while HF's NeoX default is True. Mirror the same hardcode on the HookedTransformer side. Add test_parallel_residual_identities.py asserting the identity ParallelBlockBridge documents in its own docstring (resid_post == resid_pre + attn_out + mlp_out) across all seven adapters that use it, exercising the three config-switchable families in both wirings. Fixtures are seeded tinies from local HF configs, so no hub access is needed. Fixes #1644 Co-Authored-By: Claude Opus 5 (1M context) --- .../test_parallel_residual_identities.py | 179 ++++++++++++++++++ .../test_residual_decomposition_identities.py | 6 +- .../test_loading_from_pretrained_utilities.py | 25 +++ transformer_lens/loading_from_pretrained.py | 3 +- .../model_bridge/sources/_bridge_builder.py | 2 + .../supported_architectures/neox.py | 11 +- 6 files changed, 221 insertions(+), 5 deletions(-) create mode 100644 tests/integration/model_bridge/test_parallel_residual_identities.py diff --git a/tests/integration/model_bridge/test_parallel_residual_identities.py b/tests/integration/model_bridge/test_parallel_residual_identities.py new file mode 100644 index 000000000..275c7218a --- /dev/null +++ b/tests/integration/model_bridge/test_parallel_residual_identities.py @@ -0,0 +1,179 @@ +"""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" + 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("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}", + ) diff --git a/tests/integration/model_bridge/test_residual_decomposition_identities.py b/tests/integration/model_bridge/test_residual_decomposition_identities.py index 0eac7339c..9a0ec94f6 100644 --- a/tests/integration/model_bridge/test_residual_decomposition_identities.py +++ b/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 diff --git a/tests/unit/test_loading_from_pretrained_utilities.py b/tests/unit/test_loading_from_pretrained_utilities.py index 2a4f737d7..b0c948f94 100644 --- a/tests/unit/test_loading_from_pretrained_utilities.py +++ b/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 diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py index 7411ac84e..a42ddd5bb 100644 --- a/transformer_lens/loading_from_pretrained.py +++ b/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", diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 0e5674860..a1b5399e4 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/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", diff --git a/transformer_lens/model_bridge/supported_architectures/neox.py b/transformer_lens/model_bridge/supported_architectures/neox.py index 0ca3cd6bb..f162c1332 100644 --- a/transformer_lens/model_bridge/supported_architectures/neox.py +++ b/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,13 @@ 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. + use_parallel_residual = getattr(cfg, "use_parallel_residual", 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 +145,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( From a4e0ba6d147b52eeed4b60987c00398c912fb013 Mon Sep 17 00:00:00 2001 From: Chinmayrawat15 Date: Wed, 12 Aug 2026 08:08:34 -0700 Subject: [PATCH 2/3] fix(bridge): honour parallel_attn_mlp on caller-supplied NeoX configs Review feedback from @jlarson4 on #1649. TransformerBridgeConfig has no use_parallel_residual field, so on the build_bridge_from_module(..., tl_config=...) path the getattr default fired and a config with parallel_attn_mlp=False still got ParallelBlockBridge with no hook_resid_mid. Fall back to parallel_attn_mlp before defaulting to HF's True. The existing _make_cfg fixture relied on the old hardcode, so it now states parallel_attn_mlp=True explicitly -- a caller-supplied NeoX config is otherwise indistinguishable from one that asked for sequential, since the dataclass default is False. Wrap the tiny-model seeding in test_parallel_residual_identities.py in torch.random.fork_rng(devices=[]), matching bridge.py:360. Seeding ran only on a cache miss, so it leaked RNG state into later unseeded tests depending on ordering. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_parallel_residual_identities.py | 23 +++++++------ .../test_neox_adapter.py | 33 +++++++++++++++++++ .../supported_architectures/neox.py | 7 +++- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/tests/integration/model_bridge/test_parallel_residual_identities.py b/tests/integration/model_bridge/test_parallel_residual_identities.py index 275c7218a..eada93754 100644 --- a/tests/integration/model_bridge/test_parallel_residual_identities.py +++ b/tests/integration/model_bridge/test_parallel_residual_identities.py @@ -101,16 +101,19 @@ def _run(hf_architecture: str, config_cls: Any, config_kwargs: dict[str, Any]): if key not in _CACHE: config = config_cls(**config_kwargs) config._attn_implementation = "eager" - 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) + # 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] diff --git a/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py b/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py index 75f999f35..47f028ca2 100644 --- a/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py +++ b/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" diff --git a/transformer_lens/model_bridge/supported_architectures/neox.py b/transformer_lens/model_bridge/supported_architectures/neox.py index f162c1332..7be40e58e 100644 --- a/transformer_lens/model_bridge/supported_architectures/neox.py +++ b/transformer_lens/model_bridge/supported_architectures/neox.py @@ -49,7 +49,12 @@ def __init__(self, cfg: Any) -> None: # 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. - use_parallel_residual = getattr(cfg, "use_parallel_residual", True) + # 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 From e06b3949c3277d59c31bbf3364418049b8e329aa Mon Sep 17 00:00:00 2001 From: Chinmayrawat15 Date: Wed, 12 Aug 2026 08:10:25 -0700 Subject: [PATCH 3/3] test(bridge): pin use_parallel_residual passthrough The fallback added for caller-supplied configs made the _HF_PASSTHROUGH_ATTRS entry redundant -- reverting it left every test green, since sources/transformers.py already derives parallel_attn_mlp and the fallback picks it up. That only holds while TransformerBridgeConfig keeps defaulting parallel_attn_mlp to False, so assert the adapter sees HF's own flag rather than resting on two defaults agreeing. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_parallel_residual_identities.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/integration/model_bridge/test_parallel_residual_identities.py b/tests/integration/model_bridge/test_parallel_residual_identities.py index eada93754..22b436ce5 100644 --- a/tests/integration/model_bridge/test_parallel_residual_identities.py +++ b/tests/integration/model_bridge/test_parallel_residual_identities.py @@ -162,6 +162,21 @@ def test_sequential_variant_exposes_resid_mid( 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]