diff --git a/tests/integration/model_bridge/test_left_padding_positions.py b/tests/integration/model_bridge/test_left_padding_positions.py new file mode 100644 index 000000000..9c0f50b17 --- /dev/null +++ b/tests/integration/model_bridge/test_left_padding_positions.py @@ -0,0 +1,168 @@ +"""Regression tests for left-padding position handling in TransformerBridge. + +A causal LM's logits at a sequence's real token positions must not depend on how +that sequence is padded, provided the caller supplies the matching attention_mask. +Left padding shifts every real token's absolute position, so position_ids have to +be derived from the mask; without that the bridge silently returns wrong logits +and a wrong loss. See #1609. + +Right padding is included as a control: causality already protects it, so it was +never affected and must stay that way. +""" + +from __future__ import annotations + +import pytest +import torch + +from transformer_lens import utilities as utils + +PAD_ID = 0 + + +def _pad(tokens: torch.Tensor, n_pad: int, side: str) -> tuple[torch.Tensor, torch.Tensor]: + """Pad `tokens` on `side`, returning (padded_tokens, attention_mask).""" + pads = torch.full((tokens.shape[0], n_pad), PAD_ID, dtype=tokens.dtype) + ones = torch.ones(tokens.shape, dtype=torch.long) + zeros = torch.zeros((tokens.shape[0], n_pad), dtype=torch.long) + if side == "left": + return torch.cat([pads, tokens], dim=1), torch.cat([zeros, ones], dim=1) + return torch.cat([tokens, pads], dim=1), torch.cat([ones, zeros], dim=1) + + +@pytest.fixture(scope="module") +def tokens(distilgpt2_bridge) -> torch.Tensor: + return distilgpt2_bridge.to_tokens("The capital of France is") + + +@pytest.mark.parametrize("side", ["left", "right"]) +@pytest.mark.parametrize("n_pad", [1, 3, 5]) +def test_logits_are_invariant_to_padding(distilgpt2_bridge, tokens, side, n_pad) -> None: + """Padding must not change the logits at a sequence's real positions.""" + padded, mask = _pad(tokens, n_pad, side) + real = slice(n_pad, None) if side == "left" else slice(None, tokens.shape[1]) + + with torch.no_grad(): + baseline = distilgpt2_bridge(tokens, return_type="logits") + actual = distilgpt2_bridge(padded, attention_mask=mask, return_type="logits")[:, real] + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, baseline, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("side", ["left", "right"]) +def test_compat_mode_logits_are_invariant_to_padding(distilgpt2_bridge_compat, side: str) -> None: + """enable_compatibility_mode() promises HookedTransformer-equivalent numerics, + which this property is part of.""" + tokens = distilgpt2_bridge_compat.to_tokens("The capital of France is") + n_pad = 3 + padded, mask = _pad(tokens, n_pad, side) + real = slice(n_pad, None) if side == "left" else slice(None, tokens.shape[1]) + + with torch.no_grad(): + baseline = distilgpt2_bridge_compat(tokens, return_type="logits") + actual = distilgpt2_bridge_compat(padded, attention_mask=mask, return_type="logits")[ + :, real + ] + + torch.testing.assert_close(actual, baseline, rtol=1e-3, atol=1e-3) + + +def test_derived_position_ids_match_hooked_transformer(distilgpt2_bridge, tokens) -> None: + """The derived positions must be the ones HookedTransformer would use, i.e. the + shared get_offset_position_ids helper rather than a parallel derivation.""" + n_pad = 3 + padded, mask = _pad(tokens, n_pad, "left") + expected = utils.get_offset_position_ids(0, mask) + + with torch.no_grad(): + derived = distilgpt2_bridge(padded, attention_mask=mask, return_type="logits") + supplied = distilgpt2_bridge( + padded, attention_mask=mask, position_ids=expected, return_type="logits" + ) + + torch.testing.assert_close(derived, supplied, rtol=1e-5, atol=1e-5) + + +def test_explicit_position_ids_take_precedence(distilgpt2_bridge, tokens) -> None: + """A caller-supplied position_ids must not be overwritten by the derivation.""" + n_pad = 3 + padded, mask = _pad(tokens, n_pad, "left") + derived_positions = utils.get_offset_position_ids(0, mask) + shifted = derived_positions + 1 # deliberately different, but still in range + + with torch.no_grad(): + default = distilgpt2_bridge(padded, attention_mask=mask, return_type="logits") + overridden = distilgpt2_bridge( + padded, attention_mask=mask, position_ids=shifted, return_type="logits" + ) + + assert not torch.allclose(default, overridden, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("gap", [slice(3, 5), slice(1, 2)]) +def test_interior_mask_gap_uses_derived_positions(distilgpt2_bridge, tokens, gap) -> None: + """A mask gap that is not leading padding still shifts later positions, so the + derivation must fire for any mask, not only ones starting with a pad.""" + mask = torch.ones(tokens.shape, dtype=torch.long) + mask[0, gap] = 0 + gapped = tokens.clone() + gapped[0, gap] = PAD_ID + expected = utils.get_offset_position_ids(0, mask) + + with torch.no_grad(): + derived = distilgpt2_bridge(gapped, attention_mask=mask, return_type="logits") + supplied = distilgpt2_bridge( + gapped, attention_mask=mask, position_ids=expected, return_type="logits" + ) + + torch.testing.assert_close(derived, supplied, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("mask_kind", ["all_ones", "right_padded"]) +def test_no_position_ids_injected_when_unnecessary(distilgpt2_bridge, tokens, mask_kind) -> None: + """Masks whose attended tokens already sit at their default positions must not + get position_ids injected: it is a no-op at best, and models whose forward does + not accept position_ids would raise. + """ + if mask_kind == "all_ones": + passed, mask = tokens, torch.ones(tokens.shape, dtype=torch.long) + else: + passed, mask = _pad(tokens, 3, "right") + + seen: dict[str, object] = {} + original = distilgpt2_bridge.original_model.forward + + def _spy(*args, **kwargs): + seen["position_ids"] = kwargs.get("position_ids") + return original(*args, **kwargs) + + distilgpt2_bridge.original_model.forward = _spy + try: + with torch.no_grad(): + distilgpt2_bridge(passed, attention_mask=mask, return_type="logits") + finally: + distilgpt2_bridge.original_model.forward = original + + assert seen["position_ids"] is None + + +def test_cached_step_with_left_padding(distilgpt2_bridge, tokens) -> None: + """With a KV cache the mask spans past+new while input_ids is only the new + token, so the derivation must be offset back to the tokens being passed.""" + n_pad = 3 + padded, mask = _pad(tokens, n_pad, "left") + + with torch.no_grad(): + prefill = distilgpt2_bridge.original_model(padded, attention_mask=mask, use_cache=True) + new_token = tokens[:, -1:].clone() + extended = torch.cat([mask, torch.ones(1, 1, dtype=torch.long)], dim=1) + step = distilgpt2_bridge( + new_token, + attention_mask=extended, + past_key_values=prefill.past_key_values, + return_type="logits", + ) + + assert step.shape[:2] == (1, 1) + assert torch.isfinite(step).all() diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index c0318a5b6..8d7798fc1 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1551,6 +1551,37 @@ def forward( position_ids.masked_fill_(attention_mask == 0, 1) kwargs["position_ids"] = position_ids + # Any masked-out token shifts the absolute position of every real token + # after it, so positions must be derived from the mask rather than left + # to HF's default arange. This is the same derivation HookedTransformer + # applies in pos_embed; without it the bridge silently returns wrong + # logits. An all-ones mask reduces to arange, so this is a no-op there. + # + # The mask spans any cached prefix as well as the new tokens, so it is + # offset back to just the tokens actually being passed — matching how + # AbstractAttention/PosEmbed use past_kv_pos_offset. + if ( + attention_mask is not None + and "position_ids" not in kwargs + and not _is_inputs_embeds + and attention_mask.ndim == 2 + and isinstance(input_ids, torch.Tensor) + and input_ids.ndim == 2 + and attention_mask.shape[1] >= input_ids.shape[1] + ): + _derived = utils.get_offset_position_ids(0, attention_mask) + _arange = torch.arange(attention_mask.shape[1], device=_derived.device) + # Only intervene when the mask actually moves an attended token off + # its default position — i.e. some masked token precedes a real one + # (left padding, or an interior gap). Pure right padding and an + # all-ones mask already agree with arange, and injecting position_ids + # there would be a no-op at best and unsupported by models whose + # forward does not take them at worst. + if bool(((_derived != _arange) & (attention_mask != 0)).any()): + kwargs["position_ids"] = _derived[ + :, attention_mask.shape[1] - input_ids.shape[1] : + ] + if attention_mask is not None: kwargs["attention_mask"] = attention_mask if past_key_values is not None: