[Paper] Optimizing High-Throughput Distributed Data Pipelines for Reproducible Deep Learning at Scale

Source: arXiv - 2604.21275v1

Overview

The paper tackles a painful bottleneck that many of us hit when training huge deep‑learning models on petabyte‑scale data: the data‑loading pipeline becomes the limiting factor, throttling GPU utilization and making runs non‑deterministic. By profiling a typical PyTorch + Petastorm + Parquet workflow, the authors uncover where network I/O, CPU‑bound conversions, and shared‑queue contention sap performance, and they redesign the pipeline to achieve a 6× speed‑up and deterministic training at scale.

Key Contributions

• Root‑cause profiling of distributed data pipelines, pinpointing network I/O and PyArrow‑to‑NumPy conversion as the primary culprits.
• Push‑down worker‑level transformations that move heavy data preprocessing from the driver to each data‑loader worker, eliminating redundant work across epochs.
• Fanout‑Cache local‑disk caching strategy that stores a per‑worker copy of frequently accessed Parquet shards, drastically cutting network traffic.
• Deterministic queue architecture using a round‑robin ventilator and dedicated result queues, removing race conditions in multi‑worker setups.
• Modern RNG handling to guarantee reproducible shuffling and augmentation across runs.
• Empirical validation showing GPU utilization jump from ~10 % to >60 % and end‑to‑end training time drop from 22 h to 3 h on a multi‑node GPU cluster.

Methodology

1. Baseline Profiling – The authors instrumented a standard distributed training job (Petastorm loader → PyArrow → NumPy → GPU) on a 4‑node, 8‑GPU‑per‑node cluster. They measured per‑step latency, network bytes transferred, CPU usage, and GPU occupancy.
2. Bottleneck Isolation – By toggling individual stages (e.g., reading raw Parquet vs. pre‑cached shards) they showed that network reads and the PyArrow‑to‑NumPy conversion dominate the critical path.
3. Architectural Redesign –
   • Worker‑level transforms: each DataLoader worker now performs the heavy conversion locally, caching the result on its attached SSD/NVMe.
   • Fanout‑Cache: a lightweight daemon replicates needed Parquet files to each worker’s local disk the first time they are accessed, then serves subsequent epochs from disk.
   • Deterministic Queues: a central “ventilator” assigns batches to workers in a strict round‑robin order, while each worker writes results to its own queue, preventing contention.
   • RNG overhaul: seed propagation is handled per‑worker with a cryptographically‑secure generator, ensuring identical shuffling across runs.
4. Evaluation – The optimized pipeline was benchmarked on the same hardware and dataset (tens of TB of image/video data) across multiple training runs, comparing throughput, GPU utilization, and run‑to‑run variance.

Results & Findings

The findings confirm that moving transformations to the worker level and caching locally eliminates the repeated cost of pulling and decoding the same shards each epoch. The deterministic queue design also removes the nondeterministic jitter that previously made reproducibility a nightmare.

Practical Implications

• Faster Model Iteration – Cutting training from days to hours means data scientists can experiment with larger architectures or hyper‑parameter sweeps without waiting for a week‑long job.
• Cost Savings – Higher GPU utilization translates directly into lower cloud‑GPU spend; the same hardware does six times the work.
• Reproducible Pipelines – Deterministic data loading is essential for debugging, auditability, and compliance (e.g., in regulated AI). Teams can now lock down a training run and guarantee identical results across reruns.
• Scalable Architecture Blueprint – The push‑down + fanout‑cache pattern can be adopted in any framework that uses Parquet/Arrow datasets (TensorFlow, JAX, etc.) and is especially relevant for multi‑node clusters with high‑speed local SSDs.
• Simplified Ops – By offloading heavy I/O to workers, the central parameter server or driver node is less likely to become a bottleneck, easing cluster provisioning and monitoring.

Limitations & Future Work

• Hardware Dependency – The gains rely on each worker having fast local storage (NVMe/SSD). Environments without such disks may see smaller improvements.
• Dataset Format Specificity – Optimizations target Apache Parquet + PyArrow; other formats (e.g., TFRecord, custom binary blobs) would need separate adaptation.
• Cache Warm‑up Cost – The first epoch incurs a noticeable warm‑up as data is fan‑out to workers; for very short training runs this overhead could dominate.
• Future Directions – The authors suggest extending the approach to streaming data sources, integrating with container‑orchestrated pipelines (Kubernetes), and exploring automatic transformation placement (CPU vs. GPU) via profiling‑guided compilers.

Bottom line: By re‑thinking where and how data is transformed and cached, the authors turn a sluggish, nondeterministic data pipeline into a high‑throughput, reproducible engine—an upgrade that any team training large‑scale deep models should consider.

Authors

• Kashish Mittal
• Di Yu
• Roozbeh Ketabi
• Arushi Arora
• Brendon Lapp
• Peng Zhang

Paper Information

• arXiv ID: 2604.21275v1
• Categories: cs.DC
• Published: April 23, 2026
• PDF: Download PDF