From aeed6c39c650559ebae6aefe46ae524a08ec0406 Mon Sep 17 00:00:00 2001 From: Yihang Chen Date: Mon, 2 Nov 2026 10:54:41 -0800 Subject: [PATCH] fix: keep the causal mask for broadcast 4D padding masks _apply_reconstruct_attention_mask treated every 4D mask as already encoding causal semantics. GPT-1 passes a [batch, 1, 1, seq] padding-only mask to its blocks and leaves the causal mask to the attention module, so any attention mask opened attention to later tokens. Only skip the local causal mask when the 4D mask has its own query rows. --- .../model_bridge/test_openai_gpt_adapter.py | 49 +++++++++++++++++++ .../test_mask_slicing.py | 24 +++++++++ .../generalized_components/attention.py | 11 +++-- 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/tests/integration/model_bridge/test_openai_gpt_adapter.py b/tests/integration/model_bridge/test_openai_gpt_adapter.py index 9d45543d2e..0f3d231ed6 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 0606cce0f6..ccac31c63e 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 a777a72695..016d7720cb 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