From 4560dab6cab7262e86f02f95d09d845590a820a1 Mon Sep 17 00:00:00 2001 From: Tom Aarsen Date: Fri, 24 Jul 2026 17:23:36 +0200 Subject: [PATCH 1/2] Allow position_ids_start=2 on DataCollatorWithFlattening for RoBERTa etc. --- src/transformers/data/data_collator.py | 26 +++++++++++++++++- .../modeling_flash_attention_utils.py | 4 ++- tests/trainer/test_data_collator.py | 27 +++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py index 8412ab5ae25a..b8074810daea 100644 --- a/src/transformers/data/data_collator.py +++ b/src/transformers/data/data_collator.py @@ -1371,11 +1371,33 @@ class DataCollatorWithFlattening(DefaultDataCollator): - optionally returns the kwargs contained in FlashAttentionKwargs - optionally returns seq_idx indicating which sequence each token belongs to + Args: + return_position_ids (`bool`, *optional*, defaults to `True`): + Whether to return `position_ids`, which restart at `position_ids_start` for every flattened sequence. + separator_id (`int`, *optional*, defaults to -100): + The value used to separate sequences within the concatenated `labels`. + return_flash_attn_kwargs (`bool`, *optional*, defaults to `False`): + Whether to return the sequence boundary kwargs used by FlashAttention (`cu_seq_lens_q`, `cu_seq_lens_k`, + `max_length_q` and `max_length_k`). + return_seq_idx (`bool`, *optional*, defaults to `False`): + Whether to return `seq_idx`, which indicates the sequence each token belongs to. + position_ids_start (`int`, *optional*, defaults to 0): + The position id given to the first token of every sequence. + Using `DataCollatorWithFlattening` will flatten the entire mini batch into single long sequence. Make sure your attention computation is able to handle it! + + + + + Not every model numbers positions from 0 when computing position ids itself. RoBERTa-like architectures (among + them `roberta`, `xlm_roberta`, `camembert`, `mpnet` and `esm` with absolute position embeddings) start at + `padding_idx + 1`, which is usually 2. Passing such models the default 0-based `position_ids` silently shifts + the position embeddings and degrades their outputs, so set `position_ids_start` accordingly for those models. + """ @@ -1386,6 +1408,7 @@ def __init__( separator_id=-100, return_flash_attn_kwargs=False, return_seq_idx=False, + position_ids_start=0, **kwargs, ): super().__init__(*args, **kwargs) @@ -1393,6 +1416,7 @@ def __init__( self.separator_id = separator_id self.return_flash_attn_kwargs = return_flash_attn_kwargs self.return_seq_idx = return_seq_idx + self.position_ids_start = position_ids_start self._int_64_keys = {"labels", "position_ids", "input_ids"} self._batch_dim_keys = {"labels", "position_ids", "input_ids", "seq_idx"} self._py_int_keys = {"max_length_q", "max_length_k"} @@ -1427,7 +1451,7 @@ def __call__(self, features, return_tensors=None, separator_id=None): else: batch["labels"] += [separator_id] + input_ids[1:] if self.return_position_ids: - batch["position_ids"] += list(range(len(input_ids))) + batch["position_ids"] += [self.position_ids_start + i for i in range(len(input_ids))] if self.return_seq_idx: batch["seq_idx"] += [seq_idx for _ in range(len(input_ids))] if self.return_flash_attn_kwargs: diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index e0f302f90bb9..f6b261c37665 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -474,7 +474,9 @@ def prepare_fa_kwargs_from_position_ids(position_ids): tensor_kwargs = {"dtype": torch.int32, "device": position_ids.device} position_ids = position_ids.reshape(-1) - indices_q = (position_ids == 0).nonzero().view(-1) + # Packed sequences all restart from the same first position id, but it is not always 0 + # (RoBERTa-like models start at padding_idx + 1) + indices_q = (position_ids == position_ids.min()).nonzero().view(-1) cu_seq_lens_q = torch.cat( ( diff --git a/tests/trainer/test_data_collator.py b/tests/trainer/test_data_collator.py index c550364208af..8ace704b7b1a 100644 --- a/tests/trainer/test_data_collator.py +++ b/tests/trainer/test_data_collator.py @@ -48,6 +48,8 @@ if is_torch_available(): import torch + from transformers.modeling_flash_attention_utils import _is_packed_sequence, prepare_fa_kwargs_from_position_ids + class DataCollatorTestMixin: """Mixin providing common setup and utility methods for data collator tests.""" @@ -346,6 +348,31 @@ def test_basic_flattening(self): for key in ["attention_mask", "cu_seq_lens_k", "cu_seq_lens_q", "seq_idx"]: self.assertNotIn(key, batch) + def test_position_ids_start(self): + """Test position_ids for models numbering positions from padding_idx + 1, like RoBERTa.""" + for return_tensors in ["pt", "np"]: + collator = DataCollatorWithFlattening(return_tensors=return_tensors, position_ids_start=2) + batch = collator(self._get_features()) + + self.assertEqual(batch["position_ids"][0].tolist(), [2, 3, 4, 2, 3, 4, 5, 6, 7, 2, 3, 4, 5, 6, 7, 8]) + + def test_position_ids_flash_attn_boundary_inference(self): + """Test FlashAttention sequence boundary inference from collator position_ids, 0-based or not.""" + for position_ids_start in [0, 2]: + collator = DataCollatorWithFlattening(return_tensors="pt", position_ids_start=position_ids_start) + batch = collator(self._get_features()) + + self.assertTrue(_is_packed_sequence(batch["position_ids"], batch_size=1)) + cu_seq_lens, max_lengths = prepare_fa_kwargs_from_position_ids(batch["position_ids"]) + self.assertEqual(cu_seq_lens[0].tolist(), [0, 3, 9, 16]) + self.assertEqual(cu_seq_lens[1].tolist(), [0, 3, 9, 16]) + self.assertEqual(int(max_lengths[0]), 7) + self.assertEqual(int(max_lengths[1]), 7) + + # A batch with a single sequence stays on the regular FA path + single = collator([self._get_features()[0]]) + self.assertFalse(_is_packed_sequence(single["position_ids"], batch_size=1)) + def test_flash_attn_kwargs(self): """Test flattening with Flash Attention kwargs.""" collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True) From 9d1c8e8828d279916b498c05566940afcb17369b Mon Sep 17 00:00:00 2001 From: Tom Aarsen Date: Mon, 27 Jul 2026 13:39:23 +0200 Subject: [PATCH 2/2] Add RoBERTa flash-attn integration test for position_ids_start --- tests/models/roberta/test_modeling_roberta.py | 45 ++++++++++++++++++- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/tests/models/roberta/test_modeling_roberta.py b/tests/models/roberta/test_modeling_roberta.py index ea4ea2a62845..ae380f153e27 100644 --- a/tests/models/roberta/test_modeling_roberta.py +++ b/tests/models/roberta/test_modeling_roberta.py @@ -14,8 +14,17 @@ import unittest -from transformers import RobertaConfig, is_torch_available -from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device +import pytest + +from transformers import DataCollatorWithFlattening, RobertaConfig, is_torch_available +from transformers.testing_utils import ( + TestCasePlus, + require_flash_attn, + require_torch, + require_torch_accelerator, + slow, + torch_device, +) from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester @@ -505,6 +514,38 @@ def test_create_position_ids_from_inputs_embeds(self): self.assertEqual(position_ids.shape, expected_positions.shape) self.assertTrue(torch.all(torch.eq(position_ids, expected_positions))) + @require_flash_attn + @require_torch_accelerator + @pytest.mark.flash_attn_test + def test_flash_attn_padding_free_position_ids_start(self): + """Test that flattening a batch with position_ids_start=pad_token_id + 1 matches the padded batched run.""" + config = self.model_tester.prepare_config_and_inputs()[0] + config.attn_implementation = "flash_attention_2" + model = RobertaModel(config).to(dtype=torch.bfloat16, device=torch_device).eval() + + # Sample from [2, vocab_size) so no token collides with pad_token_id, which would shift + # the position ids the model computes internally in the batched reference run + seq_len_a, seq_len_b = 3, 4 + input_ids = ids_tensor([2, seq_len_b], config.vocab_size - 2) + 2 + attention_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]], device=torch_device) + input_ids[attention_mask == 0] = config.pad_token_id + + with torch.no_grad(): + batched = model(input_ids, attention_mask=attention_mask).last_hidden_state + + collator = DataCollatorWithFlattening(position_ids_start=config.pad_token_id + 1) + packed = collator([{"input_ids": input_ids[0, :seq_len_a]}, {"input_ids": input_ids[1, :seq_len_b]}]) + self.assertEqual(packed["position_ids"][0].tolist(), [2, 3, 4, 2, 3, 4, 5]) + # Without attention_mask and cu_seq_lens, FlashAttention infers the boundaries from position_ids + with torch.no_grad(): + flattened = model( + input_ids=packed["input_ids"].to(torch_device), position_ids=packed["position_ids"].to(torch_device) + ).last_hidden_state + + # bf16 noise between the padded and packed kernel shapes reaches ~3e-2, wrong position ids differ by >1.9 + torch.testing.assert_close(batched[0, :seq_len_a], flattened[0, :seq_len_a], atol=1e-1, rtol=1e-1) + torch.testing.assert_close(batched[1, :seq_len_b], flattened[0, seq_len_a:], atol=1e-1, rtol=1e-1) + @require_torch class RobertaModelIntegrationTest(TestCasePlus):