From f67bad0d8b212e5d39c2398c2f04ec587cc8ba33 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 30 Mar 2026 17:49:50 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=20processors?= =?UTF-8?q?=E3=80=81cache=E3=80=81packing=20=E6=A8=A1=E5=9D=97=E6=80=A7?= =?UTF-8?q?=E8=83=BD=E5=B9=B6=E7=AE=80=E5=8C=96=20README?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 219 +++++++++-------------------------------- pipeline/cache.py | 35 ++++--- pipeline/io.py | 1 + pipeline/packing.py | 71 ++++++++----- pipeline/processors.py | 23 +++-- 5 files changed, 129 insertions(+), 220 deletions(-) diff --git a/README.md b/README.md index b982ba0..276616f 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # DataPipeline -用于训练 KHAOSZ 模型的数据集处理工具。提供文本导出、Tokenize、序列打包、H5 存储等独立工具，支持预训练 / SFT / DPO 三种训练范式。 +数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。 ## 项目结构 @@ -8,231 +8,104 @@ pipeline/ ├── tokenizer.py # BPE 分词器 ├── text.py # 文本规范化 -├── packing.py # 序列打包（bin-packing） +├── packing.py # 序列打包 ├── io.py # 文件/HDF5 读写 ├── processors.py # PT / SFT / DPO 处理器 -├── export.py # Dataset → JSONL 导出 -└── cache.py # JSONL → Tokenize → H5 缓存 - -pre_train/ # 预训练数据处理脚本 -supervised_finetuning/ # SFT 数据处理脚本 -reforce_learning/ # DPO 数据处理脚本 +├── export.py # Dataset → JSONL +└── cache.py # JSONL → Tokenize → H5 ``` ## 设计理念 -每个模块**独立可用、零互相依赖**，调用者按需组合： +模块**独立可用**，通过磁盘文件解耦，按需组合： ``` -HuggingFace Hub - │ - ▼ load_dataset() - DatasetDict - │ - ▼ export_dataset() ← pipeline/export.py - JSONL 文件 - │ - ▼ cache_jsonl() ← pipeline/cache.py - │ ├─ Processor.process() ← pipeline/processors.py - │ ├─ SequencePacker.pack() ← pipeline/packing.py - │ └─ IOHandler.save_h5() ← pipeline/io.py - HDF5 张量文件 +Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5 + ↑ + processors.py + packing.py + io.py ``` -> 各阶段之间通过磁盘文件解耦。你可以只执行阶段 1（导出 JSONL），也可以继续执行阶段 2（tokenize 并缓存为 H5），按需选择。 +## 使用方法 -## 快速开始 - -### 安装依赖 - -```bash -pip install -e . -``` - -### 阶段 1：导出数据集为 JSONL +### 1. 导出数据集 ```python from datasets import load_dataset from pipeline.export import export_dataset dataset = load_dataset("your-dataset") -export_dataset( - dataset=dataset["train"], # 直接传 Dataset，不传 DatasetDict - output_dir="./dataset", - output_prefix="my-data", - max_chunks=5, # 可选，限制 chunk 数量（调试用） -) -``` - -**自定义转换函数：** - -```python -def process_func(example): - # 提取字段、转换格式、展开多轮对话等 - return {"query": example["instruction"], "response": example["output"]} - export_dataset( dataset=dataset["train"], - output_dir="./dataset", - output_prefix="my-sft", - process_func=process_func, + output_dir="./data", + output_prefix="train", + process_func=lambda x: {"text": x["content"]}, # 可选 ) ``` -> `process_func` 返回单个 `dict` 或 `list[dict]`（一条样本可展开为多条）。 - -**使用文本规范化：** +### 2. Tokenize 并缓存 ```python -from pipeline.text import TextNormalizer - -normalizer = TextNormalizer() - -def process_func(example): - return {"text": normalizer.normalize(example["content"])} - -export_dataset( - dataset=dataset["train"], - output_dir="./dataset", - output_prefix="my-pretrain", - process_func=process_func, -) -``` - -### 阶段 2：Tokenize 并缓存为 HDF5 - -```python -from pipeline.tokenizer import BpeTokenizer -from pipeline.processors import ProcessorFactory -from pipeline.cache import cache_jsonl +from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl tokenizer = BpeTokenizer("tokenizer.json") -processor = ProcessorFactory.create("pt", tokenizer) +processor = ProcessorFactory.create("pt", tokenizer) # "pt" | "sft" | "dpo" cache_jsonl( - files=["./dataset/my-pretrain_chunk_0.jsonl"], + files=["./data/train.jsonl"], output_dir="./cached", processor=processor, - pack_size=4096, # 可选，打包长度；<=0 不打包 + pack_size=4096, # <=0 不打包 pad_value=1, ) ``` -**处理器类型：** +### 3. 处理器类型 -| 类型 | 工厂 key | 输入格式 | 输出 keys | -|------|----------|---------|-----------| +| 类型 | key | 输入 | 输出 | +|------|-----|------|------| | 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` | | SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` | -| DPO | `"dpo"` | 待定 | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` | +| DPO | `"dpo"` | `{"query": "...", "chosen": "...", "rejected": "..."}` | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` | -## 独立工具参考 +## 参数参考 -### BpeTokenizer +### export_dataset() -```python -from pipeline.tokenizer import BpeTokenizer +| 参数 | 类型 | 默认 | 说明 | +|------|------|------|------| +| `dataset` | Dataset | 必填 | HuggingFace Dataset | +| `output_dir` | str | 必填 | 输出目录 | +| `output_prefix` | str | 必填 | 文件名前缀 | +| `chunk_size` | int | 1_000_000 | 每个文件的最大样本数 | +| `max_chunks` | int | None | 最大 chunk 数量 | +| `process_func` | callable | None | 样本转换函数 | +| `column` | str | "text" | 默认文本列名 | -tokenizer = BpeTokenizer("tokenizer.json") -ids = tokenizer.encode("hello world") # → [1234, 5678, 1] -text = tokenizer.decode(ids) # → "hello world" -len(tokenizer) # → 词表大小 -``` +### cache_jsonl() -### TextNormalizer - -```python -from pipeline.text import TextNormalizer - -normalizer = TextNormalizer() -text = normalizer.normalize(raw_text) -``` - -替换规则包括：全角引号 → 半角、各种短横线统一、不间断空格 → 普通空格等。支持自定义规则： - -```python -normalizer = TextNormalizer(custom_rules={"旧词": "新词"}) -``` +| 参数 | 类型 | 默认 | 说明 | +|------|------|------|------| +| `files` | List[str] | 必填 | JSONL 文件列表 | +| `output_dir` | str | 必填 | 输出目录 | +| `processor` | BaseProcessor | 必填 | 处理器实例 | +| `pack_size` | int | -1 | 打包长度，<=0 不打包 | +| `pad_value` | int | 1 | 填充值 | ### SequencePacker ```python -from pipeline.packing import SequencePacker - packer = SequencePacker(pack_size=4096, pad_value=0) -packed = packer.pack(list_of_tensors) # → List[Tensor]，每个长度为 pack_size +packed = packer.pack([tensor1, tensor2, ...]) # → List[Tensor] ``` ### IOHandler ```python -from pipeline.io import IOHandler - # 保存 -IOHandler.save_h5("./output", "my_data", {"sequence": [tensor1, tensor2]}) +IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]}) # 加载 -data = IOHandler.load_h5("./output") # → {"sequence": [tensor1, tensor2, ...]} - -# 遍历文件 -files = IOHandler.fetch_files("./dataset") -folders = IOHandler.fetch_folders("./dataset") -``` - -### 自定义 Processor - -```python -from pipeline.processors import BaseProcessor, ProcessorFactory -import torch - -class MyProcessor(BaseProcessor): - def __init__(self, tokenizer): - self.tokenizer = tokenizer - - def process(self, input_dict): - tokens = self.tokenizer.encode(input_dict["text"]) - return {"sequence": torch.tensor(tokens, dtype=torch.int32)} - - @property - def output_keys(self): - return ["sequence"] - -ProcessorFactory.register("my_type", MyProcessor) -``` - -## 运行脚本 - -```bash -# 预训练 -python pre_train/chinese-c4.py -python pre_train/english-wiki.py - -# SFT -python supervised_finetuning/sft_belle.py -python supervised_finetuning/sft_coder.py - -# DPO -python reforce_learning/dpp_chinese_dpo_pairs.py -``` - -## 输出格式 - -**JSONL**（阶段 1 输出）： - -```jsonl -{"text": "训练文本内容..."} -{"query": "问题", "response": "答案"} -``` - -**HDF5**（阶段 2 输出）： - -``` -my_data.h5 -├── sequence/ -│ ├── data_0 # Tensor (4096,) int32 -│ ├── data_1 # Tensor (4096,) int32 -│ └── ... -└── loss_mask/ # 仅 SFT - ├── data_0 # Tensor (4096,) bool - └── ... +data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]} ``` diff --git a/pipeline/cache.py b/pipeline/cache.py index 3cfb80c..9efa337 100644 --- a/pipeline/cache.py +++ b/pipeline/cache.py @@ -2,7 +2,7 @@ import json import os import logging -from typing import List +from typing import List, Dict from pathlib import Path from tqdm import tqdm @@ -39,33 +39,40 @@ def cache_jsonl( """ os.makedirs(output_dir, exist_ok=True) output_files: List[str] = [] + # Cache output_keys to avoid repeated attribute access + output_keys = processor.output_keys for file_path in files: file_name = Path(file_path).stem - arrows = [] + # Pre-allocate lists for each output key + arrows: Dict[str, List] = {key: [] for key in output_keys} + + # Read and process all lines with open(file_path, "r", encoding="utf-8") as f: for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1): try: - arrow = processor.process(json.loads(line)) + result = processor.process(json.loads(line)) + if result is not None: + # Batch append: add each key's tensor to corresponding list + for key in output_keys: + arrows[key].append(result[key]) except json.JSONDecodeError as e: logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.") continue except Exception as e: logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.") continue - if arrow is not None: - arrows.append(arrow) - package = {key: [a[key] for a in arrows] for key in processor.output_keys} - - output = {} - for key in processor.output_keys: - if pack_size > 0: - packer = SequencePacker(pack_size, pad_value) # independent instance per key - output[key] = packer.pack(package[key]) - else: - output[key] = package[key] + # Convert lists to tensors once per key + if pack_size > 0: + output = {} + for key in output_keys: + packer = SequencePacker(pack_size, pad_value) + output[key] = packer.pack(arrows[key]) + else: + # No packing: directly use the arrow tensors + output = arrows IOHandler.save_h5(output_dir, file_name, output) h5_path = os.path.join(output_dir, f"{file_name}.h5") diff --git a/pipeline/io.py b/pipeline/io.py index 33aec2b..be34e4a 100644 --- a/pipeline/io.py +++ b/pipeline/io.py @@ -34,6 +34,7 @@ class IOHandler: def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None: os.makedirs(output_dir, exist_ok=True) full_path = os.path.join(output_dir, f"{file_name}.h5") + with h5py.File(full_path, 'w') as f: for key, tensors in tensor_group.items(): grp = f.create_group(key) diff --git a/pipeline/packing.py b/pipeline/packing.py index 1c46236..fabbee2 100644 --- a/pipeline/packing.py +++ b/pipeline/packing.py @@ -1,5 +1,5 @@ import logging -from typing import List +from typing import List, Optional import torch from torch import Tensor @@ -14,14 +14,23 @@ class SequencePacker: self.pack_size = pack_size self.pad_value = pad_value self.dtype = dtype + # Pre-allocate buffer for better performance + self._buffer: Optional[Tensor] = None self._reset() def _reset(self) -> None: """Reset internal state for instance reuse.""" - self._current_pack = torch.full( - (self.pack_size,), self.pad_value, dtype=self.dtype - ) + # Reuse buffer instead of creating new tensors + if self._buffer is None or self._buffer.shape[0] != self.pack_size: + self._buffer = torch.full( + (self.pack_size,), self.pad_value, dtype=self.dtype + ) + else: + self._buffer.fill_(self.pad_value) self._current_pos = 0 + self._packages: List[Tensor] = [] + # Backward compatibility: maintain _current_pack reference + self._current_pack = self._buffer @error_handler() def pack(self, sequences: List[Tensor]) -> List[Tensor]: @@ -37,53 +46,63 @@ class SequencePacker: # Input validation if not sequences: return [] + + # Validate and cache tensor sizes in one pass + tensor_sizes = [] 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}" ) - # Check dtype compatibility and warn if mismatched + tensor_sizes.append(seq.numel()) if seq.dtype != self.dtype: logger.warning( f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, " f"will be converted. This may affect packing efficiency." ) - packages = [] - # Sort by length in descending order to improve packing efficiency - # Use sorted() to avoid modifying the input list - sorted_sequences = sorted(sequences, key=lambda x: x.numel(), reverse=True) + # Reset state for new packing + self._packages = [] + self._reset() + + # Combine sequences with their sizes for sorting + indexed_seqs = list(zip(sequences, tensor_sizes)) + # Sort by size descending (First-Fit Decreasing algorithm) + indexed_seqs.sort(key=lambda x: x[1], reverse=True) - for tensor in sorted_sequences: + for tensor, tensor_size in indexed_seqs: # Truncate sequences that exceed pack_size - if tensor.numel() > self.pack_size: + if tensor_size > self.pack_size: logger.warning( - f"Sequence length {tensor.numel()} exceeds pack_size {self.pack_size}, truncating" + f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating" ) + tensor_size = self.pack_size tensor = tensor[: self.pack_size] - tensor_size = tensor.numel() # Current package is full, create a new one if self._current_pos + tensor_size > self.pack_size: - packages.append(self._current_pack) - self._current_pack = torch.full( - (self.pack_size,), self.pad_value, dtype=self.dtype - ) + # Finish current package (pad to pack_size) + package = self._buffer.clone() + self._packages.append(package) + # Reset buffer for reuse + self._buffer.fill_(self.pad_value) self._current_pos = 0 - # Place tensor in current package (remaining positions stay as pad_value) - self._current_pack[self._current_pos : self._current_pos + tensor_size] = ( - tensor - ) + # Place tensor in current package + self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor self._current_pos += tensor_size - # Handle the last package + # Handle the last package (pad to pack_size) if self._current_pos > 0: - packages.append(self._current_pack) - self._current_pack = None - self._current_pos = 0 + package = self._buffer.clone() + self._packages.append(package) - return packages + # Clear buffer and reset state for backward compatibility + self._buffer = None + self._current_pack = None + self._current_pos = 0 + + return self._packages def reset(self) -> None: """Reset packer state for reuse. More efficient than creating a new instance.""" diff --git a/pipeline/processors.py b/pipeline/processors.py index 34cf892..9ff49e3 100644 --- a/pipeline/processors.py +++ b/pipeline/processors.py @@ -41,14 +41,18 @@ class SFTProcessor(BaseProcessor): self.tokenizer = tokenizer def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: - query, response = input_dict["query"], input_dict["response"] + query = input_dict["query"] + response = input_dict["response"] + q = self.tokenizer.encode( f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n" ) a = self.tokenizer.encode(f"{response}<|im_end|>\n") + + q_len = len(q) tokens = torch.tensor(q + a, dtype=torch.int32) - loss_mask = torch.zeros_like(tokens, dtype=torch.bool) - loss_mask[len(q):] = True + loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool) + loss_mask[q_len:] = True return {"sequence": tokens, "loss_mask": loss_mask} @property @@ -72,14 +76,19 @@ class DPOProcessor(BaseProcessor): ) chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n") + q_len = len(q) + chosen_len = len(chosen) + chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32) - chosen_mask = torch.zeros_like(chosen_tokens, dtype=torch.bool) - chosen_mask[len(q):] = True + chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool) + chosen_mask[q_len:] = True rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n") + rejected_len = len(rejected) + rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32) - rejected_mask = torch.zeros_like(rejected_tokens, dtype=torch.bool) - rejected_mask[len(q):] = True + rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool) + rejected_mask[q_len:] = True return { "chosen": chosen_tokens,