Stop enable_padding_free_metadata writing seq_lengths into the caller's examples

tests/utils/test_packing.py

@@ -27,6 +27,7 @@
     patch_hybrid_linear_attention_varlen,
 )
 
+import copy
 import inspect
 import logging
 from contextlib import ExitStack
@@ -160,6 +161,80 @@ def test_mask_packed_sequence_boundaries_across_multiple_rows():
     assert torch.any(flat != -100)
 
 
+def test_enable_padding_free_metadata_does_not_mutate_examples():
+    # The wrapper derives seq_lengths from input_ids when the dataset does not
+    # carry them. It used to cache that back onto the example, so a collator
+    # wrapper wrote a new column into the caller's rows; the sibling
+    # enable_sample_packing collects the same lengths without writing back.
+    collator = _PaddingFreeCollator()
+    trainer = SimpleNamespace(
+        data_collator = collator,
+        args = SimpleNamespace(remove_unused_columns = True),
+    )
+    enable_padding_free_metadata(_DummyModel(), trainer)
+
+    examples = [
+        {"input_ids": [1, 2, 3], "labels": [1, 2, 3]},
+        {"input_ids": [4, 5], "labels": [4, 5]},
+    ]
+    before = copy.deepcopy(examples)
+
+    batch = trainer.data_collator.torch_call(examples)
+
+    assert examples == before, "collator wrapper mutated the caller's examples"
+    assert torch.equal(batch["packed_seq_lengths"], torch.tensor([3, 2], dtype = torch.int32))
+
+    # explicit seq_lengths are still honoured, and still not written back
+    explicit = [{"input_ids": [1, 2, 3], "seq_lengths": [2, 1]}]
+    explicit_before = copy.deepcopy(explicit)
+    batch = trainer.data_collator.torch_call(explicit)
+    assert explicit == explicit_before
+    assert torch.equal(batch["packed_seq_lengths"], torch.tensor([2, 1], dtype = torch.int32))
+
+
+def test_enable_padding_free_metadata_still_hands_derived_lengths_to_the_collator():
+    # Not writing to the caller's row must not mean the wrapped collator stops
+    # seeing the lengths. TRL's padding-free collator decides whether to use
+    # seq_lengths at all from examples[0], then reads it off every row, so a row
+    # that reaches it without the key (or with a null one) changes the
+    # position_ids it builds, or makes it sum(None).
+    collator = _PaddingFreeCollator()
+    trainer = SimpleNamespace(
+        data_collator = collator,
+        args = SimpleNamespace(remove_unused_columns = True),
+    )
+    enable_padding_free_metadata(_DummyModel(), trainer)
+
+    examples = [
+        {"input_ids": [1, 2, 3], "labels": [1, 2, 3]},
+        {"input_ids": [4, 5], "labels": [4, 5]},
+    ]
+    before = copy.deepcopy(examples)
+
+    trainer.data_collator.torch_call(examples)
+
+    assert examples == before, "collator wrapper mutated the caller's examples"
+    assert [row["seq_lengths"] for row in collator.seen] == [[3], [2]]
+    assert [row["labels"] for row in collator.seen] == [[1, 2, 3], [4, 5]]
+
+    # a row carrying the column with no value counts as missing here, so it has
+    # to be replaced rather than passed through
+    nulled = [{"input_ids": [1, 2], "seq_lengths": None}]
+    nulled_before = copy.deepcopy(nulled)
+
+    trainer.data_collator.torch_call(nulled)
+
+    assert nulled == nulled_before
+    assert [row["seq_lengths"] for row in collator.seen] == [[2]]
+
+    # nothing was derived, so the caller's own list goes straight through
+    explicit = [{"input_ids": [1, 2, 3], "seq_lengths": [2, 1]}]
+
+    trainer.data_collator.torch_call(explicit)
+
+    assert collator.seen is explicit
+
+
 def test_configure_sample_packing():
     config = SimpleNamespace()
     configure_sample_packing(config)
@@ -978,9 +1053,11 @@ def __init__(self):
         self.padding_free = True
         self.return_position_ids = False
         self.calls = 0
+        self.seen = None
 
     def torch_call(self, examples):
         self.calls += 1
+        self.seen = examples
         return {
             "input_ids": torch.tensor([[0]], dtype = torch.long),
             "examples_seen": self.calls,

unsloth/utils/packing.py

@@ -252,18 +252,25 @@ def enable_padding_free_metadata(model, trainer):
 
     def torch_call_with_padding_free_metadata(examples: Sequence[dict]):
         seq_lengths: list[int] = []
+        collated = examples
         if examples and isinstance(examples[0], dict):
-            for example in examples:
+            for index, example in enumerate(examples):
                 lengths = example.get("seq_lengths")
                 if lengths is None:
                     ids = example.get("input_ids")
                     if ids is None:
                         continue
                     lengths = [len(ids)]
-                    example["seq_lengths"] = lengths
+                    # The wrapped collator decides whether to use seq_lengths at
+                    # all from the first example and then reads it off every row,
+                    # so the derived lengths still have to reach it. Put them on a
+                    # shallow copy instead of the caller's own row.
+                    if collated is examples:
+                        collated = list(examples)
+                    collated[index] = {**example, "seq_lengths": lengths}
                 seq_lengths.extend(lengths)
 
-        batch = original_torch_call(examples)
+        batch = original_torch_call(collated)
         if seq_lengths:
             # Labels left alone for the same reason as enable_sample_packing: num_items_in_batch is counted off
             # this batch and the zoo's discount of the boundary targets is idempotent.