Feeding the GPUs — the distributed data pipeline

Every lesson so far assumed the data was just there. It isn't. On N ranks you must hand each one a disjoint, shuffled slice of every epoch — and you must do it fast enough that a $40/hr GPU never waits on a CPU reading a file. The input pipeline is the one subsystem that can silently cap a perfectly-tuned training run.

Two jobs, both easy to get wrong

A distributed data pipeline has exactly two responsibilities, and each has a failure mode that doesn't show up as an error — only as a slow or subtly-wrong run.

1. Partition correctly. Across one epoch, the N data-parallel ranks must collectively see the whole dataset, once, with no sample landing on two ranks. Get this wrong and you either train on duplicates (some samples weighted 2×, biasing the model) or skip data — and because the loss still goes down, nobody notices.
2. Keep up. The pipeline must deliver the next batch before the GPU finishes the current step. A GPU step is milliseconds; reading and decoding a shard from disk can be longer. If supply < demand, the GPU stalls — and you're paying full price for an idle accelerator.

The rest of this lesson is those two problems and their standard solutions: DistributedSampler for partitioning, and prefetching workers for keeping up.

Partitioning — DistributedSampler in one picture

The trick is that every rank runs the same shuffle. Seed a permutation of all indices with (base_seed + epoch) — a number all ranks agree on — so they generate an identical shuffled list without talking to each other. Then rank r takes a strided slice: indices r, r{+}N, r{+}2N, …. Strided slicing of a shared permutation is disjoint by construction and needs zero communication.

  shuffled(epoch=0):  [ 7  2  9  4  0  5  1  8  3  6 ]        N = 2 ranks
                        ▲     ▲     ▲     ▲     ▲             rank 0 = positions 0,2,4,6,8 → 7 9 0 1 3
                           ▲     ▲     ▲     ▲     ▲          rank 1 = positions 1,3,5,7,9 → 2 4 5 8 6
                        └────────── disjoint, union = whole set ──────────┘

  next epoch → reseed → different permutation → different (still disjoint) split

Two correctness details the picture hides:

• Padding for divisibility. If the dataset size isn't a multiple of N, the sampler pads by repeating a few early indices so every rank gets the same count (ranks must run the same number of steps or the collectives in lesson 02 deadlock — a rank that finishes early leaves its peers waiting on an AllReduce that never completes). Those few duplicates are the price of keeping all ranks in lockstep.
• You must call set_epoch(epoch). Forget it and the seed never changes, so every epoch sees the identical shuffle — the model sees the same batch order forever. Loss still falls; generalization quietly suffers. This is the single most common DistributedSampler bug.

Sharding the bytes, not just the indices

DistributedSampler partitions indices, which is fine when every rank can cheaply random-access any sample (a local memory-mapped array). At real scale the data is terabytes of compressed shards on a network filesystem or object store, and random per-index reads are death — every rank seeks all over every file. Two patterns fix this:

