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..22b436ce5 --- /dev/null +++ b/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}", + ) 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/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/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..7be40e58e 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,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(