Streaming Data Pipelines for PyTorch: Keep the GPUs Fed

Most "we need more GPUs" conversations end at the data loader. A training step that waits 90 ms for a batch and computes for 60 ms is a machine running at 40% of the hardware you are paying for, and no amount of torch.compile, FSDP2 sharding or FP8 will fix it. This tutorial covers the part of the stack that gets the least attention and returns the most: a streaming, resumable, sharded input pipeline for PyTorch training that keeps the accelerators fed.

Step 1: measure before you rewrite

Do not guess. Time the loader on its own, with no model in the loop:

import time, torch

def loader_throughput(loader, steps=200, warmup=20):
    it = iter(loader)
    for _ in range(warmup):
        next(it)
    torch.cuda.synchronize()
    t0, samples = time.perf_counter(), 0
    for _ in range(steps):
        batch = next(it)
        samples += batch[0].shape[0]
    dt = time.perf_counter() - t0
    return samples / dt

print(f"{loader_throughput(train_loader):.0f} samples/s from the loader alone")

Then compare against the model's standalone throughput on synthetic tensors already resident on the GPU. Three outcomes:

• Loader faster than model: you are compute-bound. Stop here; go optimize the model.
• Loader slower than model: you are input-bound. Every improvement below converts directly into training speed.
• Both fast, combined slow: you have a transfer or synchronization problem — usually a blocking .cpu(), a .item() inside the step, or non-pinned host memory.

A torch.profiler trace confirms it: long gaps between kernel bursts on the compute stream with the CPU sitting in next(iter) is the signature of an input-bound run. See Profiling PyTorch training for reading traces properly.

Step 2: fix the file format first

The single biggest determinant of pipeline throughput is not code, it is layout. Millions of small files on network storage is the worst case: every sample costs a metadata round-trip, and object stores charge per request. The fix is sharding: pack samples into a few hundred sequential archives of 100 MB–1 GB each.

Layout	Read pattern	Good for
Individual JPEG/JSON files	random, one request per sample	small local datasets only
Tar shards (WebDataset)	sequential streaming	images, audio, video, mixed media
Parquet / Arrow	columnar, sequential + projection	tabular, text, embeddings
Pre-tokenized .npy/.bin token stream	memory-mapped, near-zero cost	LLM pretraining and long fine-tunes

Two rules of thumb. Shard count should be a comfortable multiple of world_size * num_workers so every worker gets whole shards without stragglers. And do offline work offline: resizing images, tokenizing text and computing static features once at prep time is always cheaper than doing it every epoch.

Step 3: a resumable streaming loader

A map-style Dataset needs a global index and random access, which sharded remote data does not offer cheaply. Iterable-style datasets stream, but naive implementations duplicate data across ranks and cannot resume mid-epoch. torchdata's StatefulDataLoader is the PyTorch-native answer: a drop-in DataLoader replacement whose state_dict() captures worker-level progress.

import torch, torch.distributed as dist
from torch.utils.data import IterableDataset, get_worker_info
from torchdata.stateful_dataloader import StatefulDataLoader

