如何将流式数据集处理效率提升至百倍？

快速了解（TLDR） 现在只需一行代码，就能通过 load_dataset('dataset', streaming=True) 以流式方式加载数据集，无需下载！ 无需复杂配置、不占磁盘空间、不再担心 “磁盘已满” 或 429 请求过多错误，立即开始训练 TB 级数据集！ 性能非常强劲：在 64×H100、256 个并发 worker 环境下，流式加载速度甚至超过本地 SSD！ 我们优化后的流式系统：请求数减少 100 倍 → 数据解析速度提升 10 倍 → 样本处理速度翻倍 → 即使在 256 个并发 worker 下也 0 崩溃。 在机器学习中，特别是在处理 TB 级别的数据时，数据加载一直是个大难题。我们自己在训练 SmolLM3 时也深有体会，有段时间每次训练前都得等上 3 小时下载数据。 虽然 datasets 库早就支持流式加载，但在大规模训练中依然面临瓶颈。今天，这一切都变了 🔥。我们花了几个月优化后端，全面提升流式数据集的速度与效率。 那我们到底做了哪些优化？⤵️ 一样简单的 API，更强大的性能 首先最重要的一点：改进后的接口依然兼容原来的用法。你只需要加上 streaming=True，就能流式加载 Hugging Face 上的任意数据集，依旧简单直接。🚀 from datasets import load_dataset # Stream a dataset instead of downloading it dataset = load_dataset("HuggingFaceM4/FineVisionMax", split="train", streaming=True) # Get the first example print(next(iter(dataset))) 全球成千上万的 AI 开发者每天都在使用 datasets，现在他们无需改动任何代码，就能直接享受到更高的性能。 问题挑战：大规模流式加载 流式加载一直是快速了解数据集的好方法，但在训练模型时，大多数人仍然选择将数据预先下载到本地，或使用 S3 等云存储——我们在训练 SmolVLM 时也是这么做的。 我们希望改变这种情况，于是在开发 nanoVLM 时，尝试直接从 Hugging Face Hub 进行流式读取。 但很快就遇到一个严重问题：一次测试运行在不到一分钟的时间内发出了超过 10 万个请求，结果我们的 IP 被 Hub 屏蔽了！😅 问题的根源在于：每个 DataLoader 的 worker 都在独立初始化数据集，这导致大量冗余请求，形成了“请求风暴”，其中大部分其实是没必要的。 于是我们对启动逻辑进行了深度优化，最终将启动请求量减少了 100 倍。总体性能提升如下： 数据文件解析速度：提升 10 倍 启动请求效率：提高最多 100 倍 流式速度：提升最多 2 倍 在途请求效率：提升最多 2 倍 技术揭秘：我们具体改了什么？ 我们主要优化了两个阶段：启动阶段 和 流式加载阶段。 1. 启动优化 ⚡️ 初始的数据文件解析阶段会触发大量请求。我们进行了以下两项关键优化： 持久化数据文件缓存：现在所有 DataLoader worker 会共享数据文件列表缓存。第一个 worker 从 Hub 获取文件列表，其余 worker 直接从本地缓存中读取，从而几乎完全消除启动阶段的请求，大幅缩短加载时间，彻底告别“请求风暴”。 优化文件解析逻辑：我们精简了初始 worker 向 Hub 请求文件列表的 API 调用数量，将多个请求进行打包处理，进一步降低启动延迟。 2. 流式加载优化 🏎️ 为了提升训练过程中的流式吞吐量，我们新增了两个关键功能： Parquet 数据预取（Prefetching）：我们为 Parquet 格式的数据集启用了预取功能。这意味着，在模型处理当前数据块的同时，datasets 库会在后台提前加载下一块数据。这样可以让整个数据管道始终保持“满负荷”，确保 GPU 不会因等待数据而处于空闲状态，大大提升训练效率。 可配置缓冲机制（Buffering）：针对高级用户，我们开放了缓冲区的配置参数，支持自定义设置预取数量和数据块大小，方便根据自身硬件和网络情况进行 I/O 优化。