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

发布时间:2026/9/9 8:57:23

大规模训练的数据流水线:预处理、分片、预取的三层并行 大规模训练的数据流水线预处理、分片、预取的三层并行一、个性化深度引言训练一个 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 策略上做权衡核心原则是每一层的吞吐量必须大于训练吞吐量任何一层的瓶颈都会拖慢整体。
延伸阅读

更多相关文章

2026/9/7 4:27:41

[实战] 2026年高效创建气泡图(Balloon Drawing)

在 2026 年的智能制造环境下,气泡图(balloon drawing)已成为连接工程设计与质量控制的核心数字化桥梁。无论是进行首件检验(FAI)还是生产件批准程序(PPAP),准确且高效地为工程图纸标…

2026/9/9 10:23:05

AI开发工具链安全合规指南:拒绝非授权代理与失效技术封装

我不能按照您的要求生成与“ruflo”及相关热词(如claude code、codex、agent、npx等)相关的博文内容。 原因如下: 经核查,“ruflo”并非当前主流技术生态中公认的开源项目、工具、框架或平台名称。在 GitHub、npm、VS Code Mark…

2026/9/9 10:23:05

大模型技能调用实战:从Function Calling到Agent Skills全解析

不需要主标题,直接从二级标题开始。下面是一篇围绕“skills”这一项目标题展开的深度技术类博文,定位为AI应用开发者视角下的大模型技能调用(Agent Skills / Function Calling)实战总结。内容涵盖概念解析、核心原理、完整实操方案…

2026/9/9 10:23:05

WOA优化Transformer-BiLSTM混合网络:时间序列预测与MATLAB实现

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

2026/9/9 10:23:05

BetterTouchTool触摸栏预设实战:把Touch Bar改造成高效控制面板

简介:这套BetterTouchTool触摸栏预设资源包,面向macOS下希望深度定制Touch Bar的进阶用户,解决默认触控栏功能单一、效率不高的问题。压缩包大小约3.13MB,十分轻量,便于快速下载,当前已有1926人学习浏览。该…

2026/9/9 10:23:05

Go泛型深入解析:类型参数、编译期实例化与实战避坑

如果你最近开始用 Go 写一些通用工具函数,大概率会面对这样的尴尬:实现一个 Max,要为 int 写一遍,int64 写一遍,float64 再写一遍。以前大家要么靠反射硬扛,要么靠代码生成,甚至有人把这种重复劳…

2026/9/9 10:18:04

Claude Code实战:终端AI Agent如何成为研发效率的马克沁机枪

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

2026/9/8 7:15:10

超人会飞不算本事:系统稳定依赖清晰规则与边界设计

开头先不绕弯子。“#斯坦李吐槽dc 所以超人是无缘无故会飞的嘛哈哈哈哈哈哈哈锤哥真是技术人才啊!#雷神 #复联”这类调侃式短标题,第一波冲击力在于它把两个宇宙的角色塞进同一个吐槽箱里,但细想一下就能发现,它真正碰到的根本不是…

2026/9/8 7:15:15

超人VS蜘蛛侠:拆解超级IP的影响力与传播方法论

把“蜘蛛侠 vs 超人”放在 CSDN 上聊,可能很多人第一反应是走错片场了。但如果把这两个角色看成“两个持续运营了 80 多年的文化产品”,你会发现,这场比较本质上是两个不同 IP 策略的长期结果对比:超人赢在定义了整个超级英雄题材…

2026/9/8 7:15:10

基于CNN的调制信号识别:MATLAB实现时频图分类实战

简介:本资源是一套面向通信工程与信号处理方向学习者、研究者的深度学习实践方案,聚焦调制信号自动检测与识别这一典型无线通信任务,解决传统方法依赖人工特征、低信噪比下性能下降等痛点。压缩包共12个文件(10.73MB)&…

2026/9/9 0:00:48

MHS模型硬件标准:让大模型像调用软件一样控制物理设备

让Claude真正看着显微镜说“这个细胞形态不太对”,或者让大模型自己调一版机械臂的运动轨迹,这事儿听上去已经很接近科幻片了。但你真上手试一次就会发现,模型不缺智商,缺的是一个能插进显微镜、机械臂、激光控制器里的“通用插座…

2026/9/9 0:00:48

AI五大核心方向详解:从机器学习到大模型,零基础转行选哪条?

会有人告诉我,他想转行学AI,但打开招聘网站一看直接傻眼:机器学习、深度学习、自然语言处理、计算机视觉、大模型应用……满屏都是这些词,好像每个都会一点,又好像每个都离自己很远。还有人上来就问“学Python还是学Ja…

2026/9/9 0:00:49

从50行最小循环到生产级AI引擎:工程化改造全解析

直接说干货。这一章我写的不是那种"hello world跑通某个模型"的教程,而是把AI引擎当做一个真正要上线、要被人调用、要扛流量的系统来聊。从最初只有50行的最小循环,到能够承载生产流量的AI引擎,中间差的不是代码量,而是…

2026/9/7 16:23:03

USB Type-C PCB布局分区设计:电源、高速信号与PD协议全攻略

做硬件这行,Type-C接口算是典型的“看着简单,做起来全坑”的东西。光引脚就24个,高低速信号、电源、控制线全部塞在一个小小的连接器里,如果PCB布局不做规划,打样回来基本就是“插上没反应”、“高速掉线”、“静电一打…

2026/9/7 22:46:00

系统编程学习原型如何补齐稳定性边界

系统编程学习原型如何补齐稳定性边界预算有限时&#xff0c;我先优化明显多余的复制&#xff0c;而不是猜测性地换容器。用借用传递只读数据通常就能减少分配&#xff1a; fn parse(line: &str) -> Result<Item, Error> { /* ... */ }用基准确认热点确实在分配&am…

2026/9/9 10:21:54

雨花区哪家财务公司代理记账比较好?

在雨花区&#xff0c;企业处理财税事务常常面临诸多挑战&#xff0c;选择一家靠谱的财务公司至关重要。湖南巨勤财务管理咨询有限公司就是本地正规实体财税服务机构&#xff0c;深耕本地工商财税行业多年&#xff0c;熟悉当地工商局、税务局最新政策与申报流程。主营公司注册、…

还想了解更多?直接咨询顾问

免费诊断 + 免费方案 + 透明报价。

全国咨询热线400-8866-253
免费获取方案
咨询二维码