Hugging Face 流式数据集更新

Hugging Face 显著增强了 load_dataset('dataset', streaming=True) 功能，使用户能够在不本地下载的情况下在多 TB 数据集上进行训练。这些优化消除了诸如“磁盘空间不足”错误和 429 限流响应等常见问题，使得性能可以在高计算环境中匹配本地 SSD 速度。

大规模性能提升

对流式后端的改进在大规模训练中带来了显著的效率提升。在使用 256 个工作进程的 64xH100 GPU 上进行训练时，Hugging Face 报告了以下指标：

  • 启动请求: 最多提高 100 倍效率（减少请求）。
  • 数据文件解析时间: 提高 10 倍速度。
  • 流式速度: 最多提高 2 倍速度。
  • 飞行中请求: 最多提高 2 倍效率。
  • 稳定性: 在 256 个并发工作进程时零工作进程崩溃。

技术优化

性能提升分为两个主要阶段：启动和流式传输。

启动优化

为了防止每个 DataLoader 工作进程独立初始化数据集导致的“请求风暴”，Hugging Face 引入了以下措施：

  • 持久数据文件缓存: 第一个工作进程从 Hub 解析文件列表，所有后续工作进程从本地缓存读取，几乎消除了冗余的启动请求。
  • 优化的解析逻辑: 获取文件列表所需的 API 调用现在被更高效地捆绑，以减少延迟。

流式吞吐量

为了确保 GPU 不会因等待数据而闲置，添加了以下功能：

  • Parquet 预取: 库现在在模型处理当前数据块时，在后台预取下一个数据块。
  • 可配置缓冲: 高级用户现在可以调整缓冲区的块大小和预取量。例如，可以使用 pyarrow.dataset.ParquetFragmentScanOptions 将最小请求大小从默认的 32MiB 增加到 128MiB。

数据传输和存储基础设施

Hugging Face 使用 Xet，一种基于去重的存储系统，来加速上传和下载。通过利用 Parquet Content Defined Chunking (CDC)，重复数据仅传输一次，使得上传到 Hugging Face 的速度快于传统远程存储。这一点还得到了 pyspark_huggingface 包的进一步支持，该包提供了一个 Spark 数据源，用于在具有 Xet 和 Parquet CDC 支持的情况下读取和写入数据集。

自定义流式管道

对于 datasets 库不原生支持的格式或需要细粒度控制的用户，Hugging Face 改进了 huggingface_hub 库中的 HfFileSystem。这使得可以使用 .read()、.readline() 和 .seek() 高效地远程读取数据集仓库中的文件。

当 HfFileSystem 被传递给 PyTorch 的 DataLoader 时，它会重用来自 .ls() 和 .glob() 的缓存结果，从而在列出数据文件时无需额外请求。此方法已在 LeRobot 库中用于视频帧采样，以及在 WebDataset 库中用于流式传输 TAR 归档。

实际应用：nanoVLM

这些增强功能在开发下一代 SmolVLMs 的 nanoVLM 过程中经过了实战检验。Hugging Face 发现，流式传输现在的速度与从本地 SSD 读取相当，消除了之前由于将数据传输到本地 SSD 导致训练启动延迟三小时的瓶颈。

Sources