为什么GPU在等数据？LitData：AI训练数据加载加速器的完整入门指南

为什么GPU在等数据？LitData：AI训练数据加载加速器的完整入门指南

【免费下载链接】litData Speed up model training by fixing data loading. 项目地址: https://gitcode.com/gh_mirrors/li/litData

做深度学习训练的人都遇到过这一幕：GPU 利用率曲线忽高忽低，明明算力很贵，却被"等数据"拖慢了节奏。LitData 是一款专为 AI 模型训练打造的数据加载加速器，它通过预优化数据集格式和流式读取，让训练吞吐量最高提升 20 倍，让 GPU 不再干等。

🤔 为什么 GPU 会在"等数据"？

加速模型训练不只是调 kernel 那么简单。数据加载慢，通常是这三个原因造成的：

• 数据太大：数据集放不进本地磁盘，只能在线读取；
• 文件太碎：上百万个小文件，逐个读取 I/O 开销爆炸；
• 云上传得慢：从 S3 / GCS 等对象存储串行拉数据，网络成了瓶颈。

传统做法里，DataLoader 一个文件一个文件地读，GPU 算完一个 batch 就得停下来等下一个。LitData 的思路是：让数据"先跑一步"，并且批量、异步地跑。

⚡ 认识 LitData：两步解决数据瓶颈

LitData 的核心能力可以概括为两组操作：

一键安装 LitData

pip install litdata

需要更多模态支持（视频、医学影像等）时，安装附加包：

pip install 'litdata[extras]'

🚀 工作流一：零预处理，现有文件直接流式加载

如果你只是想立刻开训，不想先跑任何转换任务，可以用 StreamingRawDataset（源码位于 src/litdata/raw/dataset.py）：

• 指向一个本地或云端的图片/音频/文本文件夹即可；
• 拿到的是原始字节，用 PIL、torchaudio、json 想怎么解码就怎么解码；
• 下载完全异步且成批进行，云客户端自带自动重试；
• 文件索引（index.json.zstd）本地+远端双缓存，重启不重扫。

官方在 ImageNet val（S3 上 5 万个 JPEG）上的实测：8 个 worker 时吞吐从 4841 提升到 5718 张/秒（+18%），单进程场景（0 worker）提升可达 +35%。对"起步阶段"来说，这几乎是免费的提速。

🏎️ 工作流二：一次性 optimize，换取最高 20x 训练加速

追求极限吞吐时，走"优化 + 流式"两步路线（核心代码在 src/litdata/streaming/ 目录）：

第 1 步：优化数据（只做一次）

用 ld.optimize(...) 把原始数据转成高效的分块二进制格式（chunk-*.bin + index.json），支持多 worker 并行写入，参考示例：examples/getting_started/generate.py。

第 2 步：上传到云存储

一条 aws s3 cp --recursive 命令即可，数据从此"住在云上"。

第 3 步：训练时流式读取

from litdata import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset('s3://my-bucket/fast_data', shuffle=True, drop_last=True)
dataloader = StreamingDataLoader(dataset, batch_size=64)

for batch in dataloader:
    train_step(batch)

就这么短。相比普通 DataLoader，官方宣称训练数据加载最快提速 20 倍，且自带这些新手友好的特性：

✅ 无需本地下载：直接用云上数据训练，省磁盘、省时间 ✅ 断点续流：流式读取可暂停/恢复，训练中断不怕 ✅ 多 GPU 自动分片：自动保证每个 rank 拿到数量一致、内容不重叠的数据，开箱即用 ✅ 确定性 shuffle：同样的 seed 每次得到相同顺序，支持可恢复的 checkpoint ✅ 灵活存储：S3、GCS、Azure、Hugging Face、Lightning Studio 都能直连

💡 小贴士：优化图像数据时，建议以 JPEG 质量约 95 的形式写入（ld.Image(path=..., quality=95, format="jpeg")），比存未压缩像素小 10 倍以上，流式读取也更快。

🧩 还能做什么？LitData 的进阶能力

• 分布式数据变换 ld.map：把图像缩放、向量化、爬虫、批量推理等任务并行到多台机器上，源码见 src/litdata/processing/；
• 多模态全覆盖：图片、JPEG、音频、视频、PDF、医学影像（NIfTI）、3D 网格、Parquet、PyTorch Geometric 图数据、NumPy/Tensor 等都有对应包装类型，完整清单见 README.md 的 Modality 表格；
• 流式读 Parquet 和 Hugging Face 数据集：hf:// 开头的 URL 传进 StreamingDataset 即可自动建索引；
• 兼容 MosaicML MDS 格式：已有 MDS 分片可直接流式读取，零转换迁移；
• 按键查找与原地修补：构建键索引后可用 dataset["entity-id"] 直接取样本，甚至对已有数据集做局部更新。

📊 速度验证：基准测试怎么跑

仓库内置了完整的基准脚本，方便你在自己的数据集上复现提速效果：

• benchmarks/litdata/：LitData 的 optimize + stream 全流程（如 optimize_imagenet.py）；
• benchmarks/ffcv/：与 FFCV 对比的 ImageNet 转换/流式脚本；
• 使用方式说明见 benchmarks/README.md。

以 155 GB 的 ImageNet 为例，LitData 优化后的产物约 45 GB、耗时约 283 秒，比 Parquet 方案更小更快。

📁 项目结构导览：源码在哪里

想深入阅读源码，可以克隆仓库：

git clone https://gitcode.com/gh_mirrors/li/litData

✍️ 写在最后

GPU 很贵，别让它在"等数据"。LitData 给出的答案很直接：能不动就不动（StreamingRawDataset 直接流式读现有文件），要快就优化一次（optimize 换最高 20x 吞吐）。从 pip install litdata 开始，几分钟内你就能让训练流水线跑起来——剩下的，交给 GPU 算就好了。

【免费下载链接】litData Speed up model training by fixing data loading. 项目地址: https://gitcode.com/gh_mirrors/li/litData