fix(bridge): derive position_ids from attention_mask for left-padded input - #1610
fix(bridge): derive position_ids from attention_mask for left-padded input#1610sohv wants to merge 1 commit into
Conversation
| # 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 |
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()) |
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?
c37a8d8 to
cdc2af2
Compare
|
Thank you @jlarson4 for bringing these issues to my attention both were real. I reproduced these issues as you described and have pushed the revised code in commit cdc2af2. Instead of separately patching the two cases, I switched the derivation to Cached step. Reproduced: Interior gap. Reproduced: 6.752e+00 on gpt2 in compat mode against a 0.000e+00 unpadded control, with the gate correctly not firing. Rather than widen the predicate I dropped it — an all-ones mask gives One thing you did not flag that this surfaced: my pad-position convention was also wrong. I had Verification:
One correction to the PR description: the "this PR alone" row in the compounding table was 6.688155 and is now 5.396149, since pad positions inherit the previous index and that changes the still-unmasked aggregate #1607 produces. The "both" row is unchanged at 4.814578. |
…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). 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. The derivation fires only when the mask actually moves an attended token off its default position, i.e. when some masked token precedes a real one. That covers left padding and interior mask gaps. Pure right padding and all-ones masks already agree with arange, so they are left alone: injecting position_ids there is a no-op at best, and breaks models whose forward does not accept the argument or which compute their own position streams (multimodal mRoPE). With a KV cache the mask spans past+new while input_ids holds only the new tokens, so the derived positions are sliced back to the tokens 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, precedence of an explicit position_ids, interior mask gaps, a cached decode step, and that no position_ids are injected when the mask does not require it. 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>
cdc2af2 to
32dd94d
Compare
|
CI caught a regression in my last push so this needed a third pass. Dropping the gate entirely was an over-correction. You asked me to widen it to any mask with a gap; I widened it to any mask at all, which made the derivation fire for cases that never needed it. Right padding and all-ones masks already agree with Fixed in 32dd94d: the derivation now fires only when the mask actually moves an attended token off its default position — some masked token precedes a real one. That is exactly left padding and interior gaps, and excludes right padding and all-ones masks.
One unrelated failure in that run: |
2 participants
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: