Data Streaming for Distributed Training: How to Keep GPUs Fully Utilized

Large training datasets can leave GPUs idle and waiting on I/O. In this blog, we show how to stream sharded data from object storage to local NVMe cache.

Why GPUs go idle during large-scale training

GPUs enable fast, parallel operations and typically dominate AI compute costs, so it’s important to constantly feed them data to process. GPU starvation occurs when accelerators that are ready to train instead rest idle as the input pipeline cannot provide the next batch of training data quickly enough.

Teams may witness GPU starvation with standard data loading approaches as the accelerators are faster at processing the data than enterprise NVMe storage can supply it. A drive delivering 14 GB/s in sequential reads provides less than one five hundredth of the bandwidth of the NVIDIA HGX™ B200 or B300 7.7 TB/s onboard memory speeds.

This article will explore leveraging Verda Object Storage, data streaming and local cache memory on multi-node Instant Clusters to maximize GPU and training run efficiency.

What GPU starvation looks like and how to diagnose it

Indicators of GPU starvation include:

• A high input-wait fraction - training-loop time is blocked waiting for the DataLoader.
• Low or fluctuating GPU utilization: utilization drops between bursts of compute when the next batches aren’t ready.
• Poor throughput: low number of data inputs processed every second despite accelerator availability. Throughput may not increase relative to additional GPUs.

Instant Clusters come with a preconfigured Grafana dashboard showing GPU utilization and other metrics. Nvidia’s nvidia-smi is also useful for quick and easy diagnosis. Additional tools include PyTorch Profiler and, for more extensive system-wide performance analysis, Nvidia’s Nsight Systems. Adding metrics and print statements to training runs is also illustrative.

DataLoader limitations

PyTorch's DataLoader provides several tuning options for improving sample-based throughput. Increasing num_workers parallelizes data fetching across CPU cores; setting persistent_workers=True avoids expensive worker respawn overhead between training epochs, and prefetch_factor enables asynchronous batch staging.

However, the underlying sample-based loading model may prove difficult to scale efficiently. Increasing worker counts and prefetch depth raises the amount of data, process state and buffering held in memory at once, while additional workers can bring contention for CPU, storage bandwidth and inter-process communication. DataLoader tuning can lead to diminishing returns for large or throughput-sensitive workloads.

How data streaming works

Data streaming addresses the bottlenecks of DataLoader by using larger, sequential shards rather than staccato loading of data samples from storage. Sharding splits the dataset across distributed nodes or partitions, enabling parallel reads so each shard can be streamed independently without shared state contention. Prefetching means future batches stream asynchronously while current batches execute.

Local caching makes frequently accessed or recently fetched shards available on fast local storage such as NVMe or RAM disks, preventing slower remote fetches and network costs across epochs.

Cold and warm cache

A cold cache run means the required shards are fetched from remote object storage as training progresses and then stored locally. Subsequent warm cache access can reuse those local shards. Cache limits can ensure local storage isn’t overloaded, or additional storage can house extensive caches.

Choosing a format: MosaicML Streaming (MDS), WebDataset, LitData

These libraries use cloud streaming directly from object storage buckets in various file formats for shards including .tar (WebDataset), .mds (MDS) and .bin (LitData). We will use MDS in an end-to-end benchmarking example shortly.

MDS does the heavy lifting of sharding, local caching and distributed partitioning, reducing custom data pipeline code. WebDataset is simple and portable, using .tar shards and sequential streaming. The library achieves high I/O rates and users will take more responsibility for shard layout, distributed sampling and caching behaviour. LitData supports broader data engineering preliminaries to model training, streams raw files directly and optimizes them into chunked datasets. It can also read MDS datasets.

End-to-End on Verda

Let’s walk through a benchmarking example. Our replicable task will be to train a ResNet-18 image classifier on the iNaturalist 2021 mini training dataset of 500,000 images representing 10,000 classes. The benchmark runs training steps including forward pass, loss calculation, backpropagation and optimizer update. These operations are to measure system performance rather than train a model to convergence. We’ll explore the limitations of DataLoaders and whether a streaming data architecture can reduce GPU starvation by capturing various metrics during a measured training window.

A note on measuring GPU utilization: nvidia-smi measures the percentage of the sampling period during which at least one GPU kernel is running. We don’t learn how many streaming multiprocessors (SMs) are busy, or how efficiently they are working. It’s important therefore to interpret utilization alongside training throughput and input-wait metrics.

Setup

