Repository navigation
pack_dataset with strategy="wrapped" silently duplicates data across map batches #6669
Description
Reproduction
from datasets import Dataset
from trl import pack_dataset
ds = Dataset.from_dict({"input_ids": [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]})
packed = pack_dataset(ds, seq_length=4, strategy="wrapped", map_kwargs={"batch_size": 3})
print(packed.to_dict()["input_ids"])outputs:
[[1, 2, 3, 4], [5, 6], [1, 2, 3, 4], [5, 6]]
expected:
[[1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12]]
Rows 3 to 5 come back as a second copy of rows 0 to 2. Nothing raises, and the row and token counts are both unchanged, so the corruption is silent.
The batch_size=3 above is only to keep the snippet short. The same thing happens at the default map batch size once the dataset is larger than one batch:
from datasets import Dataset
from trl import pack_dataset
n, per_row = 2500, 4
ids = [[r * per_row + k for k in range(per_row)] for r in range(n)]
ds = Dataset.from_dict({"input_ids": ids})
packed = pack_dataset(ds, seq_length=8, strategy="wrapped")
expected = {t for row in ids for t in row}
got = {t for row in packed.to_dict()["input_ids"] for t in row}
print(f"tokens lost: {len(expected - got)} / {len(expected)}")outputs:
tokens lost: 6000 / 10000
Reproduces identically on datasets 4.7.0 (the minimum TRL requires) and 5.0.1, with pyarrow 21.0.0 and 25.0.0.
Identified Cause
_pack_wrapped computes the number of values from the array's sliced view, then builds the new array from the raw child buffer:
offsets, values = columns[0].offsets, columns[0].values
values = values[offsets[0].as_py() : offsets[-1].as_py()] # slice-corrected
num_elements = len(values) # uses the correction
offsets = np.arange(0, num_elements, seq_length, dtype=...)
offsets = np.concatenate((offsets, [num_elements]))
columns = [
type(column).from_arrays(offsets.astype(...), column.values) # ignores it
for column in columns
]column.values is the whole child buffer of the underlying table, so it still contains the rows before the array's own offset. The new offsets are then applied from position 0 instead of from the slice start. The values local is corrected and then never used.
datasets.map(batched=True) hands each batch to the packing function as a slice of the underlying table, so batches after the first can start at a non-zero offset.
This looks like a regression from #5189, which rewrote the loop as a list comprehension over all columns and switched to column.values. The version before it passed the corrected values and was right.
Scope
Only wrapped is affected. bfd and bfd_split reach the same column.values pattern but are safe, because pc.filter and pc.take materialize fresh arrays with offset 0 first. So switching strategy is a workaround.
Triggering needs the table to actually be sliced, which is why it is not constant. When chunk boundaries line up with the map batch size, combine_chunks concatenates into a fresh array and masks it, and an indices mapping from shuffle, filter, or select makes datasets build batches with take, which masks it too.
On SFTTrainer this means the default path is not affected, since shuffle_dataset=True shuffles right before packing. The combination that does corrupt is packing_strategy="wrapped" with shuffle_dataset=False and dataset_num_proc set, which loses 1000 of 10000 tokens in my test.
System Info
- Platform: macOS-26.5.2-arm64-arm-64bit
- Python version: 3.12.4
- TRL version: main @ d242e58
- PyTorch version: 2.13.0
- Transformers version: 5.14.1
- Datasets version: 5.0.1
- PyArrow version: 25.0.0
- Accelerator: none, CPU-only reproduction
Checklist
- I have checked that my issue isn't already filed (see open issues)
- I have included my system information
- Any code provided is minimal, complete, and reproducible (more on MREs)
- Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
- Any traceback provided is complete
Activity