AI训练要的数据这么多，怎么存？

在大模型时代，AI 的训练与推理已经不再是单纯的算力问题。随着模型参数规模进入百亿级、数据量级扩展到 TB~PB，数据管线与存储架构往往成为影响训练速度的决定性因素。算力再强，如果数据加载跟不上，也会出现 GPU 空闲等待 I/O 的情况，训练效率大打折扣。

张先森啊

1116人浏览 · 2025-09-24 15:49:42

张先森啊 · 2025-09-24 15:49:42 发布

📚AI Infra系列文章

AI Infra-为什么AI需要专属的基础设施？

GPU与TPU：AI计算硬件的核心对比

AI训练与推理的硬件需求有什么区别？

在大模型时代，AI 的训练与推理已经不再是单纯的算力问题。随着模型参数规模进入百亿级、数据量级扩展到 TB~PB，数据管线与存储架构往往成为影响训练速度的决定性因素。算力再强，如果数据加载跟不上，也会出现 GPU 空闲等待 I/O 的情况，训练效率大打折扣。

所有相关源码示例、流程图、模型配置与知识库构建技巧，我也将持续更新在Github：LLMHub，欢迎关注收藏！

希望大家带着下面的问题来学习，我会在文末给出答案。

大规模训练数据该如何分层存储与组织格式，才能兼顾成本与吞吐？
预处理、增广与特征生成放在离线还是在线更合适？
如何通过缓存与数据加载优化，避免GPU因I/O而闲置？

1. 分层存储与数据组织

在实际落地中，数据通常采用“热、温、冷”三层架构。

Hot 层存放的是训练时活跃的数据和索引，对 IOPS 和吞吐有极高要求，常见部署方式是本地 NVMe 或高性能分布式文件系统；Warm 层承担团队共享和版本管理，一般使用对象存储（如 Ceph、MinIO）或分布式文件系统；Cold 层则存放历史快照、原始数据和中间产物，用于后续再加工或追溯。这样的分层方式既能保证训练高效运行，又能合理控制存储成本。

在文件格式上，顺序可流式读取的容器化分片是主流选择，如 WebDataset（tar 分片）、TFRecord、Parquet 或 LMDB。合理的分片大小（通常在 100MB~2GB 之间）既能降低元数据开销，又能兼顾网络带宽与节点内存。为了支持可重复、可恢复训练，还需要建立样本级索引和全局 shuffle map。数据清洗与去重同样重要，文本可以使用 MinHash/SimHash，图像可用感知哈希，避免无效数据放大。

数据的版本管理也不可忽视。通过打标签和语义化版本号（如 imagenet-1.0.3），配合 DVC 或 LakeFS 等工具，可以让“数据—代码—模型”保持一致性和可追溯性。

2. 预处理与增广：离线 vs 在线

预处理和数据增强是训练效率的另一大关键。一个普遍的原则是：能离线就离线。

离线阶段适合完成解码、标准化、分词、图像 resize 或裁剪等固定流程，这样能显著降低训练时的 CPU 压力，并保证数据吞吐稳定。其优势是吞吐高、复用性强，缺点是灵活性略差、存储空间消耗稍多。

而在线处理则更适合需要随机性或多样化的增强手段，比如随机裁剪、颜色抖动、MixUp/CutMix、SpecAugment 等。这类操作对模型泛化能力帮助明显，且往往轻量、可矢量化，因此保留在 DataLoader 阶段即可。

对于 NLP，可以提前完成分词、分块，生成按上下文长度拼接的 packed samples；对于图像或视频，可以提前建立帧索引和 clip manifest，实现训练时的零拷贝读取。

3. 缓存策略

为了避免 GPU 因 I/O 停摆，缓存设计至关重要。

在本地机器层面，可以充分利用 Page Cache 和 mmap，大块顺序读能极大提升缓存命中率；

同时，将对象存储中的分片预拉取到本地 NVMe，作为短期缓存使用，也能明显提升首轮训练效率。常见做法是训练前做 warm-up，提前拉取头几个 epoch 所需数据，避免初期吞吐波动。

在集群层面，可以部署 Redis、RocksDB 或 Alluxio 作为热点数据和索引的共享缓存层。缓存一致性可通过数据版本号来管理，训练任务完成后再异步清理或逐版本淘汰，保证资源利用率。

一个小技巧是尽量**使用顺序读和大批量 prefetch，而不是频繁的小块随机 I/O，**这对分布式训练尤其重要。

4. 加速数据加载与管线并行

除了存储与缓存，数据加载和训练管线的并行优化同样重要。

在 PyTorch 中，可以通过调整 DataLoader 参数来提升效率：合理设置 num_workers（通常为 CPU 核数的一半到等于核数）、开启 pin_memory 与 prefetch_factor，并启用持久化 worker 避免频繁 fork。此外，结合 non_blocking=True 和混合精度训练，还能让计算与 HtoD 拷贝并行，减少等待。

loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=global_bs_per_rank,
    num_workers=os.cpu_count() // 2,
    pin_memory=True,
    prefetch_factor=4,
    persistent_workers=True,
    shuffle=False,
)

在 TensorFlow 中，tf.data 的自动调优功能可以简化并行与预取配置。

ds = (tf.data.TFRecordDataset(files, num_parallel_reads=tf.data.AUTOTUNE)
      .map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
      .shuffle(buffer_size=1<<16, reshuffle_each_iteration=True)
      .batch(batch_size, drop_remainder=True)
      .prefetch(tf.data.AUTOTUNE))