Hugging Face Streaming Datasets Update

Hugging Face has significantly enhanced the load_dataset('dataset', streaming=True) functionality, allowing users to train on multi-terabyte datasets without downloading them locally. These optimizations eliminate common issues such as "disk out of space" errors and 429 rate-limiting responses, delivering performance that can match local SSD speeds in high-compute environments.

Performance Gains at Scale

Improvements to the streaming backend have resulted in substantial efficiency gains for large-scale training. When training on 64xH100 GPUs with 256 workers, Hugging Face reported the following metrics:

  • Startup Requests: Up to 100x more efficient (fewer requests).
  • Data Files Resolution Time: 10x faster.
  • Streaming Speed: Up to 2x faster.
  • In-flight Requests: Up to 2x more efficient.
  • Stability: Zero worker crashes at 256 concurrent workers.

Technical Optimizations

The performance boost is divided into two primary phases: startup and streaming.

Startup Optimization

To prevent "request storms" where every DataLoader worker independently initializes the dataset, Hugging Face introduced:

  • Persistent Data Files Cache: The first worker resolves the file list from the Hub, and all subsequent workers read from a local cache, virtually eliminating redundant startup requests.
  • Optimized Resolution Logic: API calls required to fetch the file list are now bundled more efficiently to reduce latency.

Streaming Throughput

To ensure GPUs are not left waiting for data, the following features were added:

  • Prefetching for Parquet: The library now fetches the next chunk of data in the background while the model processes the current chunk.
  • Configurable Buffering: Advanced users can now adjust the buffer's block size and the prefetch volume. For example, the minimum request size can be increased from the default 32MiB to 128MiB using pyarrow.dataset.ParquetFragmentScanOptions.

Data Transfer and Storage Infrastructure

Hugging Face utilizes Xet, a deduplication-based storage system, to accelerate uploads and downloads. By leveraging Parquet Content Defined Chunking (CDC), duplicated data is only transferred once, making uploads to Hugging Face faster than traditional remote storage. This is further supported by the pyspark_huggingface package, which provides a Spark Data Source for reading and writing datasets with Xet and Parquet CDC support.

Custom Streaming Pipelines

For formats not natively supported by the datasets library or for users requiring granular control, Hugging Face improved the HfFileSystem in the huggingface_hub library. This allows for efficient remote reading of files from dataset repositories using .read(), .readline(), and .seek().

When HfFileSystem is passed to a PyTorch DataLoader, it reuses cached results from .ls() and .glob(), eliminating the need for additional requests when listing data files. This approach has been utilized in the LeRobot library for video frame sampling and the WebDataset library for streaming TAR archives.

Real-World Application: nanoVLM

These enhancements were battle-tested during the development of nanoVLM for the next generation of SmolVLMs. Hugging Face found that streaming is now as fast as reading from local SSDs, removing a previous bottleneck where transferring data to local SSDs delayed training starts by three hours.

Sources