Data loading pipelines

Deep dive

How it works

Parts:

• shard store — large sequential files (flat token array, tar, TFRecord, MDS)
• sampler — seeded global permutation, rank-strided, with (epoch, cursor) state
• workers — processes or threads that decode, augment and tokenize
• bounded queue — 2-4 batches in flight
• collator and H2D copy — pinned memory, side stream

1. Offline preprocessing. Tokenize text once and write uint16 ids (vocab < 65536) or uint32 to one flat file per split, with an EOS token between documents. Images: pack into shards of order of 100 MB-1 GB (tar, TFRecord, MDS) so reads are large and sequential.
2. Global order. Seeded permutation per epoch: perm = shuffle(range(N), seed + epoch). Rank r of W takes perm[r::W]. Every rank computes the same permutation with no communication.
3. Sequence packing. Cut the concatenated token stream into windows of seq_len + 1. Input = window[:-1], target = window[1:]. Documents straddle windows. That is fine for pretraining, or add a document mask so attention does not cross EOS.
4. Workers. Decode and augment in separate processes or threads, a bounded queue between them and the trainer. 2-4 batches in flight are enough when steady.
5. Host to device. Collate into pinned (page-locked) host memory, copy with non_blocking=True on a side stream, so the copy overlaps the previous step.
6. Resume. Save (epoch, cursor) per rank with the checkpoint. On restart, regenerate the permutation and skip to cursor, no reading of skipped data.

disk/object store ─► shard reader ─► decode/augment workers ─► collate (pinned) ─► H2D copy ─► GPU step
      (large sequential reads)     (N processes)              (bounded queue)     (side stream)

Decision table

Code

use std::fs::File;
use std::os::unix::fs::FileExt;                   // pread: positional read, no shared file cursor
use std::sync::mpsc::sync_channel;
use std::thread;

/// Seeded Fisher-Yates over window indices; identical on every rank for the same (seed, epoch).
fn permutation(n: usize, seed: u64, epoch: u64) -> Vec<u32> {
    let mut s = seed ^ epoch.wrapping_mul(0x9E3779B97F4A7C15) | 1;
    let mut p: Vec<u32> = (0..n as u32).collect();
    for i in (1..n).rev() {
        s ^= s << 13; s ^= s >> 7; s ^= s << 17;
        p.swap(i, (s % (i as u64 + 1)) as usize);
    }
    p
}

#[derive(Clone, Copy, Debug)]
pub struct LoaderState { pub epoch: u64, pub cursor: usize }   // goes into the checkpoint

/// Batches of (input, target) token windows from a flat little-endian u16 token file.
pub fn spawn_loader(path: &str, seq: usize, batch: usize, rank: usize, world: usize, seed: u64,
                    mut st: LoaderState) -> std::sync::mpsc::Receiver<(Vec<u16>, Vec<u16>, LoaderState)> {
    let file = File::open(path).expect("open tokens");
    let n_tokens = file.metadata().unwrap().len() as usize / 2;
    let n_windows = (n_tokens - 1) / seq;                // non-overlapping windows of seq+1 tokens
    let (tx, rx) = sync_channel(4);                      // bounded: at most 4 batches in flight
    thread::spawn(move || {
        let mut buf = vec![0u8; (seq + 1) * 2];
        loop {
            let perm = permutation(n_windows, seed, st.epoch);
            let mine: Vec<u32> = perm.into_iter().skip(rank).step_by(world).collect();
            let per_rank = mine.len() / batch * batch;    // drop_last: equal steps on all ranks
            while st.cursor + batch <= per_rank {
                let (mut x, mut y) = (Vec::with_capacity(batch * seq), Vec::with_capacity(batch * seq));
                for &w in &mine[st.cursor..st.cursor + batch] {
                    file.read_exact_at(&mut buf, (w as usize * seq * 2) as u64).unwrap();
                    let toks: Vec<u16> = buf.chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]])).collect();
                    x.extend_from_slice(&toks[..seq]);
                    y.extend_from_slice(&toks[1..]);      // next-token targets, shifted by one
                }
                st.cursor += batch;
                if tx.send((x, y, st)).is_err() { return; }   // trainer dropped the receiver
            }
            st = LoaderState { epoch: st.epoch + 1, cursor: 0 };
        }
    });
    rx
}

