为什么“大模型训练”会被数据加载拖垮

在“Ai大模型训练教程”的实战环节里，很多训练任务并不是算力不够，而是GPU 等数据：

• 训练日志里 GPU utilization 长期 30%~60%，但 dataloader time 很高。
• 每次扩容 GPU 后，吞吐几乎不涨，甚至更差（I/O 争抢、网络瓶颈）。
• 多机多卡时，样本读取/解压/解码成为热点，CPU 打满，GPU 闲着。

要解决这些问题，核心在于把数据管道做成高吞吐、可扩展、可恢复、可复现的工程系统。本文聚焦三种在大模型训练中最常用的方案：WebDataset（tar shard）、Parquet（列式存储）和Streaming Dataset（流式/边下边训），并给出具体落地步骤与注意事项。

设计目标：先定义“高吞吐数据加载”该长什么样

在选型前建议把目标量化，否则会陷入“换格式万能论”。常见目标：

1. 吞吐：每 GPU 每秒多少样本/多少 token；训练端 data_time 占比 < 10%（视模型而定）。
2. 可扩展：从单机 8 卡扩到多机 128 卡，吞吐尽量线性增长。
3. 鲁棒性：节点重启、网络抖动不崩；能断点续训。
4. 可复现：同一 seed 下样本顺序可控（至少到 shard 级别可复现）。
5. 成本：对象存储/网络 egress/本地 NVMe 成本可控。

后文所有实践都围绕这些目标展开。

方案一：WebDataset（tar shards）——大模型图文/多模态的“工业默认值”

WebDataset 适用场景

• 图像、音频、视频等二进制样本（JPEG/PNG/MP3/MP4）+ 文本标注。
• 训练数据规模很大，存储在 S3/OSS/GCS 等对象存储。
• 需要极高吞吐，且希望把“小文件地狱”变为“大文件顺序读”。

其关键思想：把大量样本打包成多个 tar（shard），训练时顺序/并行地读取 shard 并在线解码。

推荐的 shard 规格（经验值）

• 单 shard 大小：256MB ~ 2GB（常见 512MB 或 1GB）。

  • 太小：对象存储请求数暴涨；太大：shuffle 粒度变粗、重试成本高。

• 每个样本由相同 key 的多文件组成：例如

  • 00001234.jpg、00001234.txt、00001234.json
• 元数据：建议把数据来源、license、hash、width/height、caption 语言等写入 .json，方便过滤。

数据打包流程（可落地步骤）

H3 步骤 1：准备规范化样本

建议在离线阶段完成：

• 图片统一转 JPEG（质量 90~95），去掉异常色彩空间。
• 文本清洗：去控制字符、统一换行、截断极长样本。
• 生成样本级 sha1/md5 去重依据。

样本文件命名务必可排序（如 8 位或 12 位零填充）。

H3 步骤 2：打 tar shard

思路：按顺序把样本写入 tar，达到阈值就切一个 shard，并生成 index（可选）。

工程建议：

• 不要在 tar 内再套压缩（例如 tar.gz）用于训练：会导致随机跳读困难、CPU 解压更重；通常直接 tar（未压缩）+ 图像本身是压缩格式已经足够。
• shard 命名包含分片编号：shard-{00000..99999}.tar。

H3 步骤 3：上传与缓存策略

• 对象存储：S3/OSS 路径按数据集/版本/日期划分：s3://bucket/ds/v1/shard-00000.tar

• 训练端：推荐本地 NVMe 缓存（尤其多机），避免每 step 都打对象存储。

  • 常见策略：每个 worker 预取若干 shard 到本地，LRU 淘汰。

训练侧 DataLoader 关键点

H3 关键点 1：shard 级 shuffle + 样本级 shuffle 的组合

为了既吞吐高又随机性够：

• 全局：shuffle shards（例如每个 epoch 打乱 shard 列表）
• 局部：每个 worker 内部维护一个 buffer（如 1k~10k 样本）做样本级 shuffle

buffer 越大随机性越好，但内存更多、启动更慢。

H3 关键点 2：并行解码与 CPU 亲和

吞吐瓶颈常在 JPEG 解码和 tokenization：

• num_workers 不等于越大越好：要看 CPU 核数、解码耗时、进程切换成本。

• 推荐做法：先用 4/8/16 做阶梯测试，观察：

  • GPU utilization
  • dataloader queue 是否饥饿
  • CPU 使用率是否过载
• 对多路 tokenization：考虑使用 fast tokenizer、批量 tokenize、或把部分 token 预处理离线化。

H3 关键点 3：错误样本处理与可观测性