The trade is shuffle quality vs read pattern. Index-level gives a perfect global shuffle but random I/O; shard-level gives sequential I/O but only an approximate shuffle (a finite shuffle buffer can't reorder across distant shards). For pretraining on web-scale data the shard-level + buffer approach wins, because saturating read bandwidth matters more than a mathematically perfect shuffle.

Keeping up — the starvation problem

This is the part that surprises people. The GPU step is the expensive thing, so intuition says the GPU is always the bottleneck. But producing a batch — read shard, decompress, tokenize, collate, copy to GPU — runs on the CPU, and if it takes longer than a GPU step, the GPU sits idle waiting. The fix is to overlap: run several worker processes that prepare future batches ahead of time, so a finished batch is always waiting when the GPU asks.

The model is a simple producer–consumer. The GPU consumes one batch every T_c (compute time). One worker produces one batch every T_l (load time); W workers in parallel produce one every T_l/W. The GPU never starves iff supply keeps up with demand:

T_l / W ≤ T_c ⇒ GPU utilization = min(1, W · T_c / T_l)

The lesson in that formula: throwing more workers at the problem helps only up to W = T_l / T_c. Past that, the GPU is already saturated and extra workers just burn CPU and RAM. And if a single batch's T_l is enormous (a giant image decode, a slow tokenizer), even many workers may not close the gap — you have to make the per-batch work cheaper (pre-tokenize offline, cache, compress less).

The throughput sanity check

Before any of this, do the back-of-envelope. The cluster consumes tokens at a rate set by the global batch and the step time:

demand = (B_eff · seq_len · bytes_per_token) / step_time

For a large pretraining run — say B_eff = 4M tokens/step at one step every ~2 s — that's ~2M tokens/s. At ~2 bytes/token of tokenized data that's a few MB/s of token IDs, trivially served. But if the pipeline reads raw data and tokenizes on the fly, the figure is the raw bytes — often 5–20× larger — and now you're asking a network filesystem for hundreds of MB/s to GB/s, sustained, split across every node. This is exactly why production pipelines pre-tokenize and pack offline: it collapses the bytes the hot path has to move. If demand exceeds your storage bandwidth, no number of workers saves you — the bottleneck is the wire.

The checklist

1. Partition with a seeded shared shuffle + strided slice (DistributedSampler), and actually call set_epoch. Verify disjointness and equal per-rank counts.
2. Shard the bytes to match your storage: shard-level sequential reads + shuffle buffer for object-store-scale data; index-level for data that fits locally.
3. Pre-tokenize and pack offline so the hot path moves token IDs, not raw bytes, and every token in a batch does work.
4. Size the workers to W ≈ T_l/T_c; add pinned memory + double buffering so loads hide behind compute.
5. Make it resumable — record epoch + per-rank consumed offset in the checkpoint.
6. Profile it — if GPU utilization sits below ~95% and the kernels themselves are fast, suspect the input pipeline before you touch the model.

Packing — don't pad, concatenate

Documents have wildly different lengths. The naive fix is to pad every short document up to the fixed sequence length L the model expects, then mask the pad tokens in the loss. The problem is that the GPU still computes on every pad token — attention, MLP, the works — and then you throw the result away. A batch that is 40% padding is a batch where 40% of the FLOPs heat the room for nothing.

Packing kills the padding instead of masking it: concatenate multiple documents end-to-end until you fill one length-L sequence. A 200-token doc, a 350-token doc, and a 450-token doc become one packed 1000-token row — zero pad tokens.

  PADDED   (3 rows, L=1000):           PACKED  (1 row, L=1000):
    [docA 200 | pad 800 .............]   [docA 200 | docB 350 | docC 450]
    [docB 350 | pad 650 .........]       └── one sequence, 0 pad tokens ──┘
    [docC 450 | pad 550 .......]
    useful 1000 / total 3000 = 33%       useful 1000 / total 1000 = 100%

The catch is correctness: in a packed row, tokens of document A must not attend to document B, or you leak context across unrelated documents and corrupt the loss. A plain full-attention kernel would let them. So you record the per-document boundaries as a cumulative-offset array cu_seqlens (e.g. [0, 200, 550, 1000]) and hand it to a variable-length attention kernel — FlashAttention's varlen path — which makes attention block-diagonal: each document attends only within its own segment. The varlen mechanics live in the gpu_kernels/vLLM tracks; here the point is that packing is only legal because the kernel enforces the boundaries.

Define packing efficiency as the fraction of computed tokens that are real:

η = useful_tokens / total_tokens

This is a throughput lever, not a data-engineering footnote. If you do not pack and your corpus packs to η ≈ 0.6, you are paying for 1/0.6 ≈ 1.67× the GPUs to train the same number of real tokens — a 67% overhead that no kernel tuning or worker-count change in this lesson can recover. The fix is upstream, in how you build the sequences.

Resumable loader state — the actual mechanism

A 3-week run checkpoints periodically and resumes after every preemption. The model weights and optimizer state are the obvious things to persist; the data loader is the part people forget, and getting it wrong means you silently re-show data (overweighting it) or skip data (never training on it). Neither throws an error. To resume at the exact byte the run died on, the checkpoint must capture:

• (a) Epoch + global RNG seed. The seed drives the shuffle (the base_seed + epoch from the DistributedSampler section). Without it you cannot reproduce the permutation, so you cannot know what order the remaining samples come in.
• (b) Per-rank shard cursor. Which shard file each rank was reading.
• (c) Within-shard sample offset. How many samples into that shard it had consumed.
• (d) Shuffle-buffer contents / RNG state. If you use the streaming approximate shuffle (last section), the buffer is mutable state — its current contents and the PRNG that draws from it must be restored, or the post-resume order diverges from what a non-interrupted run would have seen.
• (e) Count of in-flight prefetched batches. The number of batches sitting in the prefetch queue, produced but not yet stepped on.

That last item hides the real subtlety. With asynchronous prefetch, "samples consumed by the model" ≠ "samples pulled from disk." The loader's read head runs ahead of the optimizer by the full depth of the in-flight queue (W workers × prefetch_factor batches). If you checkpoint the loader's read offset, you over-count — you record positions for batches that were fetched but whose gradients never reached optimizer.step(). On resume you'd skip exactly those queued-but-unstepped batches.

The rule: checkpoint the offset of the last batch the optimizer actually stepped on, not the loader's read head. In practice you tag each batch with its global sample offset as it leaves the loader, and the checkpoint records the tag of the batch that was live at the step() call — so the in-flight tail is correctly discarded and re-fetched on resume. This is precisely the pipeline state that lesson 26 persists alongside the weights; a loader that cannot report "I am here, on the stepped batch" can only restart at an epoch boundary, throwing away progress on every failure.

The storage tier

Where the bytes physically live reshapes the whole pipeline. The three tiers, fastest-but-smallest last:

There is a free speedup hiding in the OS: the page cache. Epoch 1 is disk-bound because every shard is read cold from storage. If the working set fits in host RAM, later epochs hit the cached pages and the loader speeds up sharply — sometimes the difference between input-bound and compute-bound is just "is epoch 1 over yet." (Single-epoch pretraining never gets this gift; see the last section.)

The decisive move is to pre-tokenize offline into large binary shards — a .bin/.idx memmap pair, or WebDataset tar shards — so training reads are large sequential reads of token IDs, not millions of small random reads of raw text, and the CPU on the hot path does no tokenization. This collapses the bytes moved (token IDs are ~2 bytes each vs. ~4–5 bytes/char of raw UTF-8) and turns random I/O into streaming I/O, which is what every tier above is fastest at.

Before committing to a tier, run the demand-bandwidth check:

required GB/s = global_tokens_per_sec · bytes_per_token

Compute the left side from your global batch and step time, then ask whether the chosen tier sustains it split across every node. If demand exceeds the storage bandwidth, no number of workers helps — the wire is the bottleneck, and the answer is pre-tokenized shards plus an NVMe cache, not more CPU.

Worker & host budgeting

The starvation discussion gave the worker count to keep the GPU fed: W ≈ T_l / T_c. That formula sizes the number of data-loading processes, but it assumes the host can actually run them. The real ceiling is hardware: CPU cores per GPU and host RAM.

Each worker holds its own in-flight batches, so memory scales with the worker count and prefetch depth:

host RAM needed ≈ num_workers · prefetch_factor · batch_bytes

And this competes for the same RAM as the pinned-memory staging buffers that make the async H2D copy work (the page-locked buffers from lesson 18) — pinned pages are non-pageable, so they are RAM the OS cannot reclaim under pressure. Over-provision workers and you don't just waste cores; you can push the host into swap or OOM.

Rule of thumb: take the cores allotted per GPU (total host cores / GPUs-per-node), reserve a couple for the main training process and the NCCL/communication threads, and split the rest across workers. On an 8-GPU node with 96 cores that's ~12 cores/GPU; leaving headroom, ~8–10 workers per GPU is a sane starting point — then measure GPU utilization and adjust, rather than maxing the slider.

Shuffle quality at scale

You cannot load a web-scale dataset into RAM to shuffle it — it's trillions of tokens. Streaming loaders use an approximate shuffle: maintain a buffer of size B, draw a random element from it to emit, and refill the vacated slot from the stream. Larger B approaches a true global shuffle but costs more RAM and makes resume heavier (the buffer is checkpoint state — item (d) two sections up); smaller B only locally reorders.

The pretraining reality changes what "good shuffle" even means. Data is seen roughly once — a single epoch over trillions of tokens — so there is no per-epoch reshuffle to smooth out a weak ordering. What matters instead is the global shard ordering and inter-shard mixing: if the data is blocked by source (all of Wikipedia, then all of GitHub, then all of arXiv, contiguous), the model trains on long correlated runs, the gradient estimate stops looking i.i.d., and convergence suffers — a slow, expensive problem you can't reshuffle your way out of mid-run. Interleave sources at shard-build time and shuffle shard order globally; the in-buffer shuffle only polishes the local mix.