ViperEkura
598e1ce4ae
- 将 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 参数
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
"""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",
|
|
]
|