生产数据必有脏样本：损坏图片、空 caption、编码错误。

实践建议：

• 解码异常：跳过并计数，不要让训练崩。
• 统计维度：每个 shard 的错误率、每类错误原因占比。
• 把“坏样本黑名单”写回离线处理流程，定期重打包。

WebDataset 常见坑与规避

• 单个 tar 内文件过多导致 tar 索引扫描变慢：控制 shard 大小与样本数（例如 1GB shard 里 50k 张小图会很痛）。
• 多机同时拉同一 shard导致热点：引入 worker 间的 shard 分配策略 + 本地缓存。
• shuffle 失效：若只顺序读 shard 且 buffer 太小，会出现训练不稳定；至少保证 shard 级 shuffle。

方案二：Parquet（列式存储）——大规模文本/结构化特征的高性价比方案

Parquet 适用场景

• 纯文本语料（instruction、对话、代码）、结构化字段（source、lang、quality_score、domain）。
• 训练需要频繁做过滤/采样：按 lang、domain、quality_score 进行动态配比。
• 希望更好地利用列式存储的压缩与投影读取（只读必要字段）。

Parquet 的优势在于：

• 列式压缩比高，I/O 更省。
• 支持 predicate pushdown（在引擎中能下推过滤）。
• 与数据湖生态兼容（Spark、DuckDB、Polars、Arrow）。

Parquet 的工程落地：如何分区与行组（row group）

H3 关键点 1：分区（partition）决定并行度与过滤效率

常见分区策略：

• lang=zh/en/...
• domain=web/code/book/...
• date=YYYY-MM-DD 或 version=v1

注意：分区太细会产生大量小文件；分区太粗过滤成本高。

经验：

• 每个 parquet 文件 256MB~1GB。
• 同一分区内控制文件数量，避免几十万小文件。

H3 关键点 2：row group 大小影响读取吞吐

• row group 太小：元数据开销大。
• 太大：并行度下降、随机访问变差。

常用范围：64MB~256MB/row group（按数据类型与压缩调整）。

文本训练中的 Parquet 典型管道

H3 步骤 1：离线写 Parquet（保留必要字段）

建议字段示例：

• text：训练文本
• source：来源
• lang：语言
• quality_score：质量分
• dedup_key：去重键
• length：字符数或 token 估计

离线阶段可先做：

• 去重（按 dedup_key 或 SimHash/MinHash）
• 质量过滤（低分丢弃或降采样）
• 语言识别与分桶

H3 步骤 2：训练时“按列读取 + 动态采样”

只读取 text 和少量控制列（如 lang/quality_score），不要把无用列拉进来。

动态配比示例：

• 50% 通用中文网页 + 30% 高质量指令 + 20% 代码
• 并且对 quality_score 低于阈值的样本以 0.2 概率保留（软过滤）

H3 步骤 3：与 tokenization 的配合

Parquet + 文本训练的瓶颈往往在 tokenizer：

• 如果 tokenizer 成本高，考虑：

  • 预先把文本切成段并缓存中间结果（例如 sentence split）。
  • 或建立“tokenized dataset”（存 token ids）但要权衡存储体积与灵活性。

Parquet 常见坑

• 对象存储上大量小 parquet：训练时 list/GET 成本巨大；务必做 compaction。
• 混合超长文本：会导致 batch padding 浪费；建议离线写入 length 并在训练中做 length-aware batching（相近长度拼 batch）。

方案三：Streaming Dataset（流式数据集）——边下载边训练、断点续训与弹性扩缩容

Streaming Dataset 适用场景

• 数据大到无法完整落盘或不想预下载。
• 需要“即开即训”，训练集随时追加（持续学习/增量数据）。
• 多机训练时希望以统一的方式做 shard 分配、缓存与断点。

Streaming 的核心能力：

• 远端 shard 的按需下载（prefetch）
• 本地缓存（cache limit、LRU）
• 弹性恢复（resume state）

工程实践：从零搭一个可用的 streaming 训练数据通道

H3 步骤 1：确定远端存储与 shard 清单

不管底层数据格式是 tar 还是 parquet，建议维护一个“manifest”（清单文件），包含：

• shard URL
• shard 大小
• 样本数估计
• hash 校验（可选）
• 数据版本号

manifest 的价值：训练端无需频繁列目录，减少对象存储开销，并提升可控性。

H3 步骤 2：配置本地缓存目录与容量

• 缓存目录最好放 NVMe：/local_nvme/cache/dataset
• 容量建议：至少能容纳“每机若干小时训练所需数据”（例如 200GB~2TB 视吞吐而定）。

