diff --git a/pipeline/packing.py b/pipeline/packing.py index 0fdc8fa..67fa6fb 100644 --- a/pipeline/packing.py +++ b/pipeline/packing.py @@ -9,64 +9,55 @@ logger = logging.getLogger(__name__) class SequencePacker: """ - Packs variable-length sequences into fixed-size tensors, suitable for - concatenating unequal-length training samples into uniform shapes - for DataLoader / model training. + Stream-concatenation packer for LLM training sequences. - Algorithm (Sorted Greedy Fill, based on First-Fit Decreasing heuristic): + Algorithm (streaming concat): - Input: sequences = [A(len=5), B(len=2), C(len=3)], pack_size = 8 + Input: sequences = [A(len=3), B(len=5), C(len=2)], pack_size = 6 1. Validate & Normalize - - Check 1D dimension, unify dtype, truncate overlong sequences with warning - - Result: [(A,5), (B,2), (C,3)] + - Check 1D dimension, unify dtype, warn on overlong sequences + - Result: [A, B, C] - 2. Sort by length descending (FFD) - - Result: [(A,5), (C,3), (B,2)] + 2. Stream into buffer, slice off full chunks + - buffer += A(3) -> [a1 a2 a3], pos=3 + - buffer += B(5) -> [a1 a2 a3 b1 b2 b3 b4 b5], pos=8 + pos >= 6 -> flush [a1 a2 a3 b1 b2 b3], buffer=[b4 b5], pos=2 + - buffer += C(2) -> [b4 b5 c1 c2], pos=4 + loop ends -> flush tail [b4 b5 c1 c2 PAD PAD] - 3. Greedy fill: write into a pre-allocated buffer sequentially, flush when full - - Write A(5) -> buffer = [A A A A A _ _ _], pos=5 - - Write C(3) -> pos+3=8 <= 8 -> buffer = [A A A A A C C C], pos=8 - - Buffer full -> flush as package[0], reset buffer & pos=0 - - Write B(2) -> buffer = [B B _ _ _ _ _ _], pos=2 - - Loop ends -> flush tail -> package[1] = [B B 0 0 0 0 0 0] + Output: [[a1 a2 a3 b1 b2 b3], [b4 b5 c1 c2 PAD PAD]] - Output: [package[0], package[1]] + Samples may be split across chunks — this is intentional and standard + practice in LLM training (TRL, Megatron-LM, etc.). Cross-group consistency: - When packing different key groups (e.g. sequences and loss_masks) - with separate pack() calls, tensors at the same index always have - identical lengths, so the descending sort produces the exact same - ordering. Element-level correspondence across groups is preserved. - - Performance: - - Pre-allocated buffer reused via fill_() to avoid repeated tensor creation - - Attributes cached as local variables inside the loop to reduce lookup overhead + Different tensor groups (e.g. input_ids, loss_masks) packed with + separate packer instances on samples with matching lengths produce + identical chunk boundaries. Element-level correspondence is preserved. """ - def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = None): + def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32): self.pack_size = pack_size self.pad_value = pad_value - self.dtype = dtype # None = follow input dtype - self._buffer: Tensor | None = None - self._pos = 0 + self.dtype = dtype + self._buffer: List[int] = [] + self._pos: int = 0 self._packages: List[Tensor] = [] def reset(self) -> None: - """Reset packer state for instance reuse, unlocking dtype.""" - self.dtype = None - self._buffer = None + """Reset packer state for instance reuse.""" + self._buffer = [] self._pos = 0 self._packages = [] @error_handler() def pack(self, sequences: List[Tensor]) -> List[Tensor]: """ - Pack sequences into fixed-size packages using First-Fit Decreasing. + Pack sequences via streaming concatenation into fixed-size chunks. - Sequences are sorted by length descending to minimize wasted padding. - All tensor groups (e.g. sequences, loss_masks) with matching per-item - lengths produce identical ordering, so cross-group correspondence is preserved. + Sequences are concatenated in order and sliced at pack_size boundaries. + The final chunk is padded with pad_value. Args: sequences: List of 1D input tensors. @@ -77,54 +68,33 @@ class SequencePacker: if not sequences: return [] - # --- validate & normalize in a single pass --- - normalized: list[tuple[Tensor, int]] = [] - target_dtype = self.dtype if self.dtype is not None else sequences[0].dtype + # --- validate & normalize --- + normalized: List[Tensor] = [] for i, seq in enumerate(sequences): if seq.dim() != 1: raise ValueError( f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}" ) - if seq.dtype != target_dtype: - seq = seq.to(target_dtype) - length = seq.numel() - if length > self.pack_size: - seq = seq[: self.pack_size] - length = self.pack_size - normalized.append((seq, length)) + if seq.dtype != self.dtype: + seq = seq.to(self.dtype) + normalized.append(seq) - # --- reset internal state --- - buf = self._buffer - if buf is None or buf.dtype != target_dtype: - buf = torch.full((self.pack_size,), self.pad_value, dtype=target_dtype) - self._buffer = buf - buf.fill_(self.pad_value) - self._pos = 0 + # --- stream into buffer, slice off full chunks --- + self._buffer = [] self._packages = [] - - # --- sort by length descending (FFD heuristic) --- - normalized.sort(key=lambda x: x[1], reverse=True) - - # --- greedy fill --- - buf = self._buffer - pos = self._pos - packages = self._packages pack_size = self.pack_size - pad_value = self.pad_value + buf = self._buffer - for tensor, length in normalized: - if pos + length > pack_size: - # flush current package - packages.append(buf.clone()) - buf.fill_(pad_value) - pos = 0 - buf[pos : pos + length] = tensor - pos += length + for seq in normalized: + buf.extend(seq.tolist()) + while len(buf) >= pack_size: + self._packages.append(torch.tensor(buf[:pack_size], dtype=self.dtype)) + buf = buf[pack_size:] - # flush the last (possibly partial) package - if pos > 0: - packages.append(buf.clone()) + # flush tail with padding + if buf: + padded = buf + [self.pad_value] * (pack_size - len(buf)) + self._packages.append(torch.tensor(padded, dtype=self.dtype)) - # write back state - self._pos = pos + self._pos = len(buf) return self._packages diff --git a/tests/test_packing.py b/tests/test_packing.py index 28130bd..a84c510 100644 --- a/tests/test_packing.py +++ b/tests/test_packing.py @@ -15,16 +15,13 @@ class TestSequencePacker: torch.tensor([6, 7, 8, 9], dtype=torch.int32), ] packages = packer.pack(sequences) - assert len(packages) >= 1 + assert len(packages) == 1 for pkg in packages: assert pkg.shape == (10,) - # Verify all original values are present - all_values = [] - for pkg in packages: - all_values.extend(pkg[pkg != 0].tolist()) - for val in [1, 2, 3, 4, 5, 6, 7, 8, 9]: - assert val in all_values + # Verify all original values are present in order + assert packages[0][:9].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 9] + assert packages[0][9] == 0 # padding def test_empty_list_input(self): packer = SequencePacker(pack_size=10) @@ -37,12 +34,13 @@ class TestSequencePacker: assert packages[0][:3].tolist() == [1, 2, 3] assert packages[0][3:].tolist() == [-1] * 7 - def test_truncate_long_sequence(self, caplog): + def test_long_sequence_split_across_chunks(self): + """Sequences longer than pack_size are split across multiple chunks.""" packer = SequencePacker(pack_size=5, pad_value=0) packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)]) - assert len(packages) == 1 + assert len(packages) == 2 assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert "truncating" in caplog.text.lower() or "exceeds" in caplog.text.lower() + assert packages[1].tolist() == [6, 7, 8, 0, 0] def test_padding_value(self): packer = SequencePacker(pack_size=8, pad_value=99) @@ -104,24 +102,58 @@ class TestSequencePacker: assert packages[1].tolist() == [11] + [-1] * 9 def test_cross_group_ordering(self): - """Tensor groups with identical per-item lengths are sorted identically.""" - packer = SequencePacker(pack_size=10, pad_value=0) - # sequences: lengths [3, 1, 4] -> after sort desc: [4, 3, 1] + """Separate packers for different dtypes produce identical chunk boundaries.""" + seq_packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32) + mask_packer = SequencePacker(pack_size=10, pad_value=False, dtype=torch.bool) + # sequences: lengths [3, 1, 4] seqs = [ torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([10], dtype=torch.int32), torch.tensor([4, 5, 6, 7], dtype=torch.int32), ] masks = [ - torch.tensor([True, True, True], dtype=torch.bool), - torch.tensor([True], dtype=torch.bool), - torch.tensor([True, True, True, True], dtype=torch.bool), + torch.tensor([False, False, True], dtype=torch.bool), + torch.tensor([False], dtype=torch.bool), + torch.tensor([False, False, False, True], dtype=torch.bool), ] - packed_seqs = packer.pack(seqs) - packer.reset() - packed_masks = packer.pack(masks) + packed_seqs = seq_packer.pack(seqs) + packed_masks = mask_packer.pack(masks) # Verify mask packer uses bool dtype assert packed_masks[0].dtype == torch.bool # Both groups should produce the same number of packages assert len(packed_seqs) == len(packed_masks) + + def test_stream_split_across_chunks(self): + """Sequences are split across chunks in streaming mode.""" + packer = SequencePacker(pack_size=5, pad_value=0) + packages = packer.pack([ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32), + ]) + assert len(packages) == 2 + # First chunk: [1, 2, 3, 4, 5] — first seq + part of second + assert packages[0].tolist() == [1, 2, 3, 4, 5] + # Second chunk: [6, 7, 8, 0, 0] — rest of second + padding + assert packages[1].tolist() == [6, 7, 8, 0, 0] + + def test_reset_method(self): + packer = SequencePacker(pack_size=10, pad_value=0) + seqs = [torch.tensor([1, 2, 3], dtype=torch.int32)] + packer.pack(seqs) + assert len(packer._packages) == 1 + packer.reset() + assert len(packer._packages) == 0 + assert packer._pos == 0 + assert packer._buffer == [] + + def test_no_sorting_needed(self): + """Streaming concat preserves input order, no sorting.""" + packer = SequencePacker(pack_size=4, pad_value=-1) + # short then long (fits in 2 chunks) + packages = packer.pack([ + torch.tensor([1], dtype=torch.int32), + torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32), + ]) + assert packages[0].tolist() == [1, 2, 3, 4] + assert packages[1].tolist() == [5, 6, 7, -1]