Hugging Face datasets 流式加载提速：请求减少最高 100 倍，吞吐翻倍

TLDR

我们提升了load_dataset('dataset', streaming=True)，只需一行代码即可流式传输数据集，无需下载！

立即在多TB级数据集上开始训练，无需复杂设置、无需下载、不会出现“磁盘空间不足”或429“停止请求！”错误。
速度超快！在64xH100上训练时，256个工作进程下载数据，速度超过了我们的本地SSD。 我们改进了流式传输，请求次数减少了100倍，→ 数据解析速度提升10倍 → 每秒样本数提升2倍，→ 在256个并发工作进程下0崩溃。

加载数据，尤其是在TB级别，是任何机器学习工作流程中的一大痛点。我们在训练SmolLM3时就深受其苦，有一次每次运行前都得等3个小时才能下载足够的数据。

在datasets库中流式传输一直可行，但使用海量数据集进行大规模训练仍然是个挑战。今天这一切改变了🔥。我们花了几个月时间改进后端，专注于流式传输数据集，使其更快、更高效。

我们具体做了什么？⤵️

流式传输：同样简单的API

首先：我们的改动是向后兼容的。你仍然可以用同样简单的streaming=True标志从Hub流式传输任何数据集。一如既往地简单。🚀

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时就是这么做的，所有数据都在S3上，直接从那里流式传输。

我们想改变这一点，所以在开发nanoVLM时决定使用Hub的流式传输。很快我们发现了一个大问题：我们的测试运行在不到一分钟内生成了超过10万个请求，导致我们的IP被Hub封禁！😅 这是因为每个DataLoader工作进程都独立初始化数据集。深入调查后，我们发现这造成了大量冗余请求，其中许多是不必要的。我们的改动最终将启动请求减少了100倍。总的来说，我们的改进带来了：

• 数据文件解析时间：快10倍
• 启动请求：效率提升高达100倍
• 流式传输速度：提升高达2倍
• 进行中的请求：效率提升高达2倍

幕后：我们改进了什么

那么，改变了什么？我们聚焦于两个阶段：启动和流式传输。

1. 启动⚡️ 最初的数据文件解析产生了大量请求。我们做了两项重大改变：

• 持久化数据文件缓存：我们现在在所有DataLoader工作进程之间缓存数据文件列表。第一个工作进程从Hub解析文件列表。所有其他工作进程直接从这个本地缓存读取，几乎消除了启动请求并大幅缩短了解析时间。不再有请求风暴！
• 优化解析逻辑：我们还最小化了初始工作进程获取文件列表所需的API调用次数。我们现在尽可能高效地捆绑必要的请求，进一步降低延迟。

2. 流式传输 🏎️ 为了提高流式传输本身的吞吐量，我们引入了两项新功能：

• Parquet 预取：我们为 Parquet 数据集启用了预取功能。这意味着当你的模型正在处理当前数据块时，datasets 库已经在后台获取下一个数据块。这保持了数据管道的满载，确保你的 GPU 永远不会因等待数据而闲置。
• 可配置缓冲：高级用户现在可以针对其特定硬件和网络设置微调流式传输性能。我们提供了配置缓冲区块大小和预取量的选项，让你能够最大程度地控制以优化 I/O。

以下是我们如何将流式传输时的最小请求大小从 32MiB（默认值）增加到 128MiB，并配置预取：

import pyarrow
import pyarrow.dataset

fragment_scan_options = pyarrow.dataset.ParquetFragmentScanOptions(
    cache_options=pyarrow.CacheOptions(
        prefetch_limit=1,
        range_size_limit=128 << 20
    ),
)
ds = load_dataset(parquet_dataset_id, streaming=True, fragment_scan_options=fragment_scan_options)

这些改进结合在一起，可以使你的数据吞吐量翻倍，让你能够更快、更高效地训练。

我们如何比普通 S3 更快：Xet

Hugging Face 使用 Xet：一种基于去重的存储，可实现快速去重上传和下载。与传统远程存储不同，Xet 上的数据传输更快，因为重复数据只传输一次。例如：将大规模数据集上传到 Hugging Face 利用 Xet 加速上传。数据集上传后，即可立即进行流式传输。

Parquet 的去重功能通过 Parquet 内容定义分块（CDC） 启用。得益于 Parquet CDC 和 Xet 去重，在 Hugging Face 上上传数据集比在任何传统远程存储上都更快。

这由我们的 pyspark_huggingface 包支持，它是一个用于读写 HF 数据集的 Spark 数据源。它包含 Parquet CDC 和 Xet 支持，可显著加速 HF 上的数据传输。

需要自定义流式管道？

datasets 中不支持某些数据文件格式，有时需要更多控制，因此我们让构建自定义流式管道变得简单。这已在 LeRobot 库中用于采样视频帧，并在 WebDataset 库中用于流式传输 TAR 归档，经过了实战检验。

我们改进了 huggingface_hub 库中的 HfFileSystem，以高效地从远程 Hugging Face 数据集仓库读取文件并流式传输数据：

from huggingface_hub import HfFileSystem

path = f"hf://datasets/{dataset_id}/{path_in_repo}"
with HfFileSystem().open(path) as f:
    # loop with .read() or .readline() to stream data
    # or do random access with .seek()

将 HfFileSystem 传递给 torch DataLoader 会重用 .ls() 和 .glob() 的缓存结果，从而在列出数据文件时无需额外请求。

将流式传输推向极限

我们现在在 nanoVLM 中使用这些流式传输增强功能来训练下一代 SmolVLM。通过这些调整，我们从流式传输中获得的性能比在我们集群的分层硬盘设置上训练更好。事实上，流式传输现在与从本地 SSD 读取数据一样快！以前，将数据传输到本地 SSD 的过程曾使我们的训练延迟三个小时。更多详情，请查看我们的 GitHub。

开始使用并体验差异

这些强大的新功能已在 datasets 和 huggingface_hub 库中落地。要利用它们，只需更新你的库并查看 文档：

pip install --upgrade datasets huggingface_hub

为了庆祝这一点，我们将 FineVision 中的所有数据源预连接并打乱为 FineVisionMax。你可以使用这个单一组合数据集来训练你的 VLM——无需手动处理多个数据集！

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)))

你可以在 nanoVLM 中看到我们如何大规模地做到这一点！

流式传输愉快！🤗