From fed6ce248c4264b37cc85337760064a02e840045 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:03:33 +0800 Subject: [PATCH 1/2] Fix masked causal loss in TransformerBridge --- .../model_bridge/test_loss_attention_mask.py | 127 ++++++++++++++++++ transformer_lens/model_bridge/bridge_core.py | 17 ++- .../model_bridge/sources/native/model.py | 3 + .../model_bridge/transformer_bridge.py | 1 + transformer_lens/utilities/lm_utils.py | 4 +- 5 files changed, 148 insertions(+), 4 deletions(-) create mode 100644 tests/unit/model_bridge/test_loss_attention_mask.py diff --git a/tests/unit/model_bridge/test_loss_attention_mask.py b/tests/unit/model_bridge/test_loss_attention_mask.py new file mode 100644 index 000000000..548bff75b --- /dev/null +++ b/tests/unit/model_bridge/test_loss_attention_mask.py @@ -0,0 +1,127 @@ +"""Regression tests for padding-aware TransformerBridge causal loss.""" + +from __future__ import annotations + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge + + +def _bridge() -> TransformerBridge: + cfg = TransformerBridgeConfig( + d_model=32, + d_head=8, + n_heads=4, + n_layers=2, + n_ctx=6, + d_vocab=32, + d_mlp=64, + act_fn="gelu", + normalization_type="LN", + seed=7, + initializer_range=0.2, + ) + return TransformerBridge.boot_native(cfg) + + +def _extract_loss(output: torch.Tensor | tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + return output[1] if isinstance(output, tuple) else output + + +@pytest.mark.parametrize("return_type", ["loss", "both"]) +def test_forward_loss_ignores_masked_padding_tokens(return_type: str) -> None: + bridge = _bridge() + attention_mask = torch.tensor( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + ] + ) + token_batches = ( + torch.tensor( + [ + [1, 2, 3, 0, 0, 0], + [4, 5, 6, 7, 8, 9], + ] + ), + torch.tensor( + [ + [1, 2, 3, 31, 30, 29], + [4, 5, 6, 7, 8, 9], + ] + ), + ) + + losses = [] + for tokens in token_batches: + output = bridge(tokens, attention_mask=attention_mask, return_type=return_type) + loss = _extract_loss(output) + logits = bridge(tokens, attention_mask=attention_mask, return_type="logits") + expected = bridge.loss_fn(logits, tokens, attention_mask=attention_mask) + + torch.testing.assert_close(loss, expected) + losses.append(loss) + + torch.testing.assert_close(losses[0], losses[1]) + + +def test_forward_loss_per_token_zeros_masked_transitions() -> None: + bridge = _bridge() + tokens = torch.tensor( + [ + [1, 2, 3, 0, 0, 0], + [4, 5, 6, 7, 8, 9], + ] + ) + attention_mask = torch.tensor( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + ] + ) + + loss = bridge( + tokens, + attention_mask=attention_mask, + return_type="loss", + loss_per_token=True, + ) + next_token_mask = torch.logical_and(attention_mask[:, :-1], attention_mask[:, 1:]) + + assert torch.count_nonzero(loss[~next_token_mask]) == 0 + + +def test_forward_loss_is_finite_with_left_padding() -> None: + bridge = _bridge() + tokens = torch.tensor( + [ + [0, 0, 0, 1, 2, 3], + [4, 5, 6, 7, 8, 9], + ] + ) + attention_mask = torch.tensor( + [ + [0, 0, 0, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + ] + ) + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + + logits = bridge( + tokens, + attention_mask=attention_mask, + position_ids=position_ids, + return_type="logits", + ) + loss = bridge( + tokens, + attention_mask=attention_mask, + position_ids=position_ids, + return_type="loss", + ) + + assert torch.isfinite(logits).all() + assert torch.isfinite(loss) diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index f361478ce..9f7da93ed 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/transformer_lens/model_bridge/bridge_core.py @@ -345,6 +345,8 @@ def loss_fn( """Cross-entropy loss matching HookedTransformer's formula (log_softmax + gather).""" if tokens.device != logits.device: tokens = tokens.to(logits.device) + if attention_mask is not None and attention_mask.device != logits.device: + attention_mask = attention_mask.to(logits.device) return lm_cross_entropy_loss(logits, tokens, attention_mask, per_token) def _finalize_return( @@ -353,6 +355,7 @@ def _finalize_return( logits: Optional[torch.Tensor], input_ids: Optional[torch.Tensor], *, + attention_mask: Optional[torch.Tensor] = None, is_audio_model: bool = False, is_visual_model: bool = False, inputs_embeds_was_used: bool = False, @@ -384,7 +387,12 @@ def _finalize_return( ) assert isinstance(logits, torch.Tensor), f"Expected logits tensor, got {type(logits)}" assert input_ids is not None, "input_ids required for return_type='loss'" - return self.loss_fn(logits, input_ids, per_token=loss_per_token) + return self.loss_fn( + logits, + input_ids, + attention_mask=attention_mask, + per_token=loss_per_token, + ) if return_type == "both": if is_audio_model: raise ValueError( @@ -404,7 +412,12 @@ def _finalize_return( ) assert isinstance(logits, torch.Tensor), f"Expected logits tensor, got {type(logits)}" assert input_ids is not None, "input_ids required for return_type='both'" - loss = self.loss_fn(logits, input_ids, per_token=loss_per_token) + loss = self.loss_fn( + logits, + input_ids, + attention_mask=attention_mask, + per_token=loss_per_token, + ) return (logits, loss) if return_type == "predictions": assert self.tokenizer is not None, "Tokenizer required for return_type='predictions'" diff --git a/transformer_lens/model_bridge/sources/native/model.py b/transformer_lens/model_bridge/sources/native/model.py index 5da5a1b6e..02a3aa9fb 100644 --- a/transformer_lens/model_bridge/sources/native/model.py +++ b/transformer_lens/model_bridge/sources/native/model.py @@ -329,6 +329,9 @@ def forward( scores = scores.masked_fill(block_mask, float("-inf")) pattern = F.softmax(scores, dim=-1) + # Fully masked padding queries softmax to NaN; overwrite masked entries + # so those rows contribute a zero attention update instead of poisoning later layers. + pattern = pattern.masked_fill(block_mask, 0.0) attn = torch.matmul(pattern, v).transpose(1, 2).contiguous().view(batch, seq, -1) out = self.o(attn) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index c0318a5b6..2c376dae2 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1658,6 +1658,7 @@ def forward( return_type, logits, input_ids, + attention_mask=attention_mask, is_audio_model=getattr(self.cfg, "is_audio_model", False), is_visual_model=getattr(self.cfg, "is_visual_model", False), inputs_embeds_was_used=_is_inputs_embeds, diff --git a/transformer_lens/utilities/lm_utils.py b/transformer_lens/utilities/lm_utils.py index a3d4f7932..db83c7ad3 100644 --- a/transformer_lens/utilities/lm_utils.py +++ b/transformer_lens/utilities/lm_utils.py @@ -17,7 +17,7 @@ def lm_cross_entropy_loss( tokens: Int[torch.Tensor, "batch pos"], attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, per_token: bool = False, -) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos"]]: +) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos_minus_one"]]: """Cross entropy loss for the language model, gives the loss for predicting the NEXT token. Args: @@ -37,7 +37,7 @@ def lm_cross_entropy_loss( # Ignore token positions which are masked out or where the next token is masked out # (generally padding tokens) next_token_mask = torch.logical_and(attention_mask[:, :-1], attention_mask[:, 1:]) - predicted_log_probs *= next_token_mask + predicted_log_probs = predicted_log_probs.masked_fill(~next_token_mask, 0.0) n_tokens = next_token_mask.sum().item() else: n_tokens = predicted_log_probs.numel() From df5f46d006a48c91e8a1057841872a26e347c4dc Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:33:29 +0800 Subject: [PATCH 2/2] Address masked loss review feedback --- .../model_bridge/test_loss_attention_mask.py | 75 ++++++++++++++++++- tests/unit/model_bridge/test_remote_bridge.py | 43 +++++++++++ tests/unit/test_lm_utils.py | 54 +++++++++++++ transformer_lens/model_bridge/bridge_core.py | 65 +++++++++++++++- .../model_bridge/remote_bridge.py | 1 + transformer_lens/utilities/lm_utils.py | 18 +++-- 6 files changed, 247 insertions(+), 9 deletions(-) create mode 100644 tests/unit/test_lm_utils.py diff --git a/tests/unit/model_bridge/test_loss_attention_mask.py b/tests/unit/model_bridge/test_loss_attention_mask.py index 548bff75b..de9151f62 100644 --- a/tests/unit/model_bridge/test_loss_attention_mask.py +++ b/tests/unit/model_bridge/test_loss_attention_mask.py @@ -4,6 +4,7 @@ import pytest import torch +import torch.nn.functional as F from transformer_lens.config import TransformerBridgeConfig from transformer_lens.model_bridge import TransformerBridge @@ -30,6 +31,16 @@ def _extract_loss(output: torch.Tensor | tuple[torch.Tensor, torch.Tensor]) -> t return output[1] if isinstance(output, tuple) else output +def _manual_masked_loss( + logits: torch.Tensor, tokens: torch.Tensor, attention_mask: torch.Tensor +) -> torch.Tensor: + transition_mask = attention_mask[:, :-1].bool() & attention_mask[:, 1:].bool() + return F.cross_entropy( + logits[:, :-1][transition_mask], + tokens[:, 1:][transition_mask], + ) + + @pytest.mark.parametrize("return_type", ["loss", "both"]) def test_forward_loss_ignores_masked_padding_tokens(return_type: str) -> None: bridge = _bridge() @@ -59,7 +70,7 @@ def test_forward_loss_ignores_masked_padding_tokens(return_type: str) -> None: output = bridge(tokens, attention_mask=attention_mask, return_type=return_type) loss = _extract_loss(output) logits = bridge(tokens, attention_mask=attention_mask, return_type="logits") - expected = bridge.loss_fn(logits, tokens, attention_mask=attention_mask) + expected = _manual_masked_loss(logits, tokens, attention_mask) torch.testing.assert_close(loss, expected) losses.append(loss) @@ -125,3 +136,65 @@ def test_forward_loss_is_finite_with_left_padding() -> None: assert torch.isfinite(logits).all() assert torch.isfinite(loss) + + +@pytest.mark.parametrize("mask_kind", ["bool", "additive"]) +@pytest.mark.parametrize("mask_layout", ["key_only", "causal"]) +def test_forward_loss_accepts_equivalent_4d_attention_mask( + mask_kind: str, mask_layout: str +) -> None: + bridge = _bridge() + tokens = torch.tensor( + [ + [1, 2, 3, 0, 0, 0], + [4, 5, 6, 7, 8, 9], + ] + ) + attention_mask = torch.tensor( + [ + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + ] + ) + blocked = ~attention_mask.bool()[:, None, None, :] + if mask_layout == "causal": + blocked = blocked | torch.ones(6, 6, dtype=torch.bool).triu(1)[None, None] + attention_mask_4d = blocked if mask_kind == "bool" else blocked.float() * -10_000.0 + + logits_2d, loss_2d = bridge( + tokens, + attention_mask=attention_mask, + return_type="both", + ) + logits_4d, loss_4d = bridge( + tokens, + attention_mask=attention_mask_4d, + return_type="both", + ) + + torch.testing.assert_close(logits_4d, logits_2d, rtol=0, atol=0) + torch.testing.assert_close(loss_4d, loss_2d) + torch.testing.assert_close(loss_4d, _manual_masked_loss(logits_4d, tokens, attention_mask)) + + +def test_loss_fn_reduces_rectangular_cached_4d_attention_mask() -> None: + bridge = _bridge() + tokens = torch.tensor([[4, 5]]) + logits = torch.zeros(1, 2, 32) + logits[0, 0, 5] = 2.0 + cache_and_new_mask = torch.tensor([[0, 1, 1, 1, 1, 1]]) + key_blocked = ~cache_and_new_mask.bool()[:, None, None, :] + query_positions = torch.tensor([4, 5]) + causal = torch.arange(6)[None, None, None, :] > query_positions[None, None, :, None] + attention_mask_4d = key_blocked | causal + + loss = bridge.loss_fn( + logits, + tokens, + attention_mask=attention_mask_4d, + per_token=True, + ) + + expected = F.cross_entropy(logits[:, 0], tokens[:, 1]) + torch.testing.assert_close(loss, expected.reshape(1, 1)) + assert loss.shape == (1, 1) diff --git a/tests/unit/model_bridge/test_remote_bridge.py b/tests/unit/model_bridge/test_remote_bridge.py index 2c41cb222..37c6c7f24 100644 --- a/tests/unit/model_bridge/test_remote_bridge.py +++ b/tests/unit/model_bridge/test_remote_bridge.py @@ -220,6 +220,49 @@ def forward( loss = bridge.forward(torch.tensor([[1, 2, 3]]), return_type="loss") assert isinstance(loss, torch.Tensor) and loss.dim() == 0 + def test_forward_loss_uses_scored_window_of_cached_attention_mask(self): + import torch + import torch.nn.functional as F + + logits = torch.zeros(1, 3, 16) + logits[0, 0, 5] = 2.0 + + class LogitsDriver(DriverBase): + supported_hook_points = frozenset({"x"}) + _supported_features = frozenset() + + def __init__(self): + super().__init__(_cfg(), tokenizer=None) + + def forward( + self, + input_ids=None, + *, + capture=(), + intervene=None, + max_new_tokens=1, + return_logits=True, + **kw, + ): + return ForwardResult(logits=logits, captured={}) + + bridge = RemoteBridge(_stub_adapter(), tokenizer=None, driver=LogitsDriver()) + tokens = torch.tensor([[4, 5, 6]]) + cache_and_new_mask = torch.tensor([[1, 1, 1, 1, 1, 0]]) + + loss = bridge.forward( + tokens, + attention_mask=cache_and_new_mask, + past_key_values=object(), + return_type="loss", + loss_per_token=True, + ) + + expected_first = F.cross_entropy(logits[:, 0], tokens[:, 1]) + expected = torch.stack((expected_first, expected_first.new_zeros(()))).unsqueeze(0) + torch.testing.assert_close(loss, expected) + assert loss.shape == (1, 2) + def test_forward_return_type_both(self): import torch diff --git a/tests/unit/test_lm_utils.py b/tests/unit/test_lm_utils.py new file mode 100644 index 000000000..ca3e6a897 --- /dev/null +++ b/tests/unit/test_lm_utils.py @@ -0,0 +1,54 @@ +"""Unit tests for language-model loss and accuracy helpers.""" + +from __future__ import annotations + +import pytest +import torch +from beartype.roar import BeartypeCallHintParamViolation + +from transformer_lens.utilities.lm_utils import lm_accuracy, lm_cross_entropy_loss + + +def test_lm_cross_entropy_loss_rejects_mismatched_attention_mask() -> None: + logits = torch.zeros(1, 2, 3) + tokens = torch.tensor([[0, 1]]) + attention_mask = torch.ones(1, 5, dtype=torch.long) + + with pytest.raises( + (AssertionError, BeartypeCallHintParamViolation), + match="attention_mask|axis 'pos'", + ): + lm_cross_entropy_loss(logits, tokens, attention_mask) + + +def test_lm_cross_entropy_loss_masks_nan_transition() -> None: + logits = torch.tensor( + [ + [ + [torch.nan, torch.nan], + [0.0, 0.0], + [0.0, 0.0], + ] + ] + ) + tokens = torch.tensor([[0, 1, 0]]) + attention_mask = torch.tensor([[0, 1, 1]]) + + per_token = lm_cross_entropy_loss(logits, tokens, attention_mask, per_token=True) + scalar = lm_cross_entropy_loss(logits, tokens, attention_mask) + expected = torch.log(torch.tensor(2.0)) + + torch.testing.assert_close(per_token, torch.stack((expected.new_zeros(()), expected))[None]) + torch.testing.assert_close(scalar, expected) + assert torch.isfinite(per_token).all() + assert torch.isfinite(scalar) + + +def test_lm_accuracy_per_token_returns_bool_pos_minus_one() -> None: + logits = torch.zeros(2, 4, 3) + tokens = torch.tensor([[0, 1, 2, 0], [2, 1, 0, 2]]) + + accuracy = lm_accuracy(logits, tokens, per_token=True) + + assert accuracy.dtype is torch.bool + assert accuracy.shape == (2, 3) diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index 9f7da93ed..2f7600396 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/transformer_lens/model_bridge/bridge_core.py @@ -345,10 +345,71 @@ def loss_fn( """Cross-entropy loss matching HookedTransformer's formula (log_softmax + gather).""" if tokens.device != logits.device: tokens = tokens.to(logits.device) - if attention_mask is not None and attention_mask.device != logits.device: - attention_mask = attention_mask.to(logits.device) + if attention_mask is not None: + if attention_mask.device != logits.device: + attention_mask = attention_mask.to(logits.device) + attention_mask = self._prepare_loss_attention_mask(attention_mask, tokens) return lm_cross_entropy_loss(logits, tokens, attention_mask, per_token) + @staticmethod + def _prepare_loss_attention_mask( + attention_mask: torch.Tensor, tokens: torch.Tensor + ) -> torch.Tensor: + """Reduce a forward attention mask to the token window scored by the loss.""" + batch, pos = tokens.shape + if attention_mask.ndim not in (2, 4): + raise ValueError( + "attention_mask must be 2D [batch, key_pos] or 4D " + f"[batch, *, query_pos, key_pos], got shape {tuple(attention_mask.shape)}" + ) + if attention_mask.shape[0] != batch: + raise ValueError( + "attention_mask batch dimension must match tokens, " + f"got {attention_mask.shape[0]} and {batch}" + ) + + if attention_mask.ndim == 2: + if attention_mask.shape[1] < pos: + raise ValueError( + "attention_mask must cover every scored token, " + f"got length {attention_mask.shape[1]} for {pos} tokens" + ) + return attention_mask[:, -pos:].bool() + + query_pos, key_pos = attention_mask.shape[-2:] + if key_pos < pos: + raise ValueError( + "attention_mask must cover every scored token, " + f"got key length {key_pos} for {pos} tokens" + ) + + blocked = attention_mask if attention_mask.dtype is torch.bool else attention_mask < -1.0 + if query_pos == 1: + # Broadcast key-only masks use one query row for the full sequence. + keep = ~blocked[..., 0, -pos:] + else: + if query_pos < pos: + raise ValueError( + "attention_mask must contain a query row for every scored token, " + f"got {query_pos} rows for {pos} tokens" + ) + # The aligned diagonal excludes causal masking while retaining padding. + diagonal = torch.diagonal( + blocked, + offset=key_pos - query_pos, + dim1=-2, + dim2=-1, + ) + if diagonal.shape[-1] < pos: + raise ValueError( + "attention_mask diagonal must cover every scored token, " + f"got length {diagonal.shape[-1]} for {pos} tokens" + ) + keep = ~diagonal[..., -pos:] + + # A token is padding only when every broadcast/head mask blocks its key. + return keep.reshape(batch, -1, pos).any(dim=1) + def _finalize_return( self, return_type: Optional[str], diff --git a/transformer_lens/model_bridge/remote_bridge.py b/transformer_lens/model_bridge/remote_bridge.py index 6518c7fe7..9f583338d 100644 --- a/transformer_lens/model_bridge/remote_bridge.py +++ b/transformer_lens/model_bridge/remote_bridge.py @@ -114,6 +114,7 @@ def forward( return_type, logits, kwargs.get("input_ids"), + attention_mask=kwargs.get("attention_mask"), is_audio_model=getattr(self.cfg, "is_audio_model", False), is_visual_model=getattr(self.cfg, "is_visual_model", False), loss_per_token=loss_per_token, diff --git a/transformer_lens/utilities/lm_utils.py b/transformer_lens/utilities/lm_utils.py index db83c7ad3..79d47c975 100644 --- a/transformer_lens/utilities/lm_utils.py +++ b/transformer_lens/utilities/lm_utils.py @@ -9,22 +9,24 @@ import torch import torch.nn.functional as F -from jaxtyping import Float, Int +from jaxtyping import Bool, Float, Int def lm_cross_entropy_loss( logits: Float[torch.Tensor, "batch pos d_vocab"], tokens: Int[torch.Tensor, "batch pos"], - attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, + attention_mask: Optional[ + Union[Bool[torch.Tensor, "batch pos"], Int[torch.Tensor, "batch pos"]] + ] = None, per_token: bool = False, -) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos_minus_one"]]: +) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos-1"]]: """Cross entropy loss for the language model, gives the loss for predicting the NEXT token. Args: logits (torch.Tensor): Logits. Shape [batch, pos, d_vocab] tokens (torch.Tensor[int64]): Input tokens. Shape [batch, pos] - attention_mask (torch.Tensor[int64], optional): Attention mask. Shape [batch, pos]. Used to - mask out padding tokens. Defaults to None. + attention_mask (torch.Tensor[int64 or bool], optional): Attention mask. Shape [batch, pos]. + Used to mask out padding tokens. Defaults to None. per_token (bool, optional): Whether to return the log probs predicted for the correct token, or the loss (ie mean of the predicted log probs). Note that the returned array has shape [batch, seq-1] as we cannot predict the first token (alternately, we ignore the final logit). Defaults to False. """ log_probs = F.log_softmax(logits, dim=-1) @@ -34,6 +36,10 @@ def lm_cross_entropy_loss( predicted_log_probs = log_probs[..., :-1, :].gather(dim=-1, index=tokens[..., 1:, None])[..., 0] if attention_mask is not None: + assert attention_mask.shape == tokens.shape, ( + "attention_mask must have the same shape as tokens, " + f"got {tuple(attention_mask.shape)} and {tuple(tokens.shape)}" + ) # Ignore token positions which are masked out or where the next token is masked out # (generally padding tokens) next_token_mask = torch.logical_and(attention_mask[:, :-1], attention_mask[:, 1:]) @@ -51,7 +57,7 @@ def lm_accuracy( logits: Float[torch.Tensor, "batch pos d_vocab"], tokens: Int[torch.Tensor, "batch pos"], per_token: bool = False, -) -> Union[Float[torch.Tensor, ""], Float[torch.Tensor, "batch pos"]]: +) -> Union[Float[torch.Tensor, ""], Bool[torch.Tensor, "batch pos-1"]]: """Cross-Entropy Accuracy for Language Modelling. We measure the accuracy on the logits for predicting the NEXT token. If per_token is True, returns the boolean for top 1 accuracy for each token in the batch. Note that this has size [batch, seq_len-1], as we cannot predict the first token.