ai-infra-interview-305

第 154 题：大模型训练的data pipeline优化？WebDataset、tfrecord？

题目

大模型训练的data pipeline优化？WebDataset、tfrecord？

---

完整讲解

一、Data Pipeline 瓶颈

大模型训练数据量大、预处理复杂（tokenize、pack、augment），若 DataLoader 跟不上，GPU 会空转、吞吐下降。瓶颈常在：磁盘 IO、CPU 预处理、Python GIL、数据从 CPU 到 GPU 的拷贝与排队。优化目标：让数据预取充足、与计算重叠、减少主线程阻塞。

二、WebDataset、tfrecord 等格式

WebDataset：基于 tar 的流式格式，按 shard 存储样本，训练时流式读、按需解压与解码，无需先全部解压到磁盘，适合超大规模数据集与对象存储。tfrecord：TensorFlow 的二进制序列格式，支持压缩、可随机访问或顺序读，多用于 CV/NLP 的预处理好结果缓存。Parquet/Arrow：列存、便于按列过滤与并行读，适合表格型或需过滤的 metadata。共同点：二进制、可压缩、利于顺序或并行 IO，减少小文件与解析开销。

三、Pipeline 优化手段

多 worker DataLoader、pin_memory、prefetch_factor；数据存 SSD 或高带宽存储、用多进程读；预处理放 DataLoader worker 或独立进程、与 GPU 计算重叠；大语料用 WebDataset 流式 + 多 shard 并行；tfrecord/parquet 做预处理好缓存，训练时只做轻量解码。

---

面试要点

• 瓶颈：磁盘 IO、CPU 预处理、GIL、CPU→GPU 拷贝；目标预取充足、与计算重叠。
• WebDataset = tar 流式、按 shard、无需全量解压；tfrecord = 二进制缓存、可压缩；Parquet/Arrow = 列存、过滤友好。
• 优化：多 worker、pin_memory、prefetch；高带宽存储；预处理与计算重叠；流式+多 shard。

---

记忆要点

1. 瓶颈 = IO + 预处理 + 拷贝；目标 = 预取 + 重叠。
2. WebDataset 流式；tfrecord 缓存；Parquet 列存；均利于大规模与并行。
3. 多 worker、pin_memory、prefetch、高带宽存储、预处理重叠。

返回模块

返回总览