class ShardStream(IterableDataset):
    """Streams pre-tokenized shards, split across ranks and workers, resumably."""

    def __init__(self, shards, seq_len=2048, seed=0):
        self.shards, self.seq_len, self.seed = sorted(shards), seq_len, seed
        self.epoch, self.shard_pos = 0, 0

    def _my_shards(self):
        rank = dist.get_rank() if dist.is_initialized() else 0
        world = dist.get_world_size() if dist.is_initialized() else 1
        info = get_worker_info()
        wid, nworkers = (info.id, info.num_workers) if info else (0, 1)
        g = torch.Generator().manual_seed(self.seed + self.epoch)
        order = torch.randperm(len(self.shards), generator=g).tolist()
        mine = order[rank::world]
        return mine[wid::nworkers]

    def __iter__(self):
        import numpy as np
        for i, shard_idx in enumerate(self._my_shards()):
            if i < self.shard_pos:          # skip shards already consumed
                continue
            self.shard_pos = i
            tokens = np.load(self.shards[shard_idx], mmap_mode="r")
            n = (len(tokens) // (self.seq_len + 1)) * (self.seq_len + 1)
            for start in range(0, n, self.seq_len + 1):
                chunk = torch.from_numpy(tokens[start:start + self.seq_len + 1].astype("int64"))
                yield chunk[:-1], chunk[1:]

    def state_dict(self):
        return {"epoch": self.epoch, "shard_pos": self.shard_pos}

    def load_state_dict(self, sd):
        self.epoch, self.shard_pos = sd["epoch"], sd["shard_pos"]

loader = StatefulDataLoader(
    ShardStream(shard_paths),
    batch_size=8,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=4,
    drop_last=True,
)

Save loader.state_dict() alongside your model and optimizer state in the same distributed checkpoint, and restore it with loader.load_state_dict(...). This is what makes spot and preemptible training honest: a resumed run continues on unseen data instead of silently replaying the first shards, which is a subtle and real cause of overfitting on interrupted jobs.

Step 4: tune the knobs in the right order

1. num_workers: start at 4–8 per GPU and raise while throughput improves. Too many workers thrash page cache and RAM; watch resident memory.
2. pin_memory=True plus non_blocking=True on the .to("cuda") call — pinned staging buffers let the copy overlap compute. Both halves are required; non_blocking without pinning does nothing.
3. persistent_workers=True so workers are not re-forked every epoch (a real cost with short epochs).
4. prefetch_factor: 2 is the default; 4–6 smooths bursty decode times at the cost of memory.
5. collate_fn: do the cheap, vectorized thing. Building tensors sample-by-sample in Python is a common hidden cost; prefer torch.from_numpy on a pre-batched array.
6. Batch on the worker, not the main process for variable-length data: yield whole batches from an iterable dataset with batch_size=None.

For long-sequence LLM training, add sequence packing — concatenating short examples up to seq_len with correct attention masking or block-diagonal FlexAttention masks. On instruction datasets with a wide length distribution, packing routinely removes 30–50% of wasted padding compute, which is a bigger win than any loader tweak.

Step 5: move decode off the CPU when it is the wall

If workers are pinned at 100% CPU doing JPEG decode or resizing, more workers will not help. Options, cheapest first:

• Pre-resize offline to the training resolution (usually enough).
• Decode on the GPU with torchvision.io.decode_jpeg(..., device="cuda") or NVIDIA DALI.
• Overlap host-to-device transfer on a dedicated torch.cuda.Stream so copies hide under the previous step's kernels.
• Cache the decoded, augmented tensors for the first epoch if augmentation is light and the dataset fits on local NVMe.

Step 6: make the pipeline reviewable

An input pipeline is production code and deserves the same discipline as the model:

• Version the data. Record the shard manifest (paths plus checksums plus counts) in the run's config. "Which data produced this checkpoint" must be answerable months later.
• Assert on batches. Shape, dtype, label range, NaN checks — once, in a smoke test, on the first batch of every run. A five-second check prevents the OOM or NaN discovered at hour 20.
• Log samples/second per rank. Divergence between ranks means a shard imbalance, not a hardware fault.
• Test resumption. Run 100 steps, kill it, resume, and confirm the loader state advances rather than restarting. Untested resume paths are usually broken resume paths.

The payoff

Reordering a small-file dataset into tar or token shards, adding a stateful loader and packing sequences regularly takes a training job from 40% to over 90% accelerator utilization. That is the same speedup people expect from a hardware upgrade, obtained for a day or two of engineering, and it compounds with everything in distributed training with FSDP2.

If your training runs are input-bound, or you need a sharded, versioned dataset built from a messy source, that is exactly our data processing services and PyTorch performance optimization work. Contact us with your current samples-per-second and GPU utilization and we will tell you where the pipeline is losing time.