Fix explicit position_ids generation on Transformers 4
Browse filesExplicit `position_ids` passed to `generate()` crash on the first decoding step under Transformers 4.57.6. The generation loop keeps the prompt-length positions, while `prepare_inputs_for_generation()` indexes the next cache position beyond their end.
Extend explicit positions after generation updates only when the parent implementation leaves them unchanged. Preserve each row's final position and advance by the actual number of accepted tokens. Transformers 5.22.1 already extends positions, so it retains its existing behavior.
Validation:
- Seven new tests pass under both Transformers 4.57.6 and 5.22.1. All seven fail under 4.57.6 before this fix.
- Coverage includes cached and uncached decoding, left padding, custom per-row positions, beam search, multi-token updates, and caller-input preservation. Generation is checked against independent full-forward decoding.
- Existing 65-case suites introduce no additional failures: 4.57.6 has 62 passes, 2 known unrelated failures and 1 skip; 5.22.1 has 62 passes and 3 known unrelated failures. Embeds-only, new-token-only prefix-cache reuse, and the existing TF5 chunked-prefill issues remain outside this change.
- Full BF16 RL700 on a B300 under Transformers 5.22.1: explicit-position cached and uncached outputs match their implicit-position controls. Full target plus DFlash2 also matches autoregressive output in six cases, including forced budgets 2, 3 and 64 after the companion convolution fix.
Companion fix for the separate DFlash short-block crash: https://github.com/z-lab/dflash/pull/180.
- modeling_ricursive.py +19 -0
- tests/test_generation_positions.py +93 -0
|
@@ -858,6 +858,25 @@ class RicursiveForCausalLM(RicursivePreTrainedModel, GenerationMixin):
|
|
| 858 |
attentions=outputs.attentions,
|
| 859 |
)
|
| 860 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 861 |
def prepare_inputs_for_generation(
|
| 862 |
self,
|
| 863 |
input_ids: torch.LongTensor,
|
|
|
|
| 858 |
attentions=outputs.attentions,
|
| 859 |
)
|
| 860 |
|
| 861 |
+
def _update_model_kwargs_for_generation(
|
| 862 |
+
self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1
|
| 863 |
+
):
|
| 864 |
+
"""Advance explicit positions when the generation backend leaves them unchanged."""
|
| 865 |
+
position_ids = model_kwargs.get("position_ids")
|
| 866 |
+
model_kwargs = super()._update_model_kwargs_for_generation(
|
| 867 |
+
outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder,
|
| 868 |
+
num_new_tokens=num_new_tokens,
|
| 869 |
+
)
|
| 870 |
+
# Transformers 4 leaves positions unchanged; Transformers 5 already extends them.
|
| 871 |
+
if position_ids is not None and model_kwargs.get("position_ids") is position_ids:
|
| 872 |
+
offsets = torch.arange(
|
| 873 |
+
1, num_new_tokens + 1, device=position_ids.device, dtype=position_ids.dtype
|
| 874 |
+
)
|
| 875 |
+
model_kwargs["position_ids"] = torch.cat(
|
| 876 |
+
[position_ids, position_ids[:, -1:] + offsets], dim=-1
|
| 877 |
+
)
|
| 878 |
+
return model_kwargs
|
| 879 |
+
|
| 880 |
def prepare_inputs_for_generation(
|
| 881 |
self,
|
| 882 |
input_ids: torch.LongTensor,
|
|
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run with pytest under Transformers 4.57.6 and 5.22.1; no model weights needed."""
|
| 2 |
+
|
| 3 |
+
import importlib.util
|
| 4 |
+
import sys
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
import pytest
|
| 8 |
+
import torch
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
SPEC = importlib.util.spec_from_file_location(
|
| 13 |
+
"ripple_generation_regression", ROOT / "modeling_ricursive.py",
|
| 14 |
+
submodule_search_locations=[str(ROOT)],
|
| 15 |
+
)
|
| 16 |
+
MODEL = importlib.util.module_from_spec(SPEC)
|
| 17 |
+
sys.modules[SPEC.name] = MODEL
|
| 18 |
+
SPEC.loader.exec_module(MODEL)
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
@pytest.fixture
|
| 22 |
+
def model():
|
| 23 |
+
torch.manual_seed(19)
|
| 24 |
+
config = MODEL.RicursiveConfig(
|
| 25 |
+
vocab_size=97, hidden_size=32, intermediate_size=64,
|
| 26 |
+
num_hidden_layers=5, num_attention_heads=4, num_key_value_heads=2,
|
| 27 |
+
head_dim=8, sliding_window=16, max_position_embeddings=128,
|
| 28 |
+
pad_token_id=0, eos_token_id=None, use_valres=True,
|
| 29 |
+
)
|
| 30 |
+
return MODEL.RicursiveForCausalLM(config).eval()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
@pytest.mark.parametrize("use_cache", [True, False])
|
| 34 |
+
@pytest.mark.parametrize("custom_positions", [False, True])
|
| 35 |
+
@torch.inference_mode()
|
| 36 |
+
def test_explicit_positions_extend_during_generation(model, use_cache, custom_positions):
|
| 37 |
+
prompt = torch.tensor([[3, 4, 5, 6, 7, 8], [0, 0, 3, 4, 5, 6]])
|
| 38 |
+
mask = prompt.ne(0).long()
|
| 39 |
+
positions = (mask.cumsum(-1) - 1).masked_fill(mask == 0, 0)
|
| 40 |
+
if custom_positions:
|
| 41 |
+
positions = torch.tensor([[2, 4, 5, 8, 9, 12], [0, 0, 17, 19, 20, 24]])
|
| 42 |
+
original_positions = positions.clone()
|
| 43 |
+
expected = prompt.clone()
|
| 44 |
+
expected_mask, expected_positions = mask.clone(), positions.clone()
|
| 45 |
+
# Independent full-forward decoding checks both offsets and newly added positions.
|
| 46 |
+
for _ in range(4):
|
| 47 |
+
logits = model(expected, attention_mask=expected_mask,
|
| 48 |
+
position_ids=expected_positions, use_cache=False).logits
|
| 49 |
+
expected = torch.cat([expected, logits[:, -1].argmax(-1, keepdim=True)], dim=-1)
|
| 50 |
+
expected_mask = torch.cat([expected_mask, torch.ones((2, 1), dtype=mask.dtype)], dim=-1)
|
| 51 |
+
expected_positions = torch.cat([expected_positions, expected_positions[:, -1:] + 1], dim=-1)
|
| 52 |
+
|
| 53 |
+
observed = []
|
| 54 |
+
def capture(_module, _args, kwargs):
|
| 55 |
+
observed.append(kwargs["position_ids"].clone())
|
| 56 |
+
handle = model.register_forward_pre_hook(capture, with_kwargs=True)
|
| 57 |
+
try:
|
| 58 |
+
actual = model.generate(prompt, attention_mask=mask, position_ids=positions,
|
| 59 |
+
max_new_tokens=4, do_sample=False, use_cache=use_cache)
|
| 60 |
+
finally:
|
| 61 |
+
handle.remove()
|
| 62 |
+
torch.testing.assert_close(actual, expected)
|
| 63 |
+
torch.testing.assert_close(positions, original_positions)
|
| 64 |
+
torch.testing.assert_close(observed[0], original_positions)
|
| 65 |
+
for step, seen in enumerate(observed[1:], start=1):
|
| 66 |
+
future = original_positions[:, -1:] + torch.arange(1, step + 1)
|
| 67 |
+
wanted = future[:, -1:] if use_cache else torch.cat([original_positions, future], dim=-1)
|
| 68 |
+
torch.testing.assert_close(seen, wanted)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
@torch.inference_mode()
|
| 72 |
+
def test_explicit_positions_with_beam_search(model):
|
| 73 |
+
prompt = torch.tensor([[3, 4, 5, 6, 7, 8]])
|
| 74 |
+
expected = model.generate(prompt, num_beams=3, max_new_tokens=4, do_sample=False)
|
| 75 |
+
actual = model.generate(prompt, position_ids=torch.arange(6).unsqueeze(0),
|
| 76 |
+
num_beams=3, max_new_tokens=4, do_sample=False)
|
| 77 |
+
torch.testing.assert_close(actual, expected)
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
@pytest.mark.parametrize("num_new_tokens", [1, 3])
|
| 81 |
+
@torch.inference_mode()
|
| 82 |
+
def test_position_update_appends_exactly_the_accepted_token_count(model, num_new_tokens):
|
| 83 |
+
prompt = torch.tensor([[3, 4, 5, 6]])
|
| 84 |
+
positions = torch.tensor([[7, 8, 10, 13]])
|
| 85 |
+
output = model(prompt, position_ids=positions, use_cache=True)
|
| 86 |
+
updated = model._update_model_kwargs_for_generation(
|
| 87 |
+
output,
|
| 88 |
+
{"position_ids": positions, "attention_mask": torch.ones_like(prompt),
|
| 89 |
+
"cache_position": torch.arange(4), "use_cache": True},
|
| 90 |
+
num_new_tokens=num_new_tokens,
|
| 91 |
+
)
|
| 92 |
+
expected = torch.cat([positions, positions[:, -1:] + torch.arange(1, num_new_tokens + 1)], dim=-1)
|
| 93 |
+
torch.testing.assert_close(updated["position_ids"], expected)
|