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 参数
This commit is contained in:
@@ -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",
|
||||
]
|
||||
Reference in New Issue
Block a user