RINAS: Training with Dataset Shuffling Can Be General and Fast

TL;DR

RINAS is a data loading framework that significantly improves the efficiency of global dataset shuffling in deep learning pipelines by enabling parallel data fetching, leading to substantial throughput gains in language and vision model training.

Contribution

The paper introduces RINAS, a novel intra-batch unordered data fetching approach that enhances global dataset shuffling efficiency and system throughput in deep learning training pipelines.

Findings

Up to 59% throughput improvement in language model training

Up to 89% throughput improvement in vision model training

Effective global shuffling with reduced data loading overhead

Abstract

Deep learning datasets are expanding at an unprecedented pace, creating new challenges for data processing in model training pipelines. A crucial aspect of these pipelines is dataset shuffling, which significantly improves unbiased learning and convergence accuracy by adhering to the principles of random sampling. However, loading shuffled data for large datasets incurs significant overhead in the deep learning pipeline and severely impacts the end-to-end training throughput. To mitigate this, current deep learning systems often resort to partial dataset shuffling, sacrificing global randomness to maintain acceptable training throughput on large datasets, still leaving global shuffling efficiency issues not fully explored. In this work, we present RINAS, a data loading framework that systematically addresses the performance bottleneck of loading global shuffled datasets. Our key contribution is to offer an intra-batch unordered data fetching approach, which unleashes unexplored parallelism of data loading. We implement RINAS under the PyTorch framework for common dataset libraries HuggingFace and TorchVision. Our experimental results show that RINAS improves the throughput of general language model training and vision model training by up to 59% and 89%, respectively.