From 598e1ce4ae7e9d2395a67c5028c2002c3b082f33 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 3 Jul 2026 16:17:27 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=E6=89=93?= =?UTF-8?q?=E5=8C=85=E6=A8=A1=E5=9D=97=EF=BC=8C=E6=96=B0=E5=A2=9E=20BFD/FF?= =?UTF-8?q?D/Greedy=20=E4=B8=89=E7=A7=8D=20bin-packing=20=E7=AE=97?= =?UTF-8?q?=E6=B3=95=EF=BC=8C=E9=BB=98=E8=AE=A4=20BFD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将 pipeline/packing.py 拆分为 packing/ 子包 (base/stream/binpack) - 新增 BfdPacker(默认)/FfDPacker/GreedyPacker，移除 StreamingPacker - 超长序列直接截断至 pack_size - group_size 语义改为"每 N 个 chunk 合并为一块"，默认 1000 - 新增 AutoTokenizer.token_to_id()，修复 ChatML 中 hacky 的 nl_id 获取 - pad_value 默认改为 2(pad_token_id)，position_ids pad=0, loss_mask pad=False - 新增 position_ids 打包后归零一致性测试 - scripts/cache_h5.py 新增 --pack-algo 参数 --- pipeline/__init__.py | 14 +- pipeline/io/export.py | 50 +++- pipeline/packing.py | 146 ----------- pipeline/packing/__init__.py | 84 +++++++ pipeline/packing/base.py | 49 ++++ pipeline/packing/binpack.py | 174 +++++++++++++ pipeline/strategies/chatml.py | 2 +- pipeline/tokenize/tokenizer.py | 6 + scripts/cache_h5.py | 25 +- tests/test_cache.py | 2 +- tests/test_io.py | 5 +- tests/test_packing.py | 447 ++++++++++++++++++++++----------- tests/test_processors.py | 3 + tests/test_strategies.py | 3 + 14 files changed, 715 insertions(+), 295 deletions(-) delete mode 100644 pipeline/packing.py create mode 100644 pipeline/packing/__init__.py create mode 100644 pipeline/packing/base.py create mode 100644 pipeline/packing/binpack.py diff --git a/pipeline/__init__.py b/pipeline/__init__.py index a1790ad..a0e9135 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -28,7 +28,13 @@ Usage:: from pipeline.pipeline import Pipeline, PipelineConfig, Stage, TransformStage from pipeline.tokenize import AutoTokenizer, ChatTemplate, train_bpe_tokenizer from pipeline.text import TextNormalizer -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + BasePacker, + pack_tensors, +) # I/O module from pipeline.io import FileScanner, HDF5Handler, export_dataset, cache_jsonl @@ -70,7 +76,11 @@ __all__ = [ "train_bpe_tokenizer", # Text processing "TextNormalizer", - "SequencePacker", + "GreedyPacker", + "FfDPacker", + "BfdPacker", + "BasePacker", + "pack_tensors", # I/O "FileScanner", "HDF5Handler", diff --git a/pipeline/io/export.py b/pipeline/io/export.py index 2061f42..7576693 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -5,13 +5,16 @@ import logging import os from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union + +import torch from datasets import Dataset +from torch import Tensor from tqdm import tqdm from pipeline.io.file_scanner import FileScanner from pipeline.io.hdf5_handler import HDF5Handler from pipeline.processors import BaseProcessor -from pipeline.packing import pack_tensors +from pipeline.packing import pack_tensors, BasePacker from pipeline.utils import error_handler logger = logging.getLogger(__name__) @@ -75,6 +78,32 @@ def export_dataset( return output_files +def merge_tensors( + tensors: List[Tensor], + group_size: int, +) -> List[Tensor]: + """Merge a list of tensors into fewer larger tensors. + + Concatenates every group_size consecutive tensors into one merged + tensor. This reduces the number of shm blocks when loading. + + Args: + tensors: List of 1D tensors. + group_size: Number of tensors to merge into each group. + + Returns: + List of merged tensors. + """ + if not tensors: + return [] + + merged: List[Tensor] = [] + for i in range(0, len(tensors), group_size): + merged.append(torch.cat(tensors[i : i + group_size])) + + return merged + + @error_handler() def cache_jsonl( files: List[str], @@ -83,6 +112,8 @@ def cache_jsonl( *, pack_size: int = -1, pad_value: int = 0, + group_size: int = 1_000, + pack_algo: Optional[str] = None, ) -> List[str]: """Tokenize JSONL files and pack them into HDF5 storage. @@ -92,6 +123,10 @@ def cache_jsonl( processor: Initialized Processor instance. pack_size: Packing length, <=0 means no packing. pad_value: Padding value. + group_size: Merge every this many packed chunks into one tensor, + <=0 means no merging. + pack_algo: Packing algorithm: 'bfd' (default), 'ffd', + 'greedy'. Only used when pack_size > 0. Returns: List of generated H5 file paths. @@ -125,16 +160,27 @@ def cache_jsonl( ) continue + if not arrows[output_keys[0]]: + logger.warning(f"No valid samples in {file_path}, skipping") + continue + if pack_size > 0: dtypes = ( dict(processor.schema.output_fields) if processor.schema is not None else None ) - output = pack_tensors(arrows, pack_size, pad_value, dtypes) + pad_values = {k: (0 if k == "position_ids" else (False if k.endswith("_mask") else pad_value)) for k in output_keys} + output = pack_tensors(arrows, pack_size, pad_value, dtypes, pad_values=pad_values, algo=pack_algo) else: output = arrows + if group_size > 0 and output[output_keys[0]]: + output = { + key: merge_tensors(tensors, group_size) + for key, tensors in output.items() + } + h5_path = HDF5Handler.save(output_dir, file_name, output) output_files.append(h5_path) logger.info(f"Saved {h5_path}") diff --git a/pipeline/packing.py b/pipeline/packing.py deleted file mode 100644 index 42531ea..0000000 --- a/pipeline/packing.py +++ /dev/null @@ -1,146 +0,0 @@ -import logging -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import Tensor - -from pipeline.utils import error_handler - -logger = logging.getLogger(__name__) - - -class SequencePacker: - """ - Stream-concatenation packer for LLM training sequences. - - Algorithm (streaming concat): - - Input: sequences = [A(len=3), B(len=5), C(len=2)], pack_size = 6 - - 1. Validate & Normalize - - Check 1D dimension, unify dtype, warn on overlong sequences - - Result: [A, B, C] - - 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] - - Output: [[a1 a2 a3 b1 b2 b3], [b4 b5 c1 c2 PAD PAD]] - - Samples may be split across chunks — this is intentional and standard - practice in LLM training (TRL, Megatron-LM, etc.). - - Cross-group consistency: - 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: Union[int, bool] = 0, - dtype: Optional[torch.dtype] = None, - ): - self.pack_size = pack_size - self.pad_value = pad_value - self.dtype = dtype - self._buffer: List = [] - self._pos: int = 0 - self._packages: List[Tensor] = [] - - def reset(self) -> 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 via streaming concatenation into fixed-size chunks. - - Sequences are concatenated in order and sliced at pack_size boundaries. - The final chunk is padded with pad_value. - - When dtype is not set at init, it is inferred from the first input tensor. - - Args: - sequences: List of 1D input tensors. - - Returns: - List of packed tensors, each with length equal to pack_size. - """ - if not sequences: - return [] - - # --- auto-infer dtype from first sequence --- - if self.dtype is None: - self.dtype = 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 != self.dtype: - seq = seq.to(self.dtype) - normalized.append(seq) - - # --- stream into buffer, slice off full chunks --- - self._buffer = [] - self._packages = [] - pack_size = self.pack_size - buf = self._buffer - - 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 tail with padding - if buf: - padded = buf + [self.pad_value] * (pack_size - len(buf)) - self._packages.append(torch.tensor(padded, dtype=self.dtype)) - - self._pos = len(buf) - return self._packages - - -def pack_tensors( - tensors: Dict[str, List[Tensor]], - pack_size: int, - pad_value: Union[int, bool] = 0, - dtypes: Optional[Dict[str, torch.dtype]] = None, -) -> Dict[str, List[Tensor]]: - """ - Pack multiple named tensor groups in parallel. - - Each group is packed independently with its own SequencePacker instance. - When dtypes is provided, packers use the declared dtype per key; - otherwise dtype is auto-inferred from the first tensor in each group. - - Args: - tensors: Dict mapping key names to lists of 1D tensors. - pack_size: Fixed chunk length. - pad_value: Padding value for non-bool tensors. - dtypes: Optional per-key dtype declarations. - - Returns: - Dict mapping key names to lists of packed tensors. - """ - if dtypes is None: - dtypes = {} - - output: Dict[str, List[Tensor]] = {} - for key, seqs in tensors.items(): - dtype = dtypes.get(key) - packer = SequencePacker(pack_size, pad_value, dtype=dtype) - output[key] = packer.pack(seqs) - return output diff --git a/pipeline/packing/__init__.py b/pipeline/packing/__init__.py new file mode 100644 index 0000000..b3a0a32 --- /dev/null +++ b/pipeline/packing/__init__.py @@ -0,0 +1,84 @@ +"""Sequence packing algorithms for LLM training data. + +Available packers: + - BfdPacker: Best-Fit Decreasing, samples never split (default) + - FfDPacker: First-Fit Decreasing, samples never split + - GreedyPacker: First-fit in input order, samples never split +""" + +from typing import Dict, List, Optional, Union + +import torch + +from pipeline.packing.base import BasePacker +from pipeline.packing.binpack import GreedyPacker, FfDPacker, BfdPacker + + +def pack_tensors( + tensors: Dict[str, List[torch.Tensor]], + pack_size: int, + pad_value: Union[int, bool] = 0, + dtypes: Optional[Dict[str, torch.dtype]] = None, + pad_values: Optional[Dict[str, Union[int, bool]]] = None, + algo: Optional[Union[str, BasePacker]] = None, +) -> Dict[str, List[torch.Tensor]]: + """Pack multiple named tensor groups in parallel. + + Each group is packed independently with its own packer instance. + + Args: + tensors: Dict mapping key names to lists of 1D tensors. + pack_size: Fixed chunk length. + pad_value: Default padding value, used for keys not in pad_values. + dtypes: Optional per-key dtype declarations. + pad_values: Optional per-key padding values (e.g. pad_token_id for + 'sequence', False for 'loss_mask', 0 for 'position_ids'). + algo: Packing algorithm to use. Can be 'bfd' (default), + 'ffd', 'greedy', or a BasePacker instance. + + Returns: + Dict mapping key names to lists of packed tensors. + """ + if dtypes is None: + dtypes = {} + if pad_values is None: + pad_values = {} + + output: Dict[str, List[torch.Tensor]] = {} + for key, seqs in tensors.items(): + key_pad = pad_values.get(key, pad_value) + actual_packer = _resolve_algo(algo, pack_size, key_pad) + dtype = dtypes.get(key) + if dtype is not None: + actual_packer.dtype = dtype + output[key] = actual_packer.pack(seqs) + return output + + +def _resolve_algo( + algo: Optional[Union[str, BasePacker]], + pack_size: int, + pad_value: Union[int, bool], +) -> BasePacker: + if algo is None or algo == "bfd": + return BfdPacker(pack_size, pad_value) + if isinstance(algo, BasePacker): + cls = type(algo) + return cls(pack_size, pad_value) + if algo == "ffd": + return FfDPacker(pack_size, pad_value) + if algo == "greedy": + return GreedyPacker(pack_size, pad_value) + raise ValueError( + f"Unknown packing algorithm: {algo}. " + f"Choose from: bfd, ffd, greedy" + ) + + +__all__ = [ + "BasePacker", + "BfdPacker", + "FfDPacker", + "GreedyPacker", + "pack_tensors", +] diff --git a/pipeline/packing/base.py b/pipeline/packing/base.py new file mode 100644 index 0000000..27ed628 --- /dev/null +++ b/pipeline/packing/base.py @@ -0,0 +1,49 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Union + +import torch +from torch import Tensor + + +class BasePacker(ABC): + """Abstract base class for sequence packing algorithms. + + All packers must implement pack() and reset(). + pack() takes a list of 1D tensors and returns a list of packed fixed-size tensors. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + self.pack_size = pack_size + self.pad_value = pad_value + self.dtype = dtype + + @abstractmethod + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + """Pack sequences into fixed-size chunks.""" + ... + + @abstractmethod + def reset(self) -> None: + """Reset packer state for instance reuse.""" + ... + + def _validate_and_normalize(self, sequences: List[Tensor]) -> List[Tensor]: + """Validate 1D tensors and unify dtype.""" + if self.dtype is None and sequences: + self.dtype = sequences[0].dtype + + 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 != self.dtype: + seq = seq.to(self.dtype) + normalized.append(seq) + return normalized diff --git a/pipeline/packing/binpack.py b/pipeline/packing/binpack.py new file mode 100644 index 0000000..20b7d12 --- /dev/null +++ b/pipeline/packing/binpack.py @@ -0,0 +1,174 @@ +from typing import List, Optional, Union + +import torch +from torch import Tensor + +from pipeline.packing.base import BasePacker +from pipeline.utils import error_handler + + +def _truncate(tokens: List, max_len: int) -> List: + return tokens[:max_len] + + +def _pad_bin(bin_list: List, target_len: int, pad_value: Union[int, bool], dtype: torch.dtype) -> Tensor: + bin_list.extend([pad_value] * (target_len - len(bin_list))) + return torch.tensor(bin_list, dtype=dtype) + + +class GreedyPacker(BasePacker): + """Greedy first-fit packer (no sorting). + + Sequences are packed in input order into the first bin with enough space. + Overlong sequences (> pack_size) are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + for seq in normalized: + seq_len = int(seq.shape[0]) + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class FfDPacker(BasePacker): + """First-Fit Decreasing (FFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the first + bin with enough space. Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + placed = False + for bin_list in self._bins: + if len(bin_list) + seq_len <= pack_size: + bin_list.extend(seq.tolist()) + placed = True + break + if not placed: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages + + +class BfdPacker(BasePacker): + """Best-Fit Decreasing (BFD) bin-packing packer. + + Sequences are sorted by descending length, then packed into the bin + that minimizes remaining space (tightest fit). + Overlong sequences are truncated to pack_size. + """ + + def __init__( + self, + pack_size: int, + pad_value: Union[int, bool] = 0, + dtype: Optional[torch.dtype] = None, + ): + super().__init__(pack_size, pad_value, dtype) + self._bins: List[List] = [] + + def reset(self) -> None: + self._bins = [] + + @error_handler() + def pack(self, sequences: List[Tensor]) -> List[Tensor]: + if not sequences: + return [] + + normalized = self._validate_and_normalize(sequences) + self._bins = [] + pack_size = self.pack_size + pad_value = self.pad_value + + indexed = [(int(s.shape[0]), s) for s in normalized] + indexed.sort(key=lambda x: x[0], reverse=True) + + for seq_len, seq in indexed: + if seq_len > pack_size: + self._bins.append(_truncate(seq.tolist(), pack_size)) + continue + best_idx = -1 + best_remain = pack_size + 1 + for i, bin_list in enumerate(self._bins): + remain = pack_size - len(bin_list) + if seq_len <= remain < best_remain: + best_remain = remain + best_idx = i + if best_idx >= 0: + self._bins[best_idx].extend(seq.tolist()) + else: + self._bins.append(list(seq.tolist())) + + packages: List[Tensor] = [] + for bin_list in self._bins: + packages.append(_pad_bin(bin_list, pack_size, pad_value, self.dtype)) + + return packages diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index b4f1493..b1cbb3f 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -20,7 +20,7 @@ class ChatMLStrategy(PromptStrategy): assistant_end: str = "<｜im▁end｜>", ): super().__init__(tokenizer) - nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] + nl_id = tokenizer.token_to_id("\n") self._user_start_ids = self._encode_format(user_start) + [nl_id] self._user_end_ids = self._encode_format(user_end) + [nl_id] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py index 83f0f63..060adbe 100644 --- a/pipeline/tokenize/tokenizer.py +++ b/pipeline/tokenize/tokenizer.py @@ -266,6 +266,12 @@ class AutoTokenizer: return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens) + def token_to_id(self, token: str) -> Optional[int]: + """Convert a token string to its integer ID.""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not initialized.") + return self._tokenizer.token_to_id(token) + def __len__(self) -> int: if self._tokenizer is None: return 0 diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 8dc8899..47b7ef9 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -43,6 +43,13 @@ def main(): default=None, help="Prompt strategy: chatml, alpaca (default: chatml)", ) + parser.add_argument( + "-a", + "--pack-algo", + default=None, + choices=[None, "bfd", "ffd", "greedy"], + help="Packing algorithm: bfd (default), ffd, greedy", + ) parser.add_argument( "-p", "--pack-size", @@ -51,7 +58,14 @@ def main(): help="Pack size, <=0 to disable (default: -1)", ) parser.add_argument( - "--pad-value", type=int, default=0, help="Padding value (default: 0)" + "--pad-value", type=int, default=2, help="Padding token ID (default: 2 = <|pad|>)" + ) + parser.add_argument( + "-g", + "--group-size", + type=int, + default=1_000, + help="Merge every N packed chunks into one tensor, <=0 to disable (default: 1000)", ) parser.add_argument( "--log-level", @@ -95,9 +109,14 @@ def main(): print(f"\nStart caching...") if args.pack_size > 0: - print(f" pack_size={args.pack_size}, pad_value={args.pad_value}") + algo = args.pack_algo or "bfd" + print(f" pack_size={args.pack_size}, pad_value={args.pad_value}, algo={algo}") else: print(f" no packing") + if args.group_size > 0: + print(f" group_size={args.group_size} chunks per tensor") + else: + print(f" no grouping") cache_jsonl( files=jsonl_files, @@ -105,6 +124,8 @@ def main(): processor=processor, pack_size=args.pack_size, pad_value=args.pad_value, + group_size=args.group_size, + pack_algo=args.pack_algo, ) print(f"\nDone! Output saved to {output_dir}") diff --git a/tests/test_cache.py b/tests/test_cache.py index c2c8585..c11e23b 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -120,4 +120,4 @@ class TestCacheJsonl: pack_size=-1, pad_value=0, ) - assert len(output_files) == 1 + assert len(output_files) == 0 diff --git a/tests/test_io.py b/tests/test_io.py index a496722..b07099d 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -145,6 +145,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + def apply_chat_template( self, messages, add_generation_prompt=True, tokenize=True ): @@ -168,7 +171,7 @@ class TestPositionIds: processor = SFTProcessor(DummyTokenizer()) out_dir = os.path.join(tmpdir, "cached") - cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1) + cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1, group_size=0) h5_path = os.path.join(out_dir, "data.h5") loaded = HDF5Handler.load(h5_path, share_memory=False) diff --git a/tests/test_packing.py b/tests/test_packing.py index 5390d2d..89810fc 100644 --- a/tests/test_packing.py +++ b/tests/test_packing.py @@ -2,12 +2,147 @@ import pytest import torch -from pipeline.packing import SequencePacker +from pipeline.packing import ( + GreedyPacker, + FfDPacker, + BfdPacker, + pack_tensors, +) -class TestSequencePacker: - def test_normal_packing(self): - packer = SequencePacker(pack_size=10, pad_value=0) +class TestBfdPacker: + def test_best_fit_tight(self): + packer = BfdPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([1, 2, 3, 4], dtype=torch.int32), + torch.tensor([5, 6, 7, 8], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 5, 6] + + def test_different_dtypes(self): + for dtype in [torch.int32, torch.int64, torch.float32]: + packer = BfdPacker(pack_size=10, dtype=dtype) + val = 1.0 if dtype == torch.float32 else 1 + packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) + assert packages[0].dtype == dtype + + def test_dtype_conversion_on_mismatch(self): + packer = BfdPacker(pack_size=10, dtype=torch.int32) + packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) + assert packages[0].dtype == torch.int32 + assert packages[0][:3].tolist() == [1, 2, 3] + + def test_non_1d_tensor_raises_error(self): + packer = BfdPacker(pack_size=10) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor([[1, 2], [3, 4]])]) + with pytest.raises(ValueError, match="Expected 1D tensor"): + packer.pack([torch.tensor(5)]) + + def test_empty_input(self): + packer = BfdPacker(pack_size=10) + assert packer.pack([]) == [] + + def test_reset(self): + packer = BfdPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = BfdPacker(pack_size=6, pad_value=-1) + packages = packer.pack( + [ + torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9], dtype=torch.int32), + ] + ) + assert len(packages) == 2 + assert packages[0].tolist() == [1, 2, 3, 4, 5, 6] + assert packages[1].tolist() == [8, 9, -1, -1, -1, -1] + + def test_uses_two_bins_when_needed(self): + packer = BfdPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([8, 9, 10], dtype=torch.int32), + torch.tensor([11, 12, 13, 14, 15, 16], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 2 + for pkg in packages: + assert pkg.shape == (10,) + + def test_minimizes_waste_vs_ffd(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 5, 5, 7, 2, 4, 1, 4, 6, 2] + ] + bfd = BfdPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(bfd.pack(sequences)) <= len(ffd.pack(sequences)) + + +class TestFfDPacker: + def test_fills_tightly(self): + packer = FfDPacker(pack_size=10, pad_value=0) + sequences = [ + torch.tensor([7, 8], dtype=torch.int32), + torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + + def test_overlong_sample_truncated(self): + packer = FfDPacker(pack_size=5, pad_value=0) + packages = packer.pack( + [torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)] + ) + assert len(packages) == 1 + assert packages[0].tolist() == [1, 2, 3, 4, 5] + + def test_sort_descending_order(self): + packer = FfDPacker(pack_size=10, pad_value=-1) + sequences = [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.int32), + torch.tensor([9, 10], dtype=torch.int32), + ] + packages = packer.pack(sequences) + assert len(packages) == 1 + assert packages[0].tolist() == [3, 4, 5, 6, 7, 8, 1, 2, 9, 10] + + def test_reduces_bins_vs_greedy(self): + sequences = [ + torch.tensor([6] * i, dtype=torch.int32) + for i in [3, 8, 2, 7, 1, 4, 5, 3, 2, 6] + ] + greedy = GreedyPacker(pack_size=10, pad_value=0) + ffd = FfDPacker(pack_size=10, pad_value=0) + assert len(ffd.pack(sequences)) <= len(greedy.pack(sequences)) + + def test_reset(self): + packer = FfDPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 + packer.reset() + assert len(packer._bins) == 0 + + def test_empty_input(self): + packer = FfDPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestGreedyPacker: + def test_basic_packing(self): + packer = GreedyPacker(pack_size=10, pad_value=0) sequences = [ torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([4, 5], dtype=torch.int32), @@ -15,154 +150,186 @@ class TestSequencePacker: ] packages = packer.pack(sequences) assert len(packages) == 1 - for pkg in packages: - assert pkg.shape == (10,) - - # Verify all original values are present in order + assert packages[0].shape == (10,) assert packages[0][:9].tolist() == [1, 2, 3, 4, 5, 6, 7, 8, 9] - assert packages[0][9] == 0 # padding + assert packages[0][9] == 0 - def test_empty_list_input(self): - packer = SequencePacker(pack_size=10) - assert packer.pack([]) == [] - - def test_single_sequence_input(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) - assert len(packages) == 1 - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [-1] * 7 - - 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) + def test_overlong_sample_truncated(self): + """Overlong sample is truncated to pack_size.""" + packer = GreedyPacker(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) == 2 + assert len(packages) == 1 assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert packages[1].tolist() == [6, 7, 8, 0, 0] - def test_padding_value(self): - packer = SequencePacker(pack_size=8, pad_value=99) - packages = packer.pack( - [ - torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([3], dtype=torch.int32), - ] - ) - assert packages[0][:3].tolist() == [1, 2, 3] - assert packages[0][3:].tolist() == [99] * 5 - - def test_different_dtypes(self): - for dtype in [torch.int32, torch.int64, torch.float32]: - packer = SequencePacker(pack_size=10, dtype=dtype) - val = 1.0 if dtype == torch.float32 else 1 - packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)]) - assert packages[0].dtype == dtype - - def test_dtype_conversion_on_mismatch(self, caplog): - """Tensors with mismatched dtype are silently converted.""" - packer = SequencePacker(pack_size=10, dtype=torch.int32) - packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)]) - assert packages[0].dtype == torch.int32 - assert packages[0][:3].tolist() == [1, 2, 3] - - def test_non_1d_tensor_raises_error(self): - packer = SequencePacker(pack_size=10) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor([[1, 2], [3, 4]])]) - with pytest.raises(ValueError, match="Expected 1D tensor"): - packer.pack([torch.tensor(5)]) - - def test_input_list_not_modified(self): - packer = SequencePacker(pack_size=10) - original = [ - torch.tensor([3], dtype=torch.int32), + def test_multiple_fill(self): + packer = GreedyPacker(pack_size=6, pad_value=0) + sequences = [ torch.tensor([1, 2], dtype=torch.int32), - torch.tensor([4, 5, 6, 7], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5, 6], dtype=torch.int32), + torch.tensor([7], dtype=torch.int32), ] - original_repr = [seq.tolist() for seq in original] - packer.pack(original) - assert [seq.tolist() for seq in original] == original_repr - - def test_exact_pack_size_fit(self): - packer = SequencePacker(pack_size=5, pad_value=0) - packages = packer.pack( - [ - torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32), - torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32), - ] - ) - assert len(packages) == 2 - assert packages[0].tolist() == [1, 2, 3, 4, 5] - assert packages[1].tolist() == [6, 7, 8, 9, 10] - - def test_multiple_packs_full_utilization(self): - packer = SequencePacker(pack_size=10, pad_value=-1) - sequences = [torch.tensor([i], dtype=torch.int32) for i in range(1, 12)] packages = packer.pack(sequences) assert len(packages) == 2 - assert packages[0].tolist() == list(range(1, 11)) - assert packages[1].tolist() == [11] + [-1] * 9 + for pkg in packages: + assert pkg.shape == (6,) - def test_cross_group_ordering(self): - """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([False, False, True], dtype=torch.bool), - torch.tensor([False], dtype=torch.bool), - torch.tensor([False, False, False, True], dtype=torch.bool), - ] - 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 + def test_reset(self): + packer = GreedyPacker(pack_size=10) + packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)]) + assert len(packer._bins) == 1 packer.reset() - assert len(packer._packages) == 0 - assert packer._pos == 0 - assert packer._buffer == [] + assert len(packer._bins) == 0 - 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), - ] + def test_empty_input(self): + packer = GreedyPacker(pack_size=10) + assert packer.pack([]) == [] + + +class TestPackTensors: + def test_default_is_bfd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, ) - assert packages[0].tolist() == [1, 2, 3, 4] - assert packages[1].tolist() == [5, 6, 7, -1] + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_greedy(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2, 3], dtype=torch.int32), + torch.tensor([4, 5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="greedy", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_ffd(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1], dtype=torch.int32), + torch.tensor([2, 3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="ffd", + ) + assert result["input_ids"][0].tolist() == [2, 3, 4, 1, 5] + + def test_bfd_explicit(self): + result = pack_tensors( + tensors={ + "input_ids": [ + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([3, 4], dtype=torch.int32), + torch.tensor([5], dtype=torch.int32), + ], + }, + pack_size=5, + pad_value=0, + algo="bfd", + ) + assert result["input_ids"][0].tolist() == [1, 2, 3, 4, 5] + + def test_unknown_algo_raises(self): + with pytest.raises(ValueError, match="Unknown packing algorithm"): + pack_tensors( + tensors={"input_ids": [torch.tensor([1, 2, 3])]}, + pack_size=10, + pad_value=0, + algo="unknown_algo", + ) + + +class TestPositionIdsPacking: + """Verify position_ids reset to zero at sample boundaries after packing.""" + + def test_position_ids_reset_in_packed_chunk(self): + """After packing multiple SFT samples, position_ids restart from 0 at each boundary.""" + seqs = [ + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32), # len=5 + torch.tensor([0, 1, 2], dtype=torch.int32), # len=3 + torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32), # len=7 + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + packed = result["position_ids"][0].tolist() + assert packed == [0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3, 4, 5, 6, -1] + + def test_position_ids_reset_with_bfd(self): + """BFD may reorder, but each sample's position_ids still start from 0.""" + seqs = [ + torch.tensor([0, 1, 2], dtype=torch.int32), + torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32), + torch.tensor([0, 1, 2, 3], dtype=torch.int32), + ] + result = pack_tensors( + tensors={"position_ids": seqs}, + pack_size=16, + pad_value=-1, + algo="bfd", + ) + packed = result["position_ids"][0].tolist() + assert packed[0] == 0 + zeros = [i for i, v in enumerate(packed) if v == 0 and (i == 0 or packed[i - 1] != 0)] + assert len(zeros) == 3 + + def test_multiple_keys_share_same_boundaries(self): + """sequence, loss_mask, position_ids share identical chunk boundaries after packing.""" + seq_a = torch.tensor([101, 102, 103, 104], dtype=torch.int32) + seq_b = torch.tensor([201, 202, 203, 204, 205, 206, 207], dtype=torch.int32) + seq_c = torch.tensor([301, 302, 303, 304, 305], dtype=torch.int32) + + mask_a = torch.tensor([False, False, True, True], dtype=torch.bool) + mask_b = torch.tensor([False, False, False, False, True, True, True], dtype=torch.bool) + mask_c = torch.tensor([False, False, False, True, True], dtype=torch.bool) + + pos_a = torch.tensor([0, 1, 2, 3], dtype=torch.int32) + pos_b = torch.tensor([0, 1, 2, 3, 4, 5, 6], dtype=torch.int32) + pos_c = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32) + + result = pack_tensors( + tensors={ + "sequence": [seq_a, seq_b, seq_c], + "loss_mask": [mask_a, mask_b, mask_c], + "position_ids": [pos_a, pos_b, pos_c], + }, + pack_size=16, + pad_value=-1, + algo="greedy", + ) + + seq_chunk = result["sequence"][0] + mask_chunk = result["loss_mask"][0] + pos_chunk = result["position_ids"][0] + + assert len(seq_chunk) == len(mask_chunk) == len(pos_chunk) == 16 + + for i in range(16): + if seq_chunk[i] == -1: + assert mask_chunk[i] == -1 + assert pos_chunk[i] == -1 + + pos_ids = pos_chunk.tolist() + zeros = [i for i, v in enumerate(pos_ids) if v == 0] + assert len(zeros) == 3 diff --git a/tests/test_processors.py b/tests/test_processors.py index e10ba3d..bec2400 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -18,6 +18,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + def apply_chat_template( self, messages, add_generation_prompt=True, tokenize=True ): diff --git a/tests/test_strategies.py b/tests/test_strategies.py index d4600b5..2447024 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -13,6 +13,9 @@ class DummyTokenizer: def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + def token_to_id(self, token: str): + return ord(token) + class DummyStrategy(PromptStrategy): def __init__(self, tokenizer):