Data loading and batching

Problem

A training loop wants one thing per step: a batch, a tensor of shape [B, ...] sitting in device memory. What it actually has is a dataset of individual examples: image files of different sizes, sentences of different lengths, rows in a table. Getting from one to the other raises four separate problems:

• Order. SGD assumes each batch is a fair random draw from the data. Reading files in their stored order (all cats, then all dogs) breaks that.
• Shape. A batch is one rectangular tensor, but examples differ in length. Something has to pad, crop or pack them.
• Speed. Decoding a JPEG or tokenizing a document runs on the CPU in Python. If it takes longer than the forward and backward pass, the GPU sits idle, and a faster GPU makes no difference.
• Batch size. Bigger batches are more efficient per example and give less noisy gradients, but they change how far each step should go and take more memory.

Take 12 token sequences with lengths 5, 23, 9, 31, 14, 7, 18, 4, 27, 11, 16, 8 and batch them in fours in the order they arrive. Each batch is padded to its longest sequence, so the model processes 304 token positions of which only 173 are real: 56.9% useful work. Nearly half the compute goes to padding.

Intuition

PyTorch splits the job into small pieces that each do one thing, so any one can be replaced:

• A Dataset answers "give me example i" (__getitem__) and "how many are there" (__len__). An IterableDataset is a stream instead: it only answers "give me the next one".
• A Sampler decides the order: it yields indices. RandomSampler is a fresh random permutation each epoch; SequentialSampler is 0, 1, 2, ….
• A BatchSampler groups those indices into lists of batch_size.
• A collate function turns a list of examples into one batch, usually by stacking tensors along a new first dimension.
• Workers are separate processes that run the fetch-and-collate step in parallel, so batch t+1, t+2, … are being prepared while the model trains on batch t.

The DataLoader wires these together. The main process owns the sampler, hands lists of indices to workers, and receives finished batches back. A worker only ever sees indices; it never decides what to load next. That is what keeps the order reproducible even when many processes load at once.

Padding is the second idea. A batch has to be rectangular, so short sequences get filler tokens that the model computes on and then masks out. How much filler depends entirely on which sequences end up together. Group similar lengths (bucketing) and each batch pads only a little. Drop the batch-of-rows idea entirely and concatenate sequences end to end into fixed-length rows (packing) and almost nothing is padding.

Mechanics

How the DataLoader is assembled

DataLoader.__init__ mostly validates arguments and fills in defaults. If you pass neither sampler nor batch_sampler, it builds a RandomSampler when shuffle=True and a SequentialSampler otherwise, then wraps it as BatchSampler(sampler, batch_size, drop_last). Passing your own sampler together with shuffle raises, and a custom batch_sampler excludes batch_size, shuffle, sampler and drop_last, because it already decides all four. With batching on, the default collate_fn is default_collate.

RandomSampler draws a torch.randperm of the dataset each epoch, so every example is seen exactly once per epoch (sampling without replacement). BatchSampler then cuts that stream into lists of batch_size; its __len__ is ⌈N / B⌉, or ⌊N / B⌋ with drop_last=True, which throws away the final short batch.

For each list of indices, a fetcher calls the dataset. _MapDatasetFetcher.fetch uses the dataset's __getitems__ if it has one (one call for the whole batch, useful when a database or file format can read many rows at once), otherwise [dataset[i] for i in indices], and passes the list to collate_fn.

Collation

default_collate calls a recursive collate that mirrors the structure of one example: a list of dicts becomes a dict of batched values, a list of tuples a list of batched fields (a namedtuple stays a namedtuple), a list of tensors one stacked tensor (Python and NumPy numbers become tensors too, strings stay a list). Stacking needs equal shapes; a list field of uneven length fails with "each element in list of batch should be of equal size" (line 208). Inside a worker, collate_tensor_fn stacks straight into a shared-memory tensor, so handing the batch back to the main process needs no extra copy.

For variable-length sequences you write your own collate_fn, usually around pad_sequence, which pads a list of [L_i, ...] tensors to the longest L with padding_value on the right or left (padding_side), and returns the lengths or an attention mask alongside so the model can ignore the filler.

Workers, prefetching and pinned memory

With num_workers=0, _get_iterator returns a single-process iterator: each next() fetches and collates in the training process itself, so loading and compute strictly take turns. With workers it returns _MultiProcessingDataLoaderIter, whose source comment (lines 798-815) draws the same queues as the diagram above. When the iterator starts, _reset primes the pipeline by sending prefetch_factor × num_workers index lists (prefetch_factor defaults to 2 once there are workers). _try_put_index hands lists to workers round robin and never lets more than that many be outstanding; each batch the loop consumes triggers one more.

Workers finish out of order, so _next_data holds early arrivals until the batch the sampler asked for next shows up. Results are always in sampler order unless you pass in_order=False, which trades reproducibility for not stalling on one slow worker.