fn main() {
    let path = std::env::temp_dir().join("tokens_demo.bin");
    let toks: Vec<u8> = (0..10_000u16).flat_map(|t| t.to_le_bytes()).collect();
    std::fs::write(&path, toks).unwrap();
    let p = path.to_str().unwrap().to_string();
    let rx = spawn_loader(&p, 64, 8, 0, 2, 1234, LoaderState { epoch: 0, cursor: 0 });
    let mut saved = None;
    for step in 0..5 {
        let (x, y, st) = rx.recv().unwrap();
        assert!(x.iter().zip(&y).all(|(a, b)| *b == *a + 1));  // targets are inputs shifted by 1
        if step == 2 { saved = Some(st); }
        println!("step {step}: first token {} state {:?}", x[0], st);
    }
    // Resume from the state saved after step 2: the next batch equals step 3 above.
    let rx2 = spawn_loader(&p, 64, 8, 0, 2, 1234, saved.unwrap());
    println!("resumed: first token {}", rx2.recv().unwrap().0[0]);
}

Efficiency

• Cost: tokens are 2 bytes each as u16, so 1 T tokens is 2 TB (arithmetic). A 2048-token window is a 4 KiB read; the GPU consumes order of 10⁵ tokens per second per device, which is under 1 MB/s, so text loading is easy if reads are mmap or sequential. Images are CPU-bound: JPEG decode plus augmentation costs order of milliseconds per image per core. The bound is storage or network throughput, or CPU decode, never the GPU.
• Levers:
  • Flat token binary with mmap and random offsets — no per-sample parsing or copies (nanoGPT) — LM pretraining.
  • Large sequential shards (tar, MDS) instead of small files — sequential reads, far fewer requests on object stores (WebDataset, Mosaic streaming) — remote storage.
  • Worker processes, prefetch, pinned memory, non-blocking host-to-device copies — overlaps load with compute (PyTorch docs) — always.
  • Pack documents to full length with attention masks — no padding tokens (Llama 3) — variable-length text.
• Naive → standard → state of the art:
  • One file open and decode per sample in the training process · GPU waits on I/O
  • Workers, prefetch, sharded files · loader hidden behind compute while workers keep up
  • mmap token binary, packed sequences, deterministic resumable streaming shards · bound by storage throughput

Implementation notes

• Positional reads (pread, or mmap) let many threads read one file without seeking. Random 128-1024 byte reads are fine when the token file sits in page cache or on NVMe. On network storage, read whole shards.
• Worker processes in Python each copy the dataset object. Keep it a file path plus offsets, not a loaded array, or memory multiplies by num_workers.
• Seed every worker's augmentation RNG from (base_seed, epoch, rank, worker_id). Forked workers otherwise inherit the same NumPy RNG state and apply identical augmentations.
• persistent_workers=True avoids re-forking workers every epoch, which can take seconds with large datasets.
• Pinned memory is a limited resource (page-locked). Pin only the collated batch buffers, and reuse them.
• Object storage: many parallel range GETs (order of 8-64 concurrent) per node. A single stream rarely exceeds order of 100 MB/s.
• Measure: time the loader alone (for batch in loader: pass) against the step time. The loader should be at least 1.5-2x faster to absorb jitter.

Pitfalls

• Shuffling only within a small buffer on sorted data (all of source A, then all of B): batches are not i.i.d., and loss curves show steps.
• Same permutation on every rank (seed without rank offset, or no DistributedSampler): every GPU trains on the same examples, an N-times smaller effective batch.
• Uneven per-rank batch counts at epoch end: one rank waits forever in an all-reduce. Use drop_last or pad.
• Resuming without the sampler state: the model sees the start of the epoch twice, a measurable loss dip and repeated data.
• Validation data inside the training stream through overlapping shards or near-duplicates. Deduplicate first (MinHash).
• uint16 token files with a vocabulary of 65536 or more: ids silently wrap.

Variants and alternatives

In production

nanoGPT (uint16 train.bin, memmap, random offsets), Megatron-LM (indexed .bin/.idx datasets, blended sources), PyTorch DataLoader + DistributedSampler, WebDataset, Mosaic streaming, tf.data, NVIDIA DALI, FFCV.

Implementations

Sources

• Murray, Šimša, Klimovic, Indyk, tf.data: A Machine Learning Data Processing Framework, VLDB 2021. paper · code
• Leclerc et al., FFCV: Accelerating Training by Removing Data Bottlenecks, CVPR 2023. paper · code
• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019. paper · code
• Llama Team, The Llama 3 Herd of Models, 2024 (document attention masks in packed sequences). paper
• PyTorch documentation, torch.utils.data.DataLoader. docs