大规模训练的数据流水线：预处理、分片、预取的三层并行

大规模训练的数据流水线预处理、分片、预取的三层并行一、个性化深度引言训练一个 13B 参数的模型数据集 2TB。最初的 Pipeline 是单机处理读文件、Tokenize、Batching数据加载占了 GPU 训练时间的 40%。A100 显卡每秒钟等着 CPU 喂数据——计算资源浪费触目惊心。重写了 Data Pipeline预处理层在读写分离的存储集群上并行完成分片层按 rank 均匀分配数据预取层用多进程提前准备下一批。数据加载耗时从 40% 降到 3%。见证奇迹的时刻不是在 GPU 利用率从 55% 飙升到 92%。而是发现很多人花几个月调模型结构却从没认真审视过 Data Pipeline。在分布式训练中数据加载效率对总训练时间的贡献经常超过模型结构优化。二、个性化原理剖析大规模训练的数据流水线分为三层三层职责预处理层离线在训练启动前完成。Tokenization、数据清洗、格式标准化。目标是输出可直接加载的二进制数据文件。分片层训练启动时全局 Shuffle 保证跨 Epoch 的数据随机性。Rank 分配确保每个 GPU 拿到互不重复的数据。负载均衡让各 rank 处理的数据量一致。预取层运行时多进程并行读取和预处理。内存缓冲池提前准备 2-3 个 batch。Pinned memory 加速 CPU→GPU 的数据传输。三、个性化代码实践 大规模训练的三层 Data Pipeline 实现。 设计理念每一层的瓶颈都必须独立测量和优化。 import os import time import queue import threading import multiprocessing as mp from dataclasses import dataclass, field from pathlib import Path from typing import Iterator, List, Optional, Tuple import torch import numpy as np from torch.utils.data import Dataset, DataLoader, IterableDataset from torch.utils.data.distributed import DistributedSampler # ═══════════════════════════════════════════ # 第一层预处理 # ═══════════════════════════════════════════ class PreprocessingPipeline: 设计原因离线预处理将原始文本转为二进制 token 文件。 一次预处理多次训练复用——避免每次训练都重复 Tokenization。 def __init__( self, tokenizer, output_dir: Path, chunk_size: int 1024 * 1024, # 1M tokens per file ): self.tokenizer tokenizer self.output_dir output_dir self.chunk_size chunk_size self.output_dir.mkdir(parentsTrue, exist_okTrue) def process_files( self, input_files: List[Path], num_workers: int 8 ) - List[Path]: 设计原因多进程并行处理充分利用多核 CPU。 输出文件为 .bintoken ids .idx文件偏移索引。 # 设计原因将文件列表均匀分给各个 worker。 chunks np.array_split(input_files, num_workers) with mp.Pool(num_workers) as pool: results pool.map(self._process_chunk, chunks) # 设计原因汇总所有 worker 的输出文件路径。 output_files [] for file_list in results: output_files.extend(file_list) # 设计原因生成全局索引文件记录每个 .bin 的元信息。 self._write_global_index(output_files) return output_files def _process_chunk( self, files: List[Path] ) - List[Path]: 设计原因每个 worker 处理一个文件子集。 独立进程运行不共享状态避免 GIL 和锁竞争。 output_files [] chunk_tokens [] chunk_id 0 for file_path in files: # 设计原因流式读取大文件避免内存爆炸。 with open(file_path, r) as f: for line in f: line line.strip() if not line: continue # 设计原因Tokenization 后直接 append。 tokens self.tokenizer.encode(line) chunk_tokens.extend(tokens) # 设计原因攒够 chunk_size 就写入文件。 if len(chunk_tokens) self.chunk_size: output_file self._write_chunk( chunk_tokens, chunk_id ) output_files.append(output_file) chunk_tokens [] chunk_id 1 # 设计原因最后一个不满的 chunk 也写入。 if chunk_tokens: output_file self._write_chunk(chunk_tokens, chunk_id) output_files.append(output_file) return output_files def _write_chunk( self, tokens: List[int], chunk_id: int ) - Path: 设计原因写入 .bin 文件uint16 紧凑存储。 使用 uint16 而非 int32节省 50% 存储空间。 大多数 tokenizer 的 vocab_size 65536uint16 足够。 output_path self.output_dir / fchunk_{chunk_id:06d}.bin np.array(tokens, dtypenp.uint16).tofile(output_path) return output_path staticmethod def _write_global_index(files: List[Path]) - None: 设计原因全局索引记录每个 chunk 的元信息。 pass # ═══════════════════════════════════════════ # 第二层分片 第三层预取 # ═══════════════════════════════════════════ class ShardedDataset(IterableDataset): 设计原因IterableDataset 支持流式读取大文件。 每个 rank 取互不重叠的数据切片。 def __init__( self, data_dir: Path, rank: int, world_size: int, seq_length: int 2048, ): self.data_dir data_dir self.rank rank self.world_size world_size self.seq_length seq_length # 设计原因获取所有 .bin 文件按 rank 均匀分配。 self.bin_files sorted(self.data_dir.glob(*.bin)) # 设计原因只处理分配给当前 rank 的文件。 self.rank_files [ f for i, f in enumerate(self.bin_files) if i % world_size rank ] def __iter__(self) - Iterator[torch.Tensor]: 设计原因迭代器按 rank_files 循环读取数据。 每个 epoch 随机打乱文件顺序shuffle files, not tokens。 while True: # 设计原因每个 epoch 打乱文件顺序 # 全局 shuffle 保证不同 epoch 的数据随机性。 np.random.shuffle(self.rank_files) for file_path in self.rank_files: # 设计原因使用 np.memmap 做零拷贝读取 # 数据直接映射到内存不需要复制。 tokens np.memmap( file_path, dtypenp.uint16, moder ) # 设计原因滑动窗口生成训练样本。 for i in range( 0, len(tokens) - self.seq_length, self.seq_length ): yield torch.from_numpy( tokens[i:i self.seq_length].astype(np.int64) ) class PrefetchDataLoader: 设计原因预取层——在 GPU 计算当前 batch 时 后台线程已经在准备下几个 batch。 def __init__( self, dataset: IterableDataset, batch_size: int, prefetch_factor: int 3, num_workers: int 4, ): self.dataset dataset self.batch_size batch_size self.prefetch_factor prefetch_factor # 设计原因预取队列缓冲 prefetch_factor 个 batch。 self._prefetch_queue queue.Queue( maxsizeprefetch_factor ) # 设计原因后台线程持续从 dataset 取数据放入缓冲。 self._loader_thread: Optional[threading.Thread] None self._stop_flag threading.Event() def start(self) - None: 设计原因启动后台预取线程。 self._loader_thread threading.Thread( targetself._prefetch_worker, daemonTrue ) self._loader_thread.start() def stop(self) - None: 设计原因停止预取线程。 self._stop_flag.set() if self._loader_thread: self._loader_thread.join(timeout5) def _prefetch_worker(self) - None: 设计原因后台工作线程。 持续从 dataset 取数据、组成 batch、放入缓冲队列。 iterator iter(self.dataset) batch [] while not self._stop_flag.is_set(): try: sample next(iterator) batch.append(sample) if len(batch) self.batch_size: # 设计原因组成 batch 后堆叠为 tensor。 stacked torch.stack(batch) # 设计原因使用 pinned memory 加速 GPU 传输。 pinned stacked.pin_memory() # 设计原因put 到缓冲队列如果队列满了会阻塞。 self._prefetch_queue.put(pinned) batch [] except StopIteration: iterator iter(self.dataset) def next_batch(self) - torch.Tensor: 设计原因训练循环中调用从预取队列获取 batch。 如果队列不为空立即返回否则阻塞等待。 return self._prefetch_queue.get() def __len__(self) - int: 设计原因返回当前缓冲队列中的 batch 数量。 return self._prefetch_queue.qsize() def create_distributed_dataloader( data_dir: Path, rank: int, world_size: int, batch_size: int, seq_length: int 2048, num_workers: int 4, prefetch_factor: int 3, ) - PrefetchDataLoader: 设计原因工厂函数一键创建分布式 DataLoader。 封装了分片 预取的完整配置。 # 第二层分片 dataset ShardedDataset( data_dirdata_dir, rankrank, world_sizeworld_size, seq_lengthseq_length, ) # 第三层预取 loader PrefetchDataLoader( datasetdataset, batch_sizebatch_size, prefetch_factorprefetch_factor, num_workersnum_workers, ) loader.start() return loader四、个性化边界权衡1. 离线预处理 vs 在线处理离线预处理一次处理可反复复用适合固定数据集。在线处理灵活适应动态数据集但增加训练时的 CPU 开销。建议数据稳定时离线处理数据持续增长时混合方案增量在线处理 定期离线归档。2. 文件大小大文件 vs 小文件大文件1GB减少文件打开次数但 Shuffle 粒度粗糙。小文件100MBShuffle 灵活但产生大量文件句柄。推荐1-5GB 的 chunk 文件单文件约 100 万条样本。3. Prefetch 数量1 vs NPrefetch 1 个 batch 内存开销最小但 GPU 可能短暂空转。Prefetch 3-5 个 batch 消除空转但多占用batch_size × prefetchN的 CPU 内存。推荐GPU 内存充足时 prefetch 3显存紧张时 prefetch 1。4. Shuffle 策略全局随机 vs 文件内随机全局随机所有样本完全打乱最优但需要大量内存全量索引。文件间随机 文件内顺序读取内存开销小但随机性降低。推荐文件间随机 文件内顺序读取内存可控且随机性够用。5. 数据校验训练前 vs 运行时训练前全量校验保证数据质量但 2TB 数据集全量扫描耗时数小时。运行时抽样校验轻量但可能遗漏问题。建议训练前做语法和格式校验运行时做长度和分布监控。五、总结大规模训练的数据流水线通过预处理、分片、预取三层并行架构将数据加载开销从训练时间的 40% 压缩到 3%。预处理层将原始文本离线转为紧凑二进制格式分片层按 world_size 均匀分配数据预取层用多进程 缓冲队列消除 GPU 空转。工程实践中需要在文件大小、Prefetch 数量、Shuffle 策略上做权衡核心原则是每一层的吞吐量必须大于训练吞吐量任何一层的瓶颈都会拖慢整体。