Skip to content

fix(bridge): gate batched-list position_ids on the target model - #1627

Open
sohv wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
sohv:fix/batched-list-position-ids-gate
Open

fix(bridge): gate batched-list position_ids on the target model#1627
sohv wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
sohv:fix/batched-list-position-ids-gate

Conversation

@sohv

@sohv sohv commented Aug 8, 2026

Copy link
Copy Markdown

Fixes #1626.

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 otherwise have returned logits.

This is the gap that was raised while reviewing #1610. That PR added _accepts_derived_position_ids() and gated the main forward() 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:

TypeError: TinyLLaDAModelLM.forward() got an unexpected keyword argument 'position_ids'
kwargs attempted: ['attention_mask', 'position_ids']

What I changed

Both sites now call the same _accepts_derived_position_ids() helper, so there's no new predicate. The attention_mask stays 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.x with the TypeError above and green with the fix. It asserts both halves: no position_ids reaches the model, and the attention_mask still 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.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant