From 15d9553380eb84b032372e0927bfc6d9729c3d62 Mon Sep 17 00:00:00 2001 From: Sohan Venkatesh Date: Sat, 8 Aug 2026 17:22:59 +0100 Subject: [PATCH] fix(bridge): gate batched-list position_ids on the target model 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 #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 #1626 Co-Authored-By: Claude Opus 5 --- .../model_bridge/test_llada_adapter.py | 49 +++++++++++++++++++ .../model_bridge/transformer_bridge.py | 15 ++++-- 2 files changed, 60 insertions(+), 4 deletions(-) diff --git a/tests/integration/model_bridge/test_llada_adapter.py b/tests/integration/model_bridge/test_llada_adapter.py index 9b33b2444..2078560d9 100644 --- a/tests/integration/model_bridge/test_llada_adapter.py +++ b/tests/integration/model_bridge/test_llada_adapter.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy +import functools import gc import math import weakref @@ -663,6 +664,54 @@ def test_left_padding_does_not_inject_unsupported_position_ids(models: TinyModel 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: + """Batched list input builds its own attention_mask and position_ids so pad + tokens don't contaminate the forward (#1626). The mask is safe for any model; + the position_ids are not, and this forward takes neither them nor **kwargs. + + A local bridge rather than the module fixture: this needs a tokenizer, and + attaching one to the shared instance would leak into the other tests. The + tokenizer is given a BOS so the path under test is reached independently of + BOS handling elsewhere. + """ + local = _build_models() + tokenizer = _offline_tokenizer() + tokenizer.bos_token = "" + local.bridge.tokenizer = tokenizer + + seen: dict = {} + original = local.bridge.original_model.forward + + # functools.wraps so inspect.signature() still resolves to the real forward: + # the gate reads that signature, and a bare (*args, **kwargs) spy would look + # like it accepts position_ids and defeat the check under test. + @functools.wraps(original) + def _spy(*args, **kwargs): + seen.clear() + seen.update(kwargs) + return original(*args, **kwargs) + + local.bridge.original_model.forward = _spy + try: + with torch.inference_mode(): + logits = local.bridge(["token_5 token_7 token_9", "token_5"], return_type="logits") + batched = dict(seen) + with torch.inference_mode(): + local.bridge("token_5 token_7 token_9", return_type="logits") + unbatched = dict(seen) + finally: + local.bridge.original_model.forward = original + + assert logits.shape[0] == 2 + assert "position_ids" not in batched + # The mask is still supplied — withholding it would reintroduce the padding + # contamination this branch exists to prevent. + assert "attention_mask" in batched + # Control: a single unbatched string never reached this branch, so the gate + # must not have changed anything for it either. + assert "position_ids" not in unbatched + + def test_run_with_cache_exposes_hooks_without_hf_output_attentions( models: TinyModels, ) -> None: diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 1095fb67b..5bb24b9eb 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -1734,7 +1734,11 @@ def forward( ).to(self.cfg.device) finally: self.tokenizer.padding_side = _prev_side - if "position_ids" not in kwargs: + # Gated on the target for the same reason the derivation below is: + # a fixed-signature forward raises TypeError on the kwarg, and a + # model that owns its own position derivation is overridden by it + # (#1626). + if "position_ids" not in kwargs and self._accepts_derived_position_ids(): position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) kwargs["position_ids"] = position_ids @@ -2142,9 +2146,12 @@ def _generate_tokens( ).to(self.cfg.device) self.tokenizer.padding_side = _prev_side forward_kwargs["attention_mask"] = attn_mask - position_ids = attn_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attn_mask == 0, 1) - 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(): + position_ids = attn_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attn_mask == 0, 1) + forward_kwargs["position_ids"] = position_ids if gen_step_idx == 0: if pixel_values is not None: forward_kwargs["pixel_values"] = pixel_values