Allow position_ids_start=2 on DataCollatorWithFlattening for RoBERTa etc.

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.
+
     <Tip warning={true}>
 
     Using `DataCollatorWithFlattening` will flatten the entire mini batch into single long sequence.
     Make sure your attention computation is able to handle it!
 
+    </Tip>
+
+    <Tip warning={true}>
+
+    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.
+
     </Tip>
     """
 
@@ -1386,13 +1408,15 @@ def __init__(
         separator_id=-100,
         return_flash_attn_kwargs=False,
         return_seq_idx=False,
+        position_ids_start=0,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.return_position_ids = return_position_ids
         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:

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(
         (

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):

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)