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