The Problem with "Download First"
The naive approach to training on a large dataset is: download everything to local disk, then train. This works at small scale. At petabyte scale, it fails for three reasons:
- Time: Downloading 1 PB from S3 at 10 GB/s takes 28 hours. Your training job hasn't started yet.
- Cost: Local NVMe storage costs 10–20× more per TB than S3. Storing a full copy locally is expensive.
- Flexibility: You often don't know which subset of the dataset you'll need until the experiment is defined.
The solution is streaming: the dataloader fetches data from remote storage on demand, in parallel with the GPU's computation. The GPU never waits for a full download — it starts training as soon as the first batch is ready.
How Streaming Works
A streaming dataloader maintains a prefetch buffer: a queue of pre-fetched batches held in memory. While the GPU processes batch N, background worker threads are downloading batches N+1, N+2, ..., N+K (where K is the buffer size). As long as the download rate exceeds the GPU's consumption rate, the GPU never stalls.
Time →
GPU: [Batch 1] [Batch 2] [Batch 3] [Batch 4] ...
Workers: ↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓
Prefetching N+1, N+2, N+3 while GPU runs N
In WebDataset, streaming directly from S3 is just a matter of pointing the dataset constructor at a brace-expanded range of shard URLs (s3://.../shard-{000000..001399}.tar) instead of local paths; the same .decode()/.to_tuple()/.map_tuple()/.batched() pipeline chain that works on local shards works unchanged on remote ones, and wrapping it in a WebLoader gives you a standard PyTorch-style iterator with background workers handling the actual downloads.
Local Caching for Multi-Epoch Training
Streaming is efficient for the first epoch. For subsequent epochs (training typically runs 10–100 epochs over the same data), re-downloading from S3 is wasteful. The solution is an LRU (Least Recently Used) cache on local NVMe:
- Epoch 1: Shard is downloaded from S3 → written to local NVMe cache → fed to GPU.
- Epoch 2+: Shard is read from local NVMe cache → fed to GPU. No S3 request.
The cache evicts the least recently used shards when it fills up, ensuring the most frequently accessed data stays local. WebDataset supports this natively — passing a cache_dir and cache_size to the same WebDataset constructor is enough to turn on local LRU caching for a streaming dataset, no separate caching layer required.
Shard shuffling is important for training quality: if shards are always read in the same order, the model sees the same sequence of samples every epoch, which can harm generalization. WebDataset's shardshuffle=True randomizes shard order at the start of each epoch.
The architectural principle here is overlap computation with I/O: the GPU should never be the thing waiting. If the GPU is waiting, you are wasting money. Design your data pipeline so that I/O is always one step ahead of compute.