Scalable and Performant Data Loading

TL;DR

SPDL addresses end-to-end data loading bottlenecks in AI workloads by coordinating data acquisition, preprocessing, and GPU transfer under GIL constraints. It introduces a thread-based, GIL-releasing pipeline architecture driven by an asynchronous scheduler and a dedicated thread pool, enabling efficient parallelism without IPC overhead. The system architectures, design principles, and I/O/GPU transfer primitives are evaluated against PyTorch DataLoader and DALI, showing substantial throughput gains, lower CPU and memory usage, and compatibility with Free-Threaded Python (FT-Python). The work demonstrates a practical, framework-agnostic path to scalable data loading that can keep GPU feeding rates aligned with faster training regimes while lowering operational costs.

Abstract

We present SPDL (Scalable and Performant Data Loading), an open-source, framework-agnostic library designed for efficiently loading array data to GPU. Data loading is often a bottleneck in AI applications, and is challenging to optimize because it requires coordination of network calls, CPU-bound tasks, and GPU device transfer. On top of that, Python's GIL (Global Interpreter Lock) makes it difficult to gain performance improvement from multi-threading. We found that when data preprocessing functions release the GIL entirely, it is possible to execute them concurrently in a thread pool, thereby improving the workflow performance. Our benchmark shows that compared to the PyTorch DataLoader, SPDL can iterate through the ImageNet dataset 74% faster while using 38% less CPU and 50GB less memory. When training ViT-B/16 model, SPDL can send data to the GPU at a speed that does not starve the training. Additionally, when using SPDL on Python 3.13t, without changing any code, the throughput is further by improved by 33%, thanks to the disabled GIL. SPDL can improve the performance of current AI model training, and receives further performance improvements when Free-Threaded Python is adopted in production systems. SPDL is available at https://github.com/facebookresearch/spdl.

Figures (15)

• Figure 1: The peak throughput of media processing functions in multi-threading and multi-processing. Images are decoded, resized and converted to RGB, then batched. The batch size is 32. ThreadPooolExecutor and ProcessPoolExecutor from concurrent.futures standard module are used. The test environment has 96 CPU cores. The time to initialize the workers, which is non-trivial for process workers, is not included.
• Figure 2: The time primitive operations take to execute. open is Python's built-in function and ImagingCore and ImagingDecoder are Pillow's low-level functions. The 0 concurrency means the image loading is executed in the main thread. Other concurrency values are the number of threads.
• Figure 3: Illustration of the GIL contention in multi-threading and different approaches for parallelization. Left: Parallelizing the entire data loading operation. The main thread and all the worker threads compete for the GIL, resulting in an ineffective use of the threads. Right: Parallelizing primitive operations separately. The scheduler thread dispatches operations that release the GIL to a thread pool. This way, the threads are not blocked. The competition to acquire the GIL is limited between the main thread and the scheduler thread. This approach also allows to configure the parallelism for different steps.
• Figure 4: The architecture of SPDL. An asynchronous event loop runs in a background thread. This event loop schedules tasks of different stages (data acquisition, processing and device transfer), and handles their completion. Asynchronous functions are executed as-is. They do not necessarily use the thread pool. Synchronous functions are delegated to the thread pool, so that they are handled in the same manner as native asynchronous functions. The data handle (Tensor/Array objects) are passed to the main thread via queue.
• Figure 5: The throughput of data loaders without downstream load.