15. Feeding GPUs: data loading for training

Accelerators are the most expensive resource in an AI system, and training runs only as fast as its slowest stage. The data pipeline’s job is to deliver every batch before the GPU asks for it; an idle GPU waiting on data is money burning.

The input pipeline is a chain of producers and consumers

Read bytes from storage, decode them (JPEG, audio, JSON), transform or augment them, tokenize, assemble a batch, copy it to the GPU, compute. These stages run concurrently, so throughput equals the throughput of the slowest stage.

The symptom of a data bottleneck: GPU utilization that dips between steps, and a profiler showing step time spent waiting on the data loader. Always check this before tuning the model.

Do the throughput arithmetic first

Required read rate = global batch size × bytes per sample ÷ seconds per step.

• Images: 1,024 images of about 150 KB each is about 150 MB per step. At 0.5 seconds per step the cluster needs about 300 MB/s of reads, plus enough CPU to decode 2,000 JPEGs per second.
• Text: a token ID stored as uint16 takes 2 bytes when the vocabulary is under 65,536, otherwise 4 bytes as uint32. One million tokens per step is only 2 to 4 MB.

The conclusion surprises people: pre-tokenized text pretraining is rarely I/O-bound, while images, audio and video often are. Know which regime you are in before optimizing.

Why millions of small files fail

Object storage charges 10 to 100 ms of latency per request. Reading millions of individual small files in random order turns training into a latency benchmark. The fix is sharding: pack thousands of samples into large shard files of roughly 100 MB to 1 GB, and read each shard sequentially.

Formats built for this: WebDataset (tar shards), MosaicML StreamingDataset (MDS shards), TFRecord, and memory-mapped token arrays for text. Ray Data and Hugging Face Datasets streaming offer similar behaviour.

Shuffling without random access

Shuffling prevents the model from learning the order of the data, and keeps batches representative. A perfect global shuffle requires random access, which sharding removes. The standard compromise is a two-level shuffle:

1. Shuffle the order of shards each epoch.
2. Stream samples into an in-memory shuffle buffer and draw randomly from it.

Seed every shuffle. A deterministic order makes runs reproducible and lets a crashed job resume at the exact sample where it stopped.

The PyTorch DataLoader, in principle

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,                  # map-style (index access) or iterable (streaming)
    batch_size=256,
    num_workers=8,            # parallel worker processes for reading and decoding
    prefetch_factor=4,        # batches each worker prepares ahead of time
    pin_memory=True,          # page-locked host memory: faster copies to the GPU
    persistent_workers=True,  # keep workers alive between epochs
)

The levers are always the same: more parallel workers, more prefetching, and less work per sample. Raise num_workers until either the GPU is saturated or the CPU is.

Removing CPU bottlenecks

• Precompute what does not change between epochs: tokenize once, resize images once, store the result.
• Decode on the GPU where possible; NVIDIA DALI moves decoding and augmentation onto the accelerator.
• Keep randomness cheap: random crops and flips are fine on the fly; expensive transforms belong offline.

Distributed training data

• Each worker process (rank) must read a distinct part of the data, by assigning shards per rank or using a distributed sampler.
• Every rank must run the same number of steps; if one runs out of data early, collective operations hang.
• Save the data loader state (which shards and offsets were consumed) in every checkpoint. Without it, a restarted job silently re-trains on some data and skips other data.

Caching and locality

Stage hot datasets onto each node’s local NVMe before training, let the operating system’s page cache keep re-read data in RAM, or use a parallel file system or caching layer between object storage and the cluster. Each step closer to the GPU removes orders of magnitude of latency.

Checkpoints are a data problem too

Large models write checkpoints of many gigabytes. Write them asynchronously so training continues, shard them so each rank writes its own slice in parallel, and copy them to durable object storage in the background.

Sequence packing for LLMs

Documents vary in length, but training sequences have a fixed length. Padding short documents wastes compute on padding tokens. Packing concatenates documents, separated by end-of-document tokens, to fill each sequence. Many setups also mask attention across document boundaries, so one document cannot attend to an unrelated neighbour.

Questions to ask

• Is the GPU waiting on data, and how do we know?
• What read throughput does one training step need, and what can storage deliver?
• How many files does an epoch open, and how large are they?
• After a crash, does training resume at exactly the right sample?
• Which preprocessing could move offline so it runs once instead of every epoch?