缓存过小会导致反复下载，吞吐不稳定且对象存储费用上升。

H3 步骤 3：prefetch 深度与并行下载

prefetch 太浅：GPU 断粮；太深：缓存压力大、启动慢。

经验调参方法：

• 先记录每 step 消耗的数据量（例如 token/s 或 samples/s）。
• 估算对象存储到训练机的稳定带宽。
• 让 prefetch 覆盖至少 30~120 秒的数据需求（训练越大越建议更长），并观察 cache hit rate。

H3 步骤 4：多机多卡的 shard 分配

要点：避免所有 rank 拉同一 shard。

实践策略：

• 以 global_rank 对 shard 列表做切片（每个 rank 负责不同 shard）。
• 或者以“节点”为单位分配 shard，节点内共享缓存（能共享更好）。
• 对于 epoch 概念：可以定义“读完 manifest 一轮”为一个 epoch，或基于 token 计数。

H3 步骤 5：断点续训（resume）

Streaming 的断点不仅是模型/优化器状态，还包括：

• 当前读到的 shard 位置
• 随机数状态（shuffle buffer）
• 采样比例状态（若有动态采样）

建议在 checkpoint 中记录数据加载器的状态或可重建信息（例如当前 shard index、已消费样本数）。

Streaming 常见坑

• 把“流式”当成“无限随机访问”：远端访问仍有延迟，必须 prefetch + 缓存。
• 忽略对象存储限流：并发 GET 太多会 503/429；需要限流与重试退避。
• 跨 region 拉数据：延迟与费用都爆炸；训练计算尽量与数据同地域。

选型建议：WebDataset vs Parquet vs Streaming 怎么选

按数据类型

• 图像/音频/视频：优先 WebDataset（tar shard）
• 结构化文本（多字段过滤/采样）：优先 Parquet
• 训练需要随开随训、增量追加、弹性扩缩容：引入 Streaming Dataset（可承载 WebDataset 或 Parquet 的 shard）

按团队工程成熟度

• 想快速落地吞吐提升：WebDataset 最容易“立竿见影”。
• 有数据湖/ETL 体系：Parquet 能把“数据治理”和“训练采样”打通。
• 需要长期运营与在线增量：Streaming 是终局形态，但要投入可观测性与缓存治理。

性能排查清单：把吞吐问题定位到“哪一段慢”

训练吞吐不理想时，按链路拆解（建议逐项打点）：

1. 远端 I/O：list/GET 时延、吞吐、错误率（503/429）。
2. 本地 I/O：NVMe 读带宽、文件系统是否成为瓶颈。
3. 解码/解析：JPEG 解码、JSON 解析、文本清洗。
4. tokenization：CPU 占用、batch tokenize 是否生效。
5. collate/padding：是否出现极端长样本拖慢 batch。
6. 拷贝到 GPU：pin memory、异步拷贝是否开启。

建议输出关键指标：

• samples/s 或 tokens/s
• data_time / compute_time
• cache_hit_rate
• bad_sample_rate

这样你会清楚：是带宽问题、CPU 问题还是数据分布问题。

一个可直接套用的工程落地路径（建议按周推进）

第 1 周：确定格式与最小闭环

• 多模态：先把 1~5TB 样本打成 WebDataset shard
• 文本：先把核心语料写成 Parquet（带 lang/quality/length）
• 用单机 8 卡跑通训练，建立吞吐基线

第 2 周：引入缓存与可观测

• 本地 NVMe 缓存 + 命中率统计
• 错样本跳过 + 错误码分桶
• dataloader 全链路耗时打点

第 3~4 周：多机扩展与 streaming

• shard 分配策略（按 rank/按节点）
• prefetch 深度调优、对象存储限流与重试
• 断点续训：checkpoint 中记录数据状态

最终目标：扩 GPU 时吞吐能跟上，训练不再“卡在数据上”。

小结

在 Ai大模型训练教程 的实战训练中，“高吞吐数据加载”不是锦上添花，而是能否把算力吃满、把成本打下来、把训练稳定性做上去的关键工程。

• WebDataset 解决海量小文件与多模态吞吐问题；
• Parquet 让文本/结构化语料的过滤、采样与治理更高效；
• Streaming Dataset 提供可扩展的边下边训、缓存、恢复与弹性能力。

如果你希望我根据你的实际情况（数据类型、规模、存储介质、训练框架、目标吞吐）给出具体参数建议（shard 大小、prefetch、num_workers、缓存容量、采样策略），可以把这些信息补充出来，我会按你的场景给一份可执行的配置与排查方案。