fix: 修复预处理流水线 4 个致命问题

- pipeline: 单条数据异常不再崩溃整条流水线, 改 log warning 后跳过
- pipeline: _align_bucket 统一用 len(ids) 填充, 修复多输出模式下长度错配
- writer: BinWriter/H5Writer 写入失败自动清理残留文件并记录详细错误
- packing: BFDPacking 真正将序列打包进 bin 而非仅重排, 减少碎片
This commit is contained in:
2026-06-18 17:38:01 +08:00
parent 376e9eba80
commit d88a41f8f1
3 changed files with 79 additions and 26 deletions
+36 -16
View File
@@ -6,8 +6,7 @@ pipeline later flattens the result into contiguous tensors.
"""
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Dict, List, Tuple
from typing import Dict, List
from astrai.factory import BaseFactory
@@ -53,6 +52,15 @@ class SimplePacking(PackingStrategy):
@PackingStrategyFactory.register("bfd")
class BFDPacking(PackingStrategy):
"""Best-Fit Decreasing bin packing.
Assigns sequences to bins using a best-fit heuristic (sorted by
decreasing length) and concatenates sequences within each bin into
a single packed sequence. Packed sequences are truncated to
*max_packed_len* so that each packed bin fits within one context
window during training.
"""
def apply(
self,
keys: Dict[str, List[List[int]]],
@@ -62,24 +70,40 @@ class BFDPacking(PackingStrategy):
sequences = keys.get("sequence", [])
if not sequences:
return keys
plan = self._plan(sequences, max_packed_len)
reordered: dict = defaultdict(list)
for orig_idx, _ in plan:
for k, vals in keys.items():
reordered[k].append(
_truncate(vals[orig_idx], max_packed_len, truncation_mode)
bins = self._plan(sequences, max_packed_len, truncation_mode)
packed: Dict[str, List[List[int]]] = {}
for k, vals in keys.items():
packed[k] = [
_truncate(
self._concat_bin(vals, bin_indices),
max_packed_len,
truncation_mode,
)
return dict(reordered)
for bin_indices in bins
]
return packed
@staticmethod
def _plan(sequences: List[List[int]], max_packed_len: int) -> List[Tuple[int, int]]:
def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
result: List[int] = []
for i in indices:
result.extend(vals[i])
return result
@staticmethod
def _plan(
sequences: List[List[int]], max_packed_len: int, truncation_mode: str
) -> List[List[int]]:
n = len(sequences)
order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
bins: List[List[int]] = []
bin_lengths: List[int] = []
for orig_idx in order:
seq_len = min(len(sequences[orig_idx]), max_packed_len)
seq_len = len(
_truncate(sequences[orig_idx], max_packed_len, truncation_mode)
)
best_bin = None
best_remain = max_packed_len + 1
for i, bl in enumerate(bin_lengths):
@@ -94,8 +118,4 @@ class BFDPacking(PackingStrategy):
bins.append([orig_idx])
bin_lengths.append(seq_len)
plan: List[Tuple[int, int]] = []
for bin_indices in bins:
for orig_idx in bin_indices:
plan.append((orig_idx, min(len(sequences[orig_idx]), max_packed_len)))
return plan
return bins