fix(bridge): gate batched-list position_ids on the target model - #1627
fix(bridge): gate batched-list position_ids on the target model#1627sohv wants to merge 2 commits into
Conversation
| forward_kwargs["position_ids"] = position_ids | ||
| # Same target gate as the forward() path: the mask is safe | ||
| # for every model, the derived positions are not (#1626). | ||
| if self._accepts_derived_position_ids(): |
There was a problem hiding this comment.
For a gate=False model that accepts position_ids (opt-125m), this gate makes cached steps fall into the torch.full(total_len - 1) fallback at line 2220, which counts pad slots, causing logit drift. Can you gate the continuation's derive/fallback injection starting at line 2213-2226 on the same helper?
| torch.testing.assert_close(bridge_logits, reference_logits, rtol=1e-5, atol=1e-6) | ||
|
|
||
|
|
||
| def test_batched_list_input_does_not_inject_unsupported_position_ids() -> None: |
There was a problem hiding this comment.
The _generate_tokens portion of this fix is untested here and neither fixed-signature architecture can reach that site. Can you add a cached-vs-uncached batched-generate parity test on a gate=False model?
|
@jlarson4 A quick request to review this PR and merge it if you have no issues. I reviewed the PR again today but let me know if there is anything to be addressed here. |
The review just posted! Sorry I didn't get it to you yesterday, was down with the flu. |
|
@jlarson4 It's totally alright and thank you for highlighting these issues. I will work on fixing them and will update the PR soon. |
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>
Review follow-up on TransformerLensOrg#1627. Gating only the batched-list prompt derivation moved the injection rather than stopping it: every branch of the cached-step block supplies position_ids, so a refused model fell through to the torch.full(total_len - 1) fallback, which counts pad slots and is therefore wrong per row for a left-padded batch. Measured on hf-internal-testing/tiny-random-OPTForCausalLM, which the gate refuses because OPTLearnedPositionalEmbedding consumes the mask and derives its own positions, while its forward would accept the kwarg: before TransformerLensOrg#1627 cached steps [[9],[2]] [[10],[3]] [[11],[4]] <- per row TransformerLensOrg#1627 as sent cached steps [[9],[9]] [[10],[10]] [[11],[11]] <- pad slots now cached steps None <- OPT derives cached-vs-uncached max |logit diff| goes 7.45e-08 -> 2.98e-01 -> 7.45e-08, so this was a regression the PR introduced and it is now removed. Gate the whole three-way block rather than each branch, so a model that owns its position derivation receives the mask alone, matching the uncached path. Adds tests/integration/model_bridge/test_batched_generate_position_ids.py. The _generate_tokens half of TransformerLensOrg#1627 had no coverage: neither fixed-signature architecture in the suite can reach that site, since LLaDA raises NotImplementedError on generate. A text-level parity test would not catch this either, because greedy argmax absorbs the drift and the decoded strings match in both states, so the tests compare logits and separately assert that no position_ids reaches a refused model. Two of them are red on 15d9553. Refs TransformerLensOrg#1626 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
15d9553 to
6b509c3
Compare
|
The cached-step fallback - Gating only the prompt derivation moved the injection instead of stopping it. Every branch of that block supplies Cached-vs-uncached max |logit diff| goes The test gap - LLaDA raises I have added I also rebased onto current |
2 participants
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.