The experiments run on an Instant Cluster with a Slinky environment, which offers a preconfigured Slurm environment on top of Kubernetes. This example was written using two nodes of 8 x NVIDIA B300 GPUs each. The experiment can be adapted to fit different clusters with a few configuration changes.

We’ll use Slurm to launch the training runs. If you’re a newcomer to the package, it’s a job scheduler ideal for distributed training across multiple GPUs.

Clone the repository

git clone https://github.com/verda-cloud/blog.git
cd blog/data-streaming

Repository layout

The code folders include:

• data_prep - Download, upload and shard conversion scripts.
• training - Parquet and MDS training benchmarks.
• slurm - Slurm job scripts.
• analysis - GPU utilization plotting code.
• results - Selected logs from our experiments.

Imports

From the repository root:

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Configure object storage

Verda Object Storage is an S3-compatible system for unstructured data such as datasets, model checkpoints, and media. Buckets and access keys are managed from the Verda console or API, and objects are readable and writable using any S3-compatible tool. To learn more about and set up Verda Object Storage, please see the documentation.

Let’s set some environment variables. Find the relevant details from your storage bucket in the console:

export VERDA_S3_ENDPOINT="https://YOUR_OBJECT_STORAGE_ENDPOINT"
export VERDA_S3_BUCKET="YOUR_BUCKET"
export VERDA_S3_ACCESS_KEY="YOUR_ACCESS_KEY"
export VERDA_S3_SECRET_KEY="YOUR_SECRET_KEY"

export DATA_PREFIX="inaturalist/data"
export DATA_MDS_PREFIX="inaturalist_mds_v1"

Data

We import the dataset from HuggingFace. We will upload both parquet files for our DataLoader baseline, as well as our MDS shards to our object storage bucket.

python3 data_prep/download_inaturalist.py 
python3 data_prep/upload_data.py ./downloads/inaturalist --prefix inaturalist # Upload to object storage
export DATA_PREFIX="inaturalist/data" # Change path as appropriate
python3 data_prep/convert_to_mds.py # Convert to MDS and upload

There should now be an inaturalist folder in the root directory.

Dataloader experiments

Our training code is in training/baseline_multinode.py for the DataLoader baseline and training/mds_multinode.py for data streaming. We can run their related Slurm jobs with these commands from the repository root:

sbatch slurm/baseline.sbatch
sbatch slurm/mds.sbatch

These .sbatch files are the job launchers, communicating to Slurm which resources to reserve, to start GPU monitoring (using nvidia-smi ) and to run the training code.

The MDS job should execute both cold and warm phases within one allocation. To follow the logs, check the job ID Slurm prints:

tail -f baseline-<JOB_ID>.out
tail -f mds-<JOB_ID>.out  

Running with fewer GPUs

For a single eight-GPU node, change the Slurm node count in slurm/baseline.sbatch to one, and keep eight processes per node. Reduce the GPU allocation and torchrun process count for smaller allocations.

Measurements

We capture various measurements including global throughput, wall time and mean input wait. These and other metrics are described in the results table below.

Results

The data shows general improvement from DataLoaders to data streaming, with some interesting and unpredictable results. Mean train time per step decreases steadily, however input-wait fraction increases for cold start data streaming. Ultimately, data streaming with warm start greatly improved on all metrics save the mean input wait and time-to-first batch.

Charts plotted from the CSV files of metrics captured by nvidia-smi during the training runs offer a more intuitive look at GPU activity:

In the DataLoaders / Parquet baseline we see drops in individual-GPU utilization, despite generally high activity during training. Some long input waits suggest uneven batch delivery.

For cold-start data streaming, the GPUs still exhibit sharp drops in activity, yet with longer bursts of high utilization above 80%. In the warm start version, GPUs consistently maintain higher utilization.

Finally, let’s compare the image throughput and wall time:

This relatively small experiment of 200 training steps gives a solid indication that implementing data streaming from object storage is a worthwhile experiment to make the most of GPU resources.

When local storage may suffice

For smaller datasets that fit comfortably on local NVMe or SSD storage, traditional PyTorch DataLoaders may be sufficient. If data can be read, transformed and loaded faster than the GPU completes each training step, the workload is compute-bound rather than input-bound, and optimizing the data pipeline may provide little benefit.

For larger workloads, Verda offers Object Storage and on-demand Instant Clusters with either preconfigured Slinky or Managed Kubernetes up to 144× NVIDIA HGX™ B300, 128× NVIDIA HGX™ B200, and 64× NVIDIA HGX™ H200 per cluster, all on InfiniBand. For custom configurations, contact us.