Each worker loop sets torch.set_num_threads(1) so N workers don't oversubscribe the CPU, then seeds Python's random and torch with base_seed + worker_id, and NumPy from a hash of the same two numbers (lines 274-282). The base seed is drawn when the iterator is created from the loader's generator, so random augmentations differ between workers and between epochs, yet a fixed generator reproduces them. Use get_worker_info() inside an IterableDataset to give each worker its own slice of the stream; without it every worker yields the full stream and you get each example num_workers times.

With pin_memory=True a background thread copies each batch into page-locked host memory. The GPU's copy engine can read pinned memory directly, so batch.to("cuda", non_blocking=True) becomes an asynchronous copy that overlaps with compute already queued on the GPU (see GPU execution). From ordinary pageable memory the driver has to stage the copy first. persistent_workers=True keeps worker processes alive between epochs (__iter__ resets the old iterator instead of building a new one), which saves start-up time when each worker holds something expensive, like an open file or a tokenizer.

Sharding across processes

In data-parallel training each of R processes must see a different part of the data and all must take the same number of steps, or the gradient all-reduce deadlocks. DistributedSampler shuffles with the same seed on every rank (seed + epoch), then pads the index list by repeating its head until the length divides by R (or with drop_last=True cuts the tail), and gives rank r every R-th index starting at r (line 134). Because the shuffle seed includes the epoch, you must call set_epoch(e) before each epoch; forget it and every epoch replays the same order.

Bucketing and packing

Bucketing sorts (or roughly groups) examples by length, cuts batches from neighbours, and then shuffles the order of the batches so training still sees lengths in random order. It is a custom batch_sampler: the dataset and model don't change. The cost is a little randomness: examples of similar length always share batches.

Packing concatenates documents into rows of exactly the context length, which is how language-model pretraining is usually batched. Two details make it correct rather than merely fast. Without an attention mask, tokens of one document attend to the previous document in the same row; a block-diagonal causal mask (or a variable-length attention kernel given the document boundaries) keeps them apart. And position ids should restart at 0 for each document, or the second document is encoded as if it began mid-sentence. Some pipelines skip both and accept the cross-talk, treating the separator token as enough; that is a deliberate tradeoff, not the default.

Batch size, noise and the learning rate

The mini-batch gradient is an average of per-example gradients, so it is an unbiased estimate of the full-data gradient whose variance shrinks as the batch grows. Sampling B of N examples without replacement, with per-example gradient variance σ²:

Var[ĝ_B] = σ² / B · (N − B) / (N − 1)

Doubling the batch halves the variance (for B much smaller than N), which is the case for a bigger step. The linear scaling rule (Goyal et al.) says: when you multiply the batch size by k, multiply the learning rate by k, and warm the learning rate up over the first few epochs because the rule fails while the weights are changing fastest (see learning-rate schedules). The rule only holds up to a point: past a critical batch size that depends on the task and grows as training proceeds (McCandlish et al.), extra examples per step stop reducing the number of steps needed, and you are paying for compute that buys nothing. In practice the learning rate is re-tuned whenever the batch size changes by a large factor.

Worked example

Padding waste

The twelve lengths from the problem section, batch size 4. In arrival order:

Sorted by length first, then batched:

Packing into rows of 32 with first-fit decreasing (take the longest remaining sequence, put it in the first row with room, open a new row if none has room) gives 6 rows: [31], [27 + 5], [23 + 9], [18 + 14], [16 + 11 + 4], [8 + 7]. That is 6 × 32 = 192 slots for 173 real tokens, 90.1%. Same data, same model: the packed version does 1.58× less work than the arrival-order one.

Counting batches and shards

10 examples with batch size 3: len(loader) is ⌈10/3⌉ = 4 (the last batch holds one example) or, with drop_last=True, ⌊10/3⌋ = 3.

Now shard those 10 examples (unshuffled, for readability) across 4 ranks with DistributedSampler:

By default each rank gets ⌈10/4⌉ = 3 indices: the list is padded to 12 by repeating its first two entries, so ranks 2 and 3 see examples 0 and 1 a second time this epoch. With drop_last=True every rank gets 2 and examples 8 and 9 are skipped this epoch (with shuffling on, which two depends on the epoch). For evaluation, where duplicates bias the metric, this matters; for training it rarely does.

How many workers?

Suppose loading one batch takes a worker 40 ms and a training step takes 15 ms, and the epoch has 12 batches. In an idealized model (fixed costs, no start-up time, prefetch factor 2), the epoch takes:

With no workers the loop alternates: 12 × (40 + 15) = 660 ms. One worker overlaps loading with compute, but it can only produce a batch every 40 ms, so the GPU still waits. Loading keeps up once W workers produce a batch at least every 15 ms, that is W ≥ 40/15, so W = 3. Then the epoch is the first load plus 12 back-to-back steps: 40 + 12 × 15 = 220 ms, and a fourth worker changes nothing. The remaining idle time is waiting for the first batch.

Gradient noise against batch size

