Optimizing ML Data Loading: S3 to GPU Zero-Copy with Vortex
This article discusses how Vortex, an open-source columnar file format, rethinks data loading for ML training to achieve high-throughput from S3 directly to GPUs. It focuses on eliminating CPU/NVMe bottlenecks through techniques like cascading lightweight encodings, layout-based segment pruning, and zero-copy memory pipelines. The core idea is to reduce data movement taxes and decision taxes, enabling faster iteration on ML models by processing data on the fly without costly reprocessing.
Read original on InfoQ ArchitectureTraditional machine learning data loading pipelines, especially when dealing with large datasets stored in object storage like S3, involve multiple steps that introduce significant overhead. Data typically moves from S3 to NVMe disk, then to RAM (where it's decompressed by the CPU), and finally to the GPU via PCIe. This multi-stage process is bottlenecked by CPU decompression and NVMe throughput, leading to underutilized, expensive GPUs and slow iteration cycles for ML engineers.
The Problem: Movement and Decision Taxes
The author identifies two primary 'taxes' paid when utilizing GPUs for ML training:
- Movement Tax: The cost and latency associated with moving data from its source (e.g., S3) through various layers (disk, CPU, RAM) to the GPU. This involves multiple copies and decompression steps that eat into valuable GPU time.
- Decision Tax: The overhead of reprocessing entire datasets when simple changes are needed (e.g., new filters, data mixes, or curriculum changes). Current tools often require full data re-tokenization, reshuffling, and reloading, which is time-consuming and hinders rapid experimentation.
Vortex: A Solution for Efficient Data Loading
Vortex, an open-source columnar file format, addresses these issues by designing a pipeline that streams data from S3 directly to GPUs. Key architectural differences from formats like Parquet include:
- Decoupled Types and Encodings: Vortex separates logical types from their physical disk encodings, offering an extensible specification that allows for various plug-ins. This enables compute on compressed data and avoids costly materialization of intermediate results.
- Lightweight Cascading Encodings: Instead of losing data meaning with block compression, Vortex uses lightweight and cascading encodings that allow computations (like random access and aggregations) directly on compressed data, often faster than on decompressed data. This enables GPU-side decompression for parallel processing.
- Zero-Copy Memory Pipelines: By leveraging alignment information baked into the file format, Vortex can allocate receiving buffers with the correct alignment, eliminating the need for extra copies when dispatching compute to SIMD operations or GPU kernels.
- Layout-based Segment Pruning: Vortex layouts act as a minimal logical plan, enabling rapid pruning of data segments based on filters and projections. Zone maps containing summary statistics (min/max) allow for early data elimination without reading the full segments, significantly reducing I/O.
Core Design Principle
Vortex aims to maximize GPU utilization by minimizing data movement and preprocessing. It achieves this by pushing query optimizations (projections, filters) down into the file format itself, allowing data to be streamed from S3 to GPU at speeds up to 60 Gbps, making ML training iterations significantly faster.