diff --git a/tests/integration/model_bridge/test_openai_gpt_adapter.py b/tests/integration/model_bridge/test_openai_gpt_adapter.py index 9d45543d2..0f3d231ed 100644 --- a/tests/integration/model_bridge/test_openai_gpt_adapter.py +++ b/tests/integration/model_bridge/test_openai_gpt_adapter.py @@ -1,9 +1,13 @@ """Integration tests for the OpenAI GPT-1 architecture adapter.""" +import copy + import pytest import torch +from transformers import OpenAIGPTConfig, OpenAIGPTLMHeadModel from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources import build_bridge_from_module MODEL = "openai-community/openai-gpt" @@ -80,3 +84,48 @@ def test_greedy_generation_is_coherent(self, bridge): ) assert isinstance(out, str) assert len(out) > len("the capital of france is") + + +@pytest.fixture(scope="module") +def tiny_models(): + """A tiny random bridge and an unwrapped copy: the bridge wraps the blocks it is given.""" + torch.manual_seed(0) + config = OpenAIGPTConfig(vocab_size=64, n_embd=32, n_layer=2, n_head=4, n_positions=16) + config.architectures = ["OpenAIGPTLMHeadModel"] + hf_model = OpenAIGPTLMHeadModel(config).eval() + reference = copy.deepcopy(hf_model) + return build_bridge_from_module(hf_model, "OpenAIGPTLMHeadModel", hf_config=config), reference + + +class TestOpenAIGPTAttentionMask: + """GPT-1 hands its blocks a [batch, 1, 1, seq] padding mask, with the causal + mask left to the attention module.""" + + PAD = 63 + + def test_right_padded_batch_matches_hf(self, tiny_models): + bridge, reference = tiny_models + tokens = torch.tensor([[11, 12, 13, 14], [21, 22, self.PAD, self.PAD]]) + mask = (tokens != self.PAD).long() + with torch.no_grad(): + logits = bridge(tokens, attention_mask=mask, prepend_bos=False) + expected = reference(tokens, attention_mask=mask).logits + assert torch.allclose(logits, expected, atol=1e-5) + + def test_left_padded_row_matches_the_row_alone(self, tiny_models): + bridge, reference = tiny_models + tokens = torch.tensor([[11, 12, 13, 14], [self.PAD, self.PAD, 21, 22]]) + with torch.no_grad(): + logits = bridge(tokens, attention_mask=(tokens != self.PAD).long(), prepend_bos=False) + expected = reference(torch.tensor([[21, 22]])).logits + assert torch.allclose(logits[1, 2:], expected[0], atol=1e-5) + + def test_mask_does_not_open_attention_to_later_tokens(self, tiny_models): + bridge, _ = tiny_models + tokens = torch.tensor([[11, 12, 13, 14]]) + with torch.no_grad(): + _, cache = bridge.run_with_cache( + tokens, attention_mask=torch.ones_like(tokens), prepend_bos=False + ) + pattern = cache["blocks.0.attn.hook_pattern"] + assert torch.all(pattern.triu(diagonal=1) == 0) diff --git a/tests/unit/model_bridge/generalized_components/test_mask_slicing.py b/tests/unit/model_bridge/generalized_components/test_mask_slicing.py index 0606cce0f..ccac31c63 100644 --- a/tests/unit/model_bridge/generalized_components/test_mask_slicing.py +++ b/tests/unit/model_bridge/generalized_components/test_mask_slicing.py @@ -168,3 +168,27 @@ def test_unresolvable_per_layer_geometry_stays_visibly_flat(self) -> None: def test_module_head_dim_resolves_when_index_missing(self) -> None: shapes = self._hook_shapes(width=16, layer_idx=None, head_dim=4) assert shapes == {"q": (2, 5, 4, 4), "z": (2, 5, 4, 4)}, shapes + + +def test_broadcast_padding_mask_keeps_the_causal_mask() -> None: + """GPT-1 passes [batch, 1, 1, kv]: 4D, but padding only.""" + bridge = _attn_bridge() + kv_len, n_heads = 4, 2 + mask = torch.zeros(1, 1, 1, kv_len) + scores = torch.zeros(1, n_heads, kv_len, kv_len) + + out = bridge._apply_reconstruct_attention_mask(scores, mask, seq_len=kv_len) + future = ~torch.tril(torch.ones(kv_len, kv_len, dtype=torch.bool)) + assert torch.all(out[..., future] < 0), out + assert torch.all(out[..., ~future] == 0), out + + +def test_single_query_row_mask_on_cached_decode_is_not_over_masked() -> None: + """One query row is also what a causal 4D mask looks like on a decode step.""" + bridge = _attn_bridge() + kv_len, n_heads = 4, 2 + mask = torch.zeros(1, 1, 1, kv_len) + scores = torch.zeros(1, n_heads, 1, kv_len) + + out = bridge._apply_reconstruct_attention_mask(scores, mask, seq_len=kv_len, q_seq_len=1) + assert torch.all(out == 0), out diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index a777a7269..016d7720c 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -585,8 +585,9 @@ def _apply_reconstruct_attention_mask( """Apply causal and optional attention masking to reconstructed scores. HuggingFace-style 4D masks already encode causal semantics, so they are - treated as authoritative. Lower-rank masks do not, so the local causal - mask is still applied before adding the caller-provided padding mask. + treated as authoritative. Lower-rank masks and 4D masks broadcast over the + query rows ([batch, 1, 1, kv_seq_len], as GPT-1 passes) do not, so the local + causal mask is still applied before adding the caller-provided padding mask. Args: attn_scores: Attention scores [batch, heads, q_seq_len, kv_seq_len]. @@ -599,7 +600,11 @@ def _apply_reconstruct_attention_mask( q_seq_len = seq_len min_dtype = torch.finfo(attn_scores.dtype).min mask_value = -torch.inf if self.compatibility_mode else min_dtype - use_direct_hf_mask = attention_mask is not None and attention_mask.ndim >= 4 + use_direct_hf_mask = ( + attention_mask is not None + and attention_mask.ndim >= 4 + and attention_mask.shape[-2] != 1 + ) # Bidirectional attention (encoders) and cross-attention have no causal # structure, so only synthesize the triangular mask for causal self-attention. apply_causal = self.is_causal and not self.is_cross_attention