Take a tiny "dataset" of 8 per-example gradients for one weight: 0.9, -0.3, 1.6, 0.2, -1.1, 0.7, 2.0, 0.4. The full gradient is their mean, 0.55. A mini-batch of size B drawn without replacement estimates it with variance σ²/B · (N − B)/(N − 1):

Going from 1 to 2 to 4 cuts the standard deviation by a bit more than √2 each time (the finite-population factor helps); at B = N the estimate is exact. The test suite checks these numbers against brute-force enumeration of every possible batch.

Implementation

A bucketing batch sampler and a padding collate function for a token dataset. The sampler shuffles within coarse length buckets and then shuffles the batch order, so lengths in a batch are similar but the training order is still random:

import random
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset, Sampler

class Tokens(Dataset):
    def __init__(self, seqs):
        self.seqs = seqs                          # list of 1-D LongTensors
    def __len__(self):
        return len(self.seqs)
    def __getitem__(self, i):
        return self.seqs[i]

class BucketBatches(Sampler):
    def __init__(self, lengths, batch_size, bucket=100, seed=0):
        self.lengths, self.batch_size, self.bucket = lengths, batch_size, bucket
        self.seed, self.epoch = seed, 0
    def set_epoch(self, epoch):
        self.epoch = epoch
    def __iter__(self):
        rng = random.Random(self.seed + self.epoch)
        idx = list(range(len(self.lengths)))
        rng.shuffle(idx)
        chunk = self.bucket * self.batch_size
        batches = []
        for start in range(0, len(idx), chunk):
            block = sorted(idx[start:start + chunk], key=self.lengths.__getitem__)
            batches += [block[i:i + self.batch_size] for i in range(0, len(block), self.batch_size)]
        rng.shuffle(batches)
        return iter(batches)
    def __len__(self):
        return sum(-(-min(self.bucket * self.batch_size, len(self.lengths) - s) // self.batch_size)
                   for s in range(0, len(self.lengths), self.bucket * self.batch_size))

def pad_collate(seqs, pad_id=0):
    ids = pad_sequence(seqs, batch_first=True, padding_value=pad_id)
    lengths = torch.tensor([len(s) for s in seqs])
    mask = torch.arange(ids.shape[1])[None, :] < lengths[:, None]
    return ids, mask

data = Tokens(seqs)
sampler = BucketBatches([len(s) for s in seqs], batch_size=32)
loader = DataLoader(data, batch_sampler=sampler, collate_fn=pad_collate,
                    num_workers=4, pin_memory=True, persistent_workers=True)

for epoch in range(epochs):
    sampler.set_epoch(epoch)
    for ids, mask in loader:
        ids = ids.to("cuda", non_blocking=True)
        mask = mask.to("cuda", non_blocking=True)
        ...

A few things to check when a loader is slow or wrong:

• Is the GPU waiting? Time next(iter) separately from the step. If fetching a batch takes longer than a step, add workers (until CPU cores run out), move work like decoding out of Python, or store data preprocessed.
• Are workers duplicating data? An IterableDataset must shard itself with get_worker_info() and, under distributed training, by rank too.
• Is the order reproducible? Pass a seeded generator to the DataLoader and call set_epoch on distributed or custom samplers.
• Is memory growing? Each worker is a separate process holding its own copy of whatever the dataset object references; a large Python list of strings in the dataset is touched (and so copied, under fork) by every worker. Arrays or memory-mapped files share better.

Tradeoffs

• Map-style vs iterable. Map-style datasets get exact shuffling, sharding and len() for free but need random access. Streams (web-scale shards, databases) need an IterableDataset, which only approximates shuffling with a buffer and must shard itself.
• Padding vs bucketing vs packing. Padding is simplest and exact. Bucketing keeps examples separate and cuts most waste, at the cost of length-correlated batches. Packing wastes almost nothing but needs a mask and position resets to be exact, and splits or drops documents longer than the context.
• More workers. They hide loading time until loading is no longer the bottleneck; past that they only cost memory and CPU, and each one adds start-up time unless workers persist.
• Pinned memory. Faster, asynchronous copies, but page-locked memory can't be swapped out; pinning too much starves the rest of the system.
• drop_last. Keeps every batch the same shape (helpful for compiled graphs and for batch-statistics layers) at the cost of skipping a few examples each epoch.
• Large batches. Better hardware use and fewer steps, but each example does less work past the critical batch size, and very large batches need learning-rate warmup and retuning. When memory is the limit, lower precision or gradient accumulation gives the same effective batch.

Connections

References

• Goyal et al. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. 2017. (The linear scaling rule and warmup.)
• McCandlish, Kaplan, Amodei and the OpenAI Dota Team. An Empirical Model of Large-Batch Training. 2018. (The critical batch size.)
• Krell, Kosec, Perez and Fitzgibbon. Efficient Sequence Packing without Cross-contamination. 2021. (Packing with masks and position resets.)
• PyTorch v2.15.0: DataLoader, _MultiProcessingDataLoaderIter, sampler.py, collate.py, _worker_loop, DistributedSampler, pad_sequence.