fix(bridge): gate batched-list position_ids on the target model - #1627
Open
sohv wants to merge 1 commit into
Open
fix(bridge): gate batched-list position_ids on the target model#1627sohv wants to merge 1 commit into
sohv wants to merge 1 commit into
Conversation
Batched list input builds an attention_mask and position_ids itself so pad tokens don't contaminate the forward. The mask is safe for any model, but the position_ids were handed over unchecked: a forward taking neither position_ids nor **kwargs raises TypeError where it would have returned logits. This is the gap jlarson4 raised while reviewing TransformerLensOrg#1610. That PR added _accepts_derived_position_ids() and gated the main forward() derivation, but these two sites were left for a follow-up because no model could be shown to fail there. The LLaDA test harness builds a fixed-signature forward in process, which reproduces it: TypeError: TinyLLaDAModelLM.forward() got an unexpected keyword argument 'position_ids' Gate both sites on the same helper. The attention_mask stays unconditional -- it is safe everywhere, and withholding it would reintroduce the padding contamination this branch exists to prevent. A single unbatched string was never affected, and the test asserts that alongside the batched case. The regression test wraps its forward spy in functools.wraps: the gate reads that forward's signature, so a bare (*args, **kwargs) wrapper would look like it accepts position_ids and silently defeat the check under test. Fixes TransformerLensOrg#1626 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 task
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
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.
Fixes #1626.
Batched list input builds an
attention_maskandposition_idsitself so pad tokens don't contaminate the forward. The mask is safe for any model, but theposition_idswere handed over unchecked — a forward taking neitherposition_idsnor**kwargsraisesTypeErrorwhere it would otherwise have returned logits.This is the gap that was raised while reviewing #1610. That PR added
_accepts_derived_position_ids()and gated the mainforward()derivation, but we left these two sites for a follow-up because I couldn't produce a model that demonstrably failed there. The LLaDA test harness builds a fixed-signature forward in-process, which reproduces it:What I changed
Both sites now call the same
_accepts_derived_position_ids()helper, so there's no new predicate. Theattention_maskstays unconditional -it's safe for every model, and withholding it would reintroduce the padding contamination the branch exists to prevent.Verification
One regression test in the LLaDA suite, red on
dev-4.xwith theTypeErrorabove and green with the fix. It asserts both halves: noposition_idsreaches the model, and theattention_maskstill does. A single unbatched string was never affected and is covered as a control.Full unit + integration + acceptance: 6594 passed, 1 failed. The failure is
test_bridge_hooked_parity_multi_step_optimization, which fails identically on unmodified code and is--ignored on the macOS CI job. mypy clean.