From d922f52dfcad169944bdd690a708b74df516ea2a Mon Sep 17 00:00:00 2001 From: Jin Soo Ihm Date: Fri, 21 Aug 2026 15:41:27 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../unit_tests/components/data/test_grain_data.py | 14 ++++++++++++++ torchtitan/components/data/collators.py | 2 +- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/components/data/test_grain_data.py b/tests/unit_tests/components/data/test_grain_data.py index 97a9a92ebf..6e83497f4b 100644 --- a/tests/unit_tests/components/data/test_grain_data.py +++ b/tests/unit_tests/components/data/test_grain_data.py @@ -1010,6 +1010,20 @@ def test_unpacked_text_collator_creates_range_positions(): assert (labels[3:] == IGNORE_INDEX).all() +def test_unpacked_text_collator_pads_positions_within_context_window(): + # pad_len (18-3=15) exceeds max_context_length (9); padding positions must + # stay within the RoPE window instead of running 0..pad_len-1. + sequence = TextSequence( + input_ids=np.asarray([1, 2, 3]), + labels=np.asarray([2, 3, 4]), + ) + + inputs, _ = TextCollator.Config().build(context=CONTEXT)([sequence]) + + assert len(inputs["positions"]) == CONTEXT.num_tokens_per_batch + assert int(inputs["positions"].max()) < CONTEXT.max_context_length + + def test_pack_then_pack_then_collate_preserves_aligned_pairs(): documents = SingleDatasetConfig( source=RowsSourceConfig( diff --git a/torchtitan/components/data/collators.py b/torchtitan/components/data/collators.py index 099a8a62f6..3262ff0b85 100644 --- a/torchtitan/components/data/collators.py +++ b/torchtitan/components/data/collators.py @@ -70,7 +70,7 @@ def __call__(self, rows: Sequence[TextSequence]) -> TrainerBatch: input_ids = torch.nn.functional.pad(input_ids, (0, pad_len)) labels = torch.nn.functional.pad(labels, (0, pad_len), value=IGNORE_INDEX) positions = torch.cat( - [positions, torch.arange(pad_len, dtype=positions.dtype)] + [positions, torch.zeros(pad_len, dtype=positions.dtype)] ) return {