refactor: 重构打包模块，新增 BFD/FFD/Greedy 三种 bin-packing 算法，默认 BFD

- 将 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

@@ -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",

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}")

pipeline/packing.py

@@ -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

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",
+]

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

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

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]

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

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}")
 

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

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)

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

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
     ):

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):