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,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
|
||||
Reference in New Issue
Block a user