fix(bridge): derive position_ids from attention_mask for left-padded input - #1610
Open
sohv wants to merge 1 commit into
Open
fix(bridge): derive position_ids from attention_mask for left-padded input#1610sohv wants to merge 1 commit into
sohv wants to merge 1 commit into
Conversation
jlarson4
reviewed
Aug 5, 2026
| # than left to HF's default arange. HookedTransformer does this via | ||
| # pos_embed; without it the bridge silently returns wrong logits. | ||
| if ( | ||
| attention_mask is not None |
Collaborator
There was a problem hiding this comment.
A manual forward(new_token, attention_mask=<full mask>, past_key_values=cache) with left padding returns logits, but raises RuntimeError: The size of tensor a (19) must match the size of tensor b (10) because the derived position_ids spans past+new while input_ids is only the new token. Can the derivation be limited to the tokens actually being passed, the way get_offset_position_ids does it (utilities/tensors.py:131), with a test covering a cached step?
| and "position_ids" not in kwargs | ||
| and not _is_inputs_embeds | ||
| and attention_mask.ndim == 2 | ||
| and bool((attention_mask[:, 0] == 0).any()) |
Collaborator
There was a problem hiding this comment.
The gate tests only column 0, so a mask with an interior gap and no leading padding still diverges from HookedTransformer. I measured 3.631e+00 on gpt2 in compat mode against a 0.000e+00 unpadded control. Would it be possible to widen this to any mask with a gap?
…input TransformerBridge.forward() did not derive position_ids from a supplied attention_mask, so masked-out tokens silently shifted the absolute position of every real token after them — no error, no NaN, just wrong logits and a wrong loss. On gpt2 the loss for one prompt moved from 4.503170 unpadded to 11.154946 with three left pads, while HookedTransformer stays invariant (drift ~1e-06). Right padding was never affected, since causality already protects it. transformer_bridge.py derived position_ids only for batched *list* input, so pre-tokenized tensors fell through to HF's plain arange and the offset was never removed. This reuses utils.get_offset_position_ids — the same helper PosEmbed and AbstractAttention already use — so the bridge shares HookedTransformer's position derivation rather than paralleling it. An explicitly supplied position_ids still wins, and an all-ones mask reduces to arange, so this is a no-op when there is no padding. The derivation is offset by any cached prefix: with past_key_values the mask spans past+new while input_ids holds only the new tokens, so positions are sliced back to the tokens actually being passed. The bridge was also inconsistent with itself before this — the same batch gave different logits depending on whether it was passed as strings or token IDs (max |logit diff| 4.142e+01) — and enable_compatibility_mode(), which documents "HookedTransformer-equivalent numerics", diverged on left-padded input while matching exactly on unpadded input. Adds integration regression tests: logit invariance under both padding sides, the same property in compatibility mode, agreement with the shared helper's derivation, precedence of an explicit position_ids, interior mask gaps, and a cached decode step. Right-padding cases are controls that pass with and without the fix. They live in the integration tier because left padding produces a fully masked query row, which the Native attention path turns into NaN until the masked-softmax fix in TransformerLensOrg#1608 lands. Fixes TransformerLensOrg#1609. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
sohv
force-pushed
the
fix/bridge-left-padding-positions
branch
from
August 5, 2026 21:35
c37a8d8 to
cdc2af2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #1609.
TransformerBridge.forward()did not deriveposition_idsfrom a suppliedattention_mask, so left-padded input silently got the wrong absolute positions — no error, no NaN, just wrong logits and a wrong loss.On gpt2, one prompt, mask supplied:
Right padding was never affected (drift ≤ 9.5e-07) — causality already protects it.
transformer_bridge.pyderivedposition_idsonly for batched list input, so pre-tokenized tensors fell through to HF's plainarangeand the padding offset was never removed. This extends the same correction that branch already applies. An explicitly suppliedposition_idsstill wins.Two consequences this also fixes:
4.142e+01).enable_compatibility_mode(), which documents "HookedTransformer-equivalent numerics", matched HT exactly on unpadded input (0.000e+00) but diverged on left-padded input.Tests
Adds
tests/integration/model_bridge/test_left_padding_positions.py: logit invariance under both padding sides, the same property in compatibility mode, and agreement between the derived and an explicitly suppliedposition_ids.Red-before / green-after: 9 passed with the fix, 5 failed without it. The right-padding cases are controls — they pass in both states, so the tests are specific to the bug rather than to padding in general.
These sit in the integration tier rather than the unit tier deliberately: left padding produces a fully masked query row, which the Native attention path turns into
NaNuntil the masked-softmax fix in #1608 lands, soboot_nativecannot express the property yet.Verification
tests/unit/model_bridge+tests/unit/test_tokenizer_padding_side.py: 3955 passed, 27 skipped, 10 xfailedpycln/isort/blackclean;mypycleanRelationship to #1607 / #1608
Independent bugs on the same path that compound. This is measured across four states (gpt2, aggregate loss on a left-padded batch; HT reference
4.814578):dev-4.xThis PR fixes the logits; #1608 fixes the loss aggregation. Batched loss is only correct with both, so reviewing this one in isolation will still show a wrong aggregate. No file overlap, so they merge in either order.
Type of change
Checklist: