PyTorch From First Principles
Chapter 10 of 16

DataLoader: Where Is the Training Loop Actually Waiting?

By AI, PyTorch

Explain this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Apply this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Here are two controlled training pipelines. They use the same batch size, the same machine, and the same synthetic post-batch workload. Each sample also carries the same nominal two-millisecond cost: in one pipeline that cost is waiting, while in the other it is fixed CPU work.

The tensor construction around that controlled cost is the same in both cases. What changes is the resource those two milliseconds consume.

Both are given the same treatment — raise num_workers from 0 to 8 — and measured the same way.

=== producer=sleep  2.0 ms/sample  batch=32  consumer=8 ms/batch ===
workers    ex/s  wait_mean  wait_p50  wait_p95  wait_max    work wait_frac
      0     414       69.8      69.8      70.4      70.7     7.5      0.90
      1     452       62.6      62.3      65.9      66.4     8.1      0.89
      2     890       27.3      25.3      56.8      57.7     8.6      0.76
      4    1769        8.4       2.1      30.3      30.5     9.7      0.46
      8    3049        0.9       0.9       2.7       3.5     9.6      0.09

=== producer=spin  2.0 ms/sample  batch=32  consumer=8 ms/batch ===
workers    ex/s  wait_mean  wait_p50  wait_p95  wait_max    work wait_frac
      0     454       62.8      62.1      66.5      89.1     7.8      0.89
      1     451       56.9      56.4      62.4      79.7    14.0      0.80
      2     456       49.5      47.7     103.8     107.2    20.7      0.71
      4     458       38.0      21.9     121.1     134.2    31.9      0.54
      8     449       28.1      24.0      88.2     110.2    43.2      0.39

The first pipeline got 7.4 times faster. The second did not get faster at all: 454 examples per second at zero workers, 449 at eight.

Now look at the column that most people would have watched. In the second table, wait_mean — how long the training loop sat blocked waiting for its next batch — fell from 62.8 ms to 28.1 ms. The fraction of each step spent waiting for data dropped from 0.89 to 0.39. By that measurement the input pipeline improved enormously. By the measurement that pays for the electricity, nothing happened.

The time did not disappear. It moved. The work column — time spent on the training step itself, after the batch arrived — rose from 7.8 ms to 43.2 ms. Eight worker processes competing for the same cores made the training step five times slower, and the step’s own slowness is what stopped it from waiting.

This chapter is about not being fooled by that, and the way not to be fooled by it is to stop treating DataLoader as a function with tuning parameters and start treating it as a system with stages you can measure.

Where we are

Chapter 5 left the model completely described. A module is a tree of registered children, parameters and buffers, and you can ask the framework what it thinks it owns rather than inferring it from the source. Four separate structures — Python reachability, registered state, autograd dependency, optimizer membership — can be inspected independently, and every one of them is checkable.

But every training loop in this book so far has begun with a line like this:

x, y = X_train, y_train

The data was already there. Already decoded, already batched, already the right dtype, already resident in memory, and available at exactly the moment the loop asked for it. That assumption is doing an enormous amount of quiet work, and real training does not get it for free. Before every single optimizer step, something has to produce the next batch.

So the question for this chapter:

What has to happen before the next batch reaches the training step, and where exactly is the system waiting when it arrives too slowly?

By the end you should be able to take a slow or unstable pipeline, split a training step into stages, measure the visible wait at each boundary, name the stage that cannot keep up, change one mechanism, and prove from end-to-end evidence whether the bottleneck moved or merely relocated.

The environment these numbers came from

Every measurement in this chapter was produced on one machine, and none of them are values you should expect to reproduce:

PyTorch      2.13.0
Python       3.12.3
OS           Linux (container), fork start method
CPU          1 vCPU
CUDA         not available
threads      torch.get_num_threads() == 1

One core is a deliberately harsh constraint, and it is why the two tables at the top of this chapter diverge so cleanly. A CPU-bound producer on a single core cannot possibly scale, and a latency-bound producer on a single core scales beautifully. On your eight-core laptop both curves will be gentler and the crossover will sit somewhere else.

That is the point rather than a caveat. The useful artifact here is the measurement procedure, not the numbers it produced on somebody else’s hardware.

The system you are actually training

The training diagram everyone draws looks like this:

batch → model → loss → backward → optimizer

The system that runs looks more like this, and the arrow in the middle is the one that matters:

PRODUCER

storage / source
      ↓
sample production
(__getitem__ or __iter__)
      ↓
read / decode / parse / transform
      ↓
collate / batch construction
      ↓
worker process
      │
      └──────── process boundary ────────→ main training process
                                              ↓
                                        prepared batch

and once a batch has crossed into the training process, it still has further to go:

CPU batch → host-to-device transfer → forward → backward → optimizer step

Two things produce work and one thing consumes it, with a queue between them. Every DataLoader argument that people treat as a magic speed knob is a property of some part of that picture:

Argument What it changes in the system
num_workers how many processes produce samples concurrently
prefetch_factor how much prepared work may sit queued ahead of the consumer
persistent_workers whether producer processes survive between iterators
collate_fn part of the production path, not something after it
batch_size fetches per step, collation events, and the consumer’s workload
pin_memory the kind of host memory a finished batch lives in

None of them can make a permanently slower producer outrun the consumer. They change concurrency, queue depth, process lifetime and memory placement. Holding that distinction is most of the diagnostic value in this chapter.

Four quantities that are not the same thing

Before measuring anything, separate four ideas that ordinary conversation collapses into “speed”:

PRODUCER THROUGHPUT   how fast batches can be prepared
CONSUMER THROUGHPUT   how fast training can use them
BUFFERING             how much temporary mismatch can be hidden
END-TO-END THROUGHPUT how many examples actually complete per second

If the producer is slower than the consumer, the consumer eventually waits, and no amount of queue depth prevents that: a buffer drains at the difference between the two rates. If the producer is faster than the consumer, more production concurrency achieves nothing except consuming resources — which is exactly the second table at the top of this chapter, where the extra resources came out of the consumer’s own budget.

Prefetching hides temporary mismatch. It cannot hide a permanent one.

Low accelerator utilization is a symptom with many causes. It is not a diagnosis, and it is not evidence that the input pipeline is slow.

The training loop may be idle because batches arrive late, because transfers are expensive, because the model’s kernels are too small to fill the device, because CPU-side Python between kernels is slow, or because the workload is simply too small. Those need different repairs. So do not begin from “utilization is low, therefore raise num_workers.” Begin from a measurement of where the step’s time actually goes.

Split the step

The smallest useful instrument separates the visible wait for the next batch from everything after it.

from time import perf_counter

t0 = perf_counter()
batch = next(iterator)      # how long the consumer was blocked
t1 = perf_counter()
train_step(batch)           # what it did once it had work
t2 = perf_counter()

That is the whole idea. Everything else is bookkeeping around it, and the bookkeeping matters because three details will otherwise ruin the measurement: startup is not steady state, averages hide stalls, and a ratio of two numbers is not a verdict.

Here is the harness used for every timing table in this chapter.

import time
import statistics
from time import perf_counter

import torch
from torch.utils.data import Dataset, DataLoader

# ~24,500 iterations of this loop cost about 1 ms of CPU on the reference
# machine. Calibrate it on yours before comparing any absolute numbers.
ITERS_PER_MS = 24_500

def spin(milliseconds):
    """Consume a fixed amount of CPU work. Needs a core; sleeping does not."""
    total = 0.0
    for i in range(int(milliseconds * ITERS_PER_MS)):
        total += i * 0.5
    return total

class Producer(Dataset):
    """A sample source whose cost is controllable and deliberately visible."""

    def __init__(self, size=8192, dim=128, cost_ms=2.0, mode="sleep"):
        self.size, self.dim = size, dim
        self.cost_ms, self.mode = cost_ms, mode

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        if self.mode == "sleep":
            time.sleep(self.cost_ms / 1000)   # waiting on something external
        else:
            spin(self.cost_ms)                # work that needs a CPU
        x = torch.randn(self.dim)
        y = torch.randint(0, 10, ()).long()
        return x, y

def measure(loader, batches=30, warmup=8, consume=lambda b: None):
    """Split each step into the visible wait for the batch and the work after it."""
    it = iter(loader)
    waits, works = [], []

    for i in range(batches + warmup):
        t0 = perf_counter()
        try:
            batch = next(it)
        except StopIteration:
            break
        t1 = perf_counter()
        consume(batch)
        t2 = perf_counter()

        if i >= warmup:
            waits.append(t1 - t0)
            works.append(t2 - t1)

    n = len(waits)
    total = sum(waits) + sum(works)
    ordered = sorted(waits)
    return {
        "batches": n,
        "wait_mean_ms": 1000 * statistics.fmean(waits),
        "wait_p50_ms": 1000 * statistics.median(waits),
        "wait_p95_ms": 1000 * ordered[min(n - 1, int(0.95 * n))],
        "wait_max_ms": 1000 * max(waits),
        "work_mean_ms": 1000 * statistics.fmean(works),
        "wait_fraction": sum(waits) / total,
        # Correct for the controlled benchmark below because every measured
# batch is full. General-purpose loader benchmarks should count the
# actual number of examples returned.
"examples_per_s": n * loader.batch_size / total,
    }

Two design decisions in there are worth defending.

spin consumes fixed work, not fixed wall-clock time. An earlier version of this harness busy-waited until perf_counter() reached a deadline. That is wrong for exactly the situation we want to study: when eight processes contend for one core, a two-millisecond deadline is satisfied while consuming far less than two milliseconds of CPU, so the “CPU-bound” producer quietly degrades into a sleeping one and the contention we were trying to measure disappears from the results. Fixed iteration counts do not have that problem. When the machine is oversubscribed, the same work simply takes longer, which is what we want to see.

sleep and spin are the two ends of a real spectrum. Sample production in practice is some mixture of waiting for storage or a network and burning CPU on decode, parsing and transforms. Those two components respond to worker concurrency in opposite ways, and separating them is the fastest way to build the right intuition. A JPEG decode is mostly spin. A read from object storage is mostly sleep. A small file on a cold network filesystem is mostly sleep with a spin tail.

The wait ratio is a question, not an answer

It is tempting to reduce the measurement to one number:

wait_fraction = data_seconds / (data_seconds + work_seconds)

and then to attach thresholds: below 0.1 healthy, above 0.5 broken. Resist that. Go back to the spin table at the top of the chapter and read the last two columns together:

workers    ex/s   work  wait_frac
      0     454    7.8       0.89
      8     449   43.2       0.39

The wait fraction more than halved. End-to-end throughput did not move. A metric that improves by 56% while the system does not change at all is not a health indicator; it is a ratio, and ratios move when either side moves.

The ratio is still worth computing, because it tells you which side of the boundary to investigate first. It is just not a verdict, and it cannot be compared across configurations without also comparing the throughput.

There is a second reason to distrust a single threshold, which is that the same number means different things at different scales. A pipeline spending 10% of its time waiting for data is worth several days of work if the run takes a week. A pipeline spending 30% of its time waiting is irrelevant if the experiment finishes in four minutes. The number acquires meaning from the job, not from a table in a book.

Compare the visible wait against useful work for this workload, then test whether reducing the larger component improves end-to-end examples per second. The throughput is the verdict; the ratio only chooses where to look.

num_workers changes concurrency, not speed

With num_workers=0, DataLoader performs sample fetching and collation in the main process, inside the next() call. There is no DataLoader worker process preparing future batches while that main thread is occupied.

That makes the mode easy to reason about: time spent synchronously producing the next batch directly delays the host from advancing the training loop. Other asynchronous work already issued elsewhere — for example CUDA work — is a separate question and should not be confused with DataLoader prefetch.

With num_workers greater than zero, DataLoader starts worker processes and moves data-production work across that process boundary.

For a map-style dataset, the main process’s sampler generates indices and dispatches those indices to workers. For an IterableDataset, each worker instead advances its own replica of the dataset iterator.

In both cases, workers can prepare data while the main training process consumes earlier batches, but the distinction between index dispatch and independent iterable replicas becomes crucial when we investigate duplication later. The consumer can then train on batch N while workers prepare N+1, N+2 and so on.

That is the entire mechanism, and it explains both tables.

When workers help

=== producer=sleep  2.0 ms/sample  batch=32  consumer=8 ms/batch ===
workers    ex/s  wait_mean  wait_p50  wait_p95  wait_max    work wait_frac
      0     414       69.8      69.8      70.4      70.7     7.5      0.90
      1     452       62.6      62.3      65.9      66.4     8.1      0.89
      2     890       27.3      25.3      56.8      57.7     8.6      0.76
      4    1769        8.4       2.1      30.3      30.5     9.7      0.46
      8    3049        0.9       0.9       2.7       3.5     9.6      0.09
     16    3106        1.1       0.3       4.2      11.3     9.2      0.11

Seven and a half times the throughput, on a machine with one core. That number is worth sitting with, because it directly contradicts the most common rule of thumb in this area. Worker count is not bounded by core count when the producer is waiting rather than computing — a sleeping process needs no CPU, and thirty-two sleeps of two milliseconds each can overlap almost perfectly.

Notice also where the scaling stops. Between 8 and 16 workers, throughput moves from 3049 to 3106, which is nothing. The reason is visible in the same row: wait_mean is already 0.9 ms and work is 9.6 ms. The producer has become faster than the consumer. A batch is always sitting in the queue when the training loop asks for one, so the loop’s rate is now set entirely by its own work:

32 examples / 10.5 ms per step ≈ 3050 examples/s

which is the measured value. Once you are on that side of the boundary, further producer concurrency buys nothing and costs memory and processes. Recognizing that state — near-zero wait, throughput matching consumer work — is how you know to stop tuning the loader and go look at the model.

When workers do not help

=== producer=spin  2.0 ms/sample  batch=32  consumer=8 ms/batch ===
workers    ex/s  wait_mean  wait_p50  wait_p95  wait_max    work wait_frac
      0     454       62.8      62.1      66.5      89.1     7.8      0.89
      1     451       56.9      56.4      62.4      79.7    14.0      0.80
      2     456       49.5      47.7     103.8     107.2    20.7      0.71
      4     458       38.0      21.9     121.1     134.2    31.9      0.54
      8     449       28.1      24.0      88.2     110.2    43.2      0.39
     16     441       19.5      16.4      52.6      52.8    53.1      0.27

Completely flat. And this one you can predict from first principles before running it, which is worth doing because it converts a mysterious result into arithmetic.

Producing one batch needs 32 samples × 2 ms = 64 ms of CPU. The training step needs another 8 ms. That is 72 ms of CPU work per batch, and the machine has one core, so:

32 examples / 72 ms ≈ 444 examples/s

Every row of that table sits between 441 and 458. Concurrency cannot create a second core. It rearranges when the single core does the work, and the arrangement is worse in two specific ways.

The work column is the first. As producers multiply, they take the core away from the training step, so the step itself slows from 7.8 ms to 53.1 ms. The consumer did not change. Its environment did.

The tail is the second, and it is uglier. At zero workers the maximum observed wait was 89 ms and the median was 62 ms — slow, but perfectly regular. At four workers the median dropped to 22 ms while p95 rose to 121 ms and the maximum to 135 ms. The pipeline became bursty: sometimes a batch is instantly available, sometimes the scheduler has not given the relevant worker any time and the loop stalls for six times the median. Averages went down. Predictability collapsed.

Averages hide starvation. Record the median, p95 and maximum wait alongside the mean, because a pipeline with a 3 ms mean and a 200 ms tail starves an accelerator in a way that no mean can show you.

The interpretation, and what to do with it

Two producers with identical nominal cost, identical machine, identical knob, opposite outcomes. The reasoning that separates them is not about DataLoader at all:

Concurrency can turn idle waiting into throughput. It cannot create more of a resource that is already saturated.

If the producer is waiting on something outside the process, more workers overlap more waiting, and scaling can continue far beyond core count. If the producer is computing, more workers compete for a fixed resource, and scaling stops at that resource — and then keeps going, downward, in the tail.

So the useful conclusion from a worker sweep is never “four is the right number.” It is a characterization:

OBSERVATION       the training loop blocks 62 ms per step waiting for batches
HYPOTHESIS        sample production cannot keep up
EXPERIMENT        sweep num_workers 0 → 1 → 2 → 4 → 8, recording
                  examples/s, wait mean/p95, and post-batch work time
RESULT            wait falls, work rises, examples/s is unchanged
INTERPRETATION    production is CPU-limited on this machine; the workers and
                  the training step are competing for the same cores
NEXT              decompose sample production into read/decode/transform/collate
                  and move deterministic work out of the hot path

That last step is where the real repair usually lives, and it is a different chapter’s territory only in part: Chapter 7 owns what the transforms mean, but the cost of running them is measurable here.

Thread oversubscription is a specific case of the same thing

The spin producer competes for the core in an obvious way. Real workers often do it invisibly, because the libraries they call have their own thread pools. A worker calling into a threaded linear-algebra backend, an image codec, or a numerical library may spawn as many threads as the machine has cores — inside each of eight worker processes.

8 worker processes × 8 library threads = 64 runnable threads on 8 cores

The symptom is the spin table: more workers, flat or falling throughput, inflated step time, ugly tails, and a machine that looks fully utilized while accomplishing little.

PyTorch’s own worker implementation already addresses one common form of oversubscription: in PyTorch 2.13, each DataLoader worker calls:

Treat this as an experiment, not a recommendation. It tests one hypothesis — that threads inside workers are oversubscribing the machine — and it does so cheaply. If throughput improves, you have learned something specific. If it does not, you have eliminated a suspect.

Two honest caveats. torch.set_num_threads(1) is not universally correct; a pipeline whose workers legitimately depend on multi-threaded tensor work can be slowed by it. And PyTorch does not control every thread pool in the process. Image and numerical libraries frequently manage their own, configured through their own APIs or through environment variables that must be set before the library is imported. If constraining PyTorch’s threads changes nothing, that is evidence about PyTorch’s threads, not about oversubscription in general.

Startup is not steady state

Every measurement above discarded the first eight batches. Here is why.

=== time to first batch vs steady-state batch ===
workers  first_batch_ms  steady_mean_ms
      0            71.6            69.9
      2           128.9            35.6
      8           279.7             6.6

With eight workers the first batch took 280 ms and a typical later batch took 6.6 ms — a factor of forty-two. A benchmark that timed only the first few batches would report that eight workers are four times worse than zero, which is the exact opposite of the truth.

The cost is real, though. Starting eight processes, importing modules in each, constructing eight dataset replicas and filling the prefetch queue takes time, and that time is not free just because it is transient. Whether it matters depends on the shape of the job:

startup latency        pay it once per iterator; dominates short jobs and
                       frequently re-created loaders
steady-state throughput  what a long training run is actually limited by
tail / stall behavior    what starves an accelerator even at good averages

These are three different quantities, and a change can improve one while damaging another. Adding workers here improved steady state by 10× and made startup 4× worse. For a week-long run that is obviously correct. For a validation loop that runs for twelve batches every epoch it may not be.

So: warm up before measuring steady state, and measure startup separately if your job pays it often.

persistent_workers is about lifecycle

By default, worker processes are created when an iterator is created and shut down when it is exhausted. A training loop that writes for batch in loader: once per epoch therefore pays the startup cost above once per epoch.

persistent_workers=True keeps them alive between iterators.

=== 5 short epochs, wall-clock seconds ===
  persistent_workers=False 1.17 s
  persistent_workers=True  0.73 s

Short epochs, four workers, startup paid once before timing began. Thirty-eight percent, from a setting that changes no computation whatsoever. The saving scales with how often you create iterators and how expensive worker startup is, which is why the benchmark has to span multiple epochs — measuring a single iterator would have shown nothing.

But this is not a free performance switch, because keeping a process alive keeps its state alive too.

import os

class Counting(Dataset):
    def __init__(self):
        self.calls = 0

    def __len__(self):
        return 4

    def __getitem__(self, i):
        self.calls += 1
        return torch.tensor([os.getpid() % 10000, self.calls])

persistent_workers=False
   epoch 0 (pid, calls) = [[766, 1], [766, 2], [766, 3], [766, 4]]
   epoch 1 (pid, calls) = [[770, 1], [770, 2], [770, 3], [770, 4]]
   epoch 2 (pid, calls) = [[774, 1], [774, 2], [774, 3], [774, 4]]
   parent dataset .calls = 0

persistent_workers=True
   epoch 0 (pid, calls) = [[778, 1], [778, 2], [778, 3], [778, 4]]
   epoch 1 (pid, calls) = [[778, 5], [778, 6], [778, 7], [778, 8]]
   epoch 2 (pid, calls) = [[778, 9], [778, 10], [778, 11], [778, 12]]
   parent dataset .calls = 0

Three things in that output deserve attention.

The process IDs change in the first block and do not in the second. That is the lifecycle, directly observed.

The counter resets in the first block and accumulates in the second. If a dataset caches, counts, advances a cursor, holds an open handle or adapts its behavior based on how much it has produced, persistent_workers=True changes what it does on epoch two. That is a semantic change delivered by a performance setting.

And parent dataset .calls is zero in both cases. The dataset object in your training script was never the object doing the work. Workers operate on their own replicas, and mutations there do not travel back. This is the single most common source of confusion about datasets under multiprocessing: instrumentation that accumulates into self works perfectly at num_workers=0 and silently reports nothing at num_workers=4.

The question to ask is therefore two questions:

Is repeated worker startup a measurable cost for this job, and is worker-local state safe to carry across epochs for this dataset?

prefetch_factor is queue depth

prefetch_factor controls how many batches each worker may prepare ahead of the consumer. It has real semantics worth confirming rather than assuming:

{'num_workers': 0} -> loader.prefetch_factor = None
{'num_workers': 2} -> loader.prefetch_factor = 2
num_workers=0 + prefetch_factor=2 -> ValueError prefetch_factor option could
    only be specified in multiprocessing.let num_workers > 0 to enable
    multiprocessing, otherwise set prefetch_factor to None.
num_workers=0 + persistent_workers=True -> ValueError persistent_workers
    option needs num_workers > 0

So the default is None, which resolves to 2 per worker when workers exist, and the argument is rejected outright when there are none. Both prefetch_factor and persistent_workers are meaningless without a worker process to apply them to, and PyTorch says so rather than ignoring you.

Now the part that matters. Deeper queues are widely believed to be faster:

=== prefetch_factor (sleep producer, 2 workers) ===
prefetch    ex/s  wait_mean  wait_p95  wait_max
       1     894       27.8      47.3      49.8
       2     903       26.1      53.2      54.2
       4     912       27.0      54.8      55.4
       8     906       27.3      56.1      57.2

Flat, within noise, across an eight-fold change in queue depth. The tail is very slightly worse at the top end.

This is exactly what the mechanism predicts. Prefetching hides variance in production latency: it lets a worker that got ahead cover for a worker that fell behind. This producer sleeps for precisely two milliseconds every time, so there is no variance to hide, and a deeper queue simply holds more batches in memory to no purpose. Two workers producing at a fixed rate below the consumer’s demand will starve the consumer identically whether the queue can hold two batches or sixteen.

Queue depth buys tolerance for irregular production. It cannot raise the average rate of a producer that is permanently too slow.

Increase it when the tail of your wait distribution is bad and the mean is fine — that is a variance problem, and it is the problem prefetching solves. Increase it blindly and you have bought memory pressure.

On which: the folk formula

in-flight batch memory ≈ num_workers × prefetch_factor × batch bytes

is a reasonable order-of-magnitude intuition for how much batch payload can be simultaneously alive, and it is not a RAM formula. It omits worker process overhead, the Python interpreter in each worker, dataset state in each replica, intermediate objects inside __getitem__ and collate_fn, and the copies involved in moving data between processes. Use it to understand why three settings all push memory in the same direction, not to predict a number.

If you want the payload part measured rather than estimated:

def tensor_payload_bytes(obj):
    if torch.is_tensor(obj):
        return obj.numel() * obj.element_size()
    if isinstance(obj, dict):
        return sum(tensor_payload_bytes(v) for v in obj.values())
    if isinstance(obj, (tuple, list)):
        return sum(tensor_payload_bytes(v) for v in obj)
    return 0

That measures the tensor payload reachable from one batch object. It is not process memory, and saying so plainly is the difference between a useful estimate and a misleading one.

When a pipeline runs out of shared memory or RAM, the three settings that drive in-flight data are num_workers, prefetch_factor and batch_size. Reduce them one at a time. If the failure disappears when you reduce one of them, you have evidence rather than a guess. Containers add one more thing to check, because shared-memory segments live in /dev/shm and container defaults for it are often far smaller than the host’s.

The producer includes collate_fn

A common and expensive mistake is to profile Dataset.__getitem__ and conclude that you have profiled data loading. Batch construction also involves stacking, padding, sorting, packing, type conversion and whatever Python you wrote to do it.

=== a cheap sample source with an expensive collate_fn ===
workers   collate    ex/s  wait_mean
      0   default    3389        1.9
      0      slow     857       29.8
      2   default    3080        2.7
      2      slow     831       18.7

The sample source here costs 0.05 ms per sample — 1.6 ms for a batch of 32. The collate_fn costs 30 ms. A profile of __getitem__ would report a fast, healthy dataset, and throughput drops by a factor of four.

Note also that two workers did not rescue it, for the reason established earlier: this collate is CPU work on a one-core machine, so moving it into a worker process moves it, and does not parallelize it.

Measure the stage, not the API name you suspect.

The general form of this, once the producer is known to be the bottleneck, is to decompose sample production into its parts and time them separately:

read → decode / parse → transform → collate

That decomposition works for images, text, audio, tabular data and remote object storage alike, and it usually finds one of two things: a stage that is much more expensive than expected, or a deterministic stage that is being recomputed every epoch for no reason. The second can be especially valuable because deterministic work may be a candidate for precomputation or caching rather than recomputation on every epoch.

Whether moving it out of the hot path is worthwhile depends on cache size, storage bandwidth, invalidation requirements and whether the transformation really is independent of training-time context.

Timing inside a worker needs care, for the reason the persistence experiment demonstrated: accumulating into self does not travel back to the main process. Return the timings with the sample instead.

class ProfiledDataset(Dataset):
    def __getitem__(self, index):
        t0 = perf_counter()
        raw = self.read(index)
        t1 = perf_counter()
        item = self.transform(raw)
        t2 = perf_counter()
        return {"x": item, "read_s": t1 - t0, "transform_s": t2 - t1}

The default collate will stack those floats into tensors, and the main process can aggregate them.

The process boundary

Chapter 5 was about disagreements between structures inside one process. Workers introduce a boundary between processes, and it changes several things that are easy to assume away: exceptions, dataset state, random number generators, file handles, database connections, caches, and how much of any of it is shared.

Start with the loud failure, because it has a clean procedure.

A worker crash, and the instrument that diagnoses it

Here is a dataset with one bad sample — index 37 decodes to None, which is what a corrupt file or a missing record looks like by the time it reaches your code.

def decode(i):
    if i == 37:
        return None
    return torch.randn(4)

class D(Dataset):
    def __len__(self):
        return 64

    def __getitem__(self, i):
        return decode(i) * 2.0

With workers, the failure arrives wrapped:

  File ".../torch/utils/data/dataloader.py", line 1567, in _process_data
    data.reraise()
  File ".../torch/_utils.py", line 794, in reraise
    raise exception
TypeError: Caught TypeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File ".../torch/utils/data/_utils/worker.py", line 374, in _worker_loop
    data = fetcher.fetch(index)
  File ".../torch/utils/data/_utils/fetch.py", line 54, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "probe9.py", line 13, in __getitem__
    return decode(i) * 2.0
TypeError: unsupported operand type(s) for *: 'NoneType' and 'float'

That is the good case. PyTorch caught the exception in the worker, serialized it, and re-raised it in the main process with the original traceback attached. You get the failing line.

The bad case looks like this, and every PyTorch user has met it:

  File ".../torch/utils/data/dataloader.py", line 1304, in _try_get_data
    raise RuntimeError(
RuntimeError: DataLoader worker (pid(s) 803) exited unexpectedly

That message appears when the worker process died rather than raised — killed by the OOM killer, terminated by a signal, or brought down by a segmentation fault inside a C extension. There was no Python exception to forward, so there is nothing to tell you.

The procedure is the same in both cases, and it is the diagnostic move this book keeps returning to: remove a boundary and see whether the failure survives.

reproduce with num_workers=0
        ↓
does the underlying dataset operation itself fail?
        ↓
yes  →  the bug is in sample production. You now have a direct traceback
        pointing at the failing line, and you can find the failing index.
no   →  the bug is in something the multiprocessing boundary introduced.
        Try num_workers=1 next: the smallest failing worker count is
        informative, and one worker eliminates worker-to-worker contention.

For the None sample above, num_workers=0 produces the same TypeError with a four-frame traceback and no wrapper. The failure survived removing the boundary, so multiprocessing was never involved — it only obscured the report.

When the failure does not survive, the suspects are specific: a library that is not safe across fork, a resource handle created in the parent and used in a child, a database or network client shared rather than opened per worker, a lock inherited in a locked state, memory pressure that only appears once state is multiplied by worker count, or — on spawn-based platforms — module-level code that re-executes on import in the child.

That last one is worth stating concretely because it is especially visible under the spawn start method.

Windows and macOS use spawn by default. Unix defaults depend on Python version and can also be configured explicitly: current PyTorch documentation notes fork for Python versions before 3.14 and forkserver from Python 3.14. On spawn, the child process imports your main module. If constructing the loader and starting iteration happens at module level, the child does it too, recursively.

def main():
    loader = DataLoader(MyDataset(), batch_size=64, num_workers=4)
    for batch in loader:
        train(batch)

if __name__ == "__main__":
    main()

The guard is not decoration. It is the thing that stops the child from re-running your program.

Worker-owned resources follow the same logic: create them lazily, inside the worker, on first use.

class RemoteDataset(Dataset):
    def __init__(self, keys):
        self.keys = keys
        self.client = None          # created per worker, on demand

    def _get_client(self):
        if self.client is None:
            self.client = open_client()
        return self.client

    def __getitem__(self, index):
        return self._get_client().fetch(self.keys[index])

Because each worker has its own dataset replica, this pattern makes each worker construct its own client lazily when it first needs one. No live client object is created in the parent and then inherited or serialized into the worker.

Whether that is sufficient still depends on the client library’s own multiprocessing guarantees. Whether that is sufficient depends on the library’s own multiprocessing guarantees, which you have to check rather than assume.

One related rule that is worth keeping even though it looks arbitrary: do not return CUDA tensors from workers. The reason is not that CUDA tensors are categorically forbidden in multiprocessing. The problem is that sharing CUDA tensors and CUDA state across process boundaries introduces additional lifetime, synchronization and multiprocessing constraints that ordinary input loading should not require you to manage.

PyTorch therefore recommends the simpler multi-process loading architecture: The architecture that stays simple is the obvious one:

workers produce CPU batches → main process moves them to the device → model

Dataset state under fork

You will see the claim that every worker gets a full copy of the dataset in RAM. That is too strong to be useful, and it is platform-dependent.

On fork-based systems, the child initially shares the parent’s memory pages copy-on-write, so a large read-only structure is not immediately duplicated. But CPython’s reference counts live in the object headers, so merely touching Python objects writes to their pages and causes them to be copied. A large Python list of a million path strings will therefore be progressively duplicated as workers walk it, which looks like a slow memory leak that scales with worker count. On spawn-based systems the dataset is serialized to each child instead, and the cost is paid up front.

The claim that survives both platforms is narrower and still actionable:

Large Python object graphs attached to a dataset become expensive when multiplied across worker processes, in a way that depends on the start method. Measure the scaling on your target platform rather than reasoning about it.

The repair, when it is needed, is to make the state compact rather than to abandon workers: arrays instead of lists of boxed objects, memory-mapped files, columnar formats, or an index that is loaded per worker instead of inherited.

Concurrency can change correctness

Everything so far has been about speed. Here is a failure where enabling workers changes what the program computes, with no exception and no warning.

A map-style dataset exposes keyed or indexed sample access. With multi-process loading, the main process’s sampler generates the indices and dispatches them to workers.

That means workers do not independently traverse the whole dataset, so simply increasing num_workers does not multiply the dataset in the way our naive IterableDataset will.

The sampler may still deliberately yield the same index more than once — for example under replacement sampling or a custom sampling policy. In that case the duplication comes from the sampler, not from every worker independently replaying the dataset.

An IterableDataset is a stream. It implements __iter__ and produces samples in whatever order it likes, which is what you want for data that does not fit in memory, arrives from a network, or has no meaningful length. Each worker gets its own replica of the dataset — and calls __iter__ on it.

from torch.utils.data import IterableDataset

class Naive(IterableDataset):
    def __init__(self, n):
        self.n = n

    def __iter__(self):
        yield from range(self.n)

Twelve items. Count what comes out.

num_workers=0: emitted=12 unique=12 max_copies=1
   order: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
num_workers=1: emitted=12 unique=12 max_copies=1
   order: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
num_workers=2: emitted=24 unique=12 max_copies=2
   order: [0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7, 8, 9, 10, 11, 8, 9, 10, 11]
num_workers=3: emitted=36 unique=12 max_copies=3
   order: [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, ...]

With two workers, an “epoch” contains every example twice. With three, three times. The dataset was asked to produce its stream once per worker, and it obligingly did, because nothing in __iter__ knows that it is one of several replicas.

Nothing raised. The loss went down. The epoch took the expected amount of time — in fact it looked faster per epoch in examples-per-second terms, because it was processing more examples. Every number on the dashboard is plausible. The model is being trained on a dataset that is silently k times its actual size, with each example repeated k times per epoch, which changes the effective sampling distribution, the meaning of “epoch”, the comparability of every run against every other run, and the relationship between the training and validation sets if the split was ever performed downstream.

This is Chapter 5’s lesson in a different system. There, four structures disagreed about what belonged to the model. Here, the number of workers changes what a pass over the data means, and only one of those meanings is the one you wrote down.

When concurrency is added to a producer, verify that the parallel producers partition the work rather than reproduce it. The evidence is a count of unique samples per epoch, not a loss curve.

Note that num_workers=1 is clean, which makes it a poor test. The bug requires two producers to be visible, and a pipeline validated at one worker and deployed at eight will be wrong.

The repair

get_worker_info() returns None in the main process and a record describing the current worker inside one.

from torch.utils.data import get_worker_info

class Sharded(IterableDataset):
    def __init__(self, n):
        self.n = n

    def __iter__(self):
        info = get_worker_info()
        if info is None:
            yield from range(self.n)
        else:
            yield from range(info.id, self.n, info.num_workers)

Each worker takes a strided slice of the stream. Worker 0 takes items 0, 3, 6, 9; worker 1 takes 1, 4, 7, 10; and so on. The slices are disjoint and cover everything.

workers=0: emitted=12 unique=12 max_copies=1
   order: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
workers=2: emitted=12 unique=12 max_copies=1
   order: [0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 9, 11]
workers=3: emitted=12 unique=12 max_copies=1
   order: [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]
workers=5: emitted=12 unique=12 max_copies=1
   order: [0, 5, 10, 1, 6, 11, 2, 7, 3, 8, 4, 9]

Twelve samples, twelve unique, at every worker count including one that does not divide evenly. The range(start, stop, step) form handles the uneven case without arithmetic on your part, which is worth preferring over computing contiguous block boundaries by hand.

The interleaved output order is worth noticing rather than fixing. This particular run produced an interleaving of the worker shards rather than the original stream order.

Do not make correctness depend on a presumed round-robin worker schedule. Current DataLoader uses in_order=True by default for multi-process loading, which enforces its expected first-in/first-out batch ordering. Setting in_order=False explicitly allows batches to be returned out of order and can affect reproducibility and observed data distribution. If your data has meaningful sequence, striding across workers destroys it and you want contiguous shards instead — which is a decision about your data, made deliberately, rather than something to discover later.

One more consequence, easy to miss. Batches are formed within a worker, from that worker’s shard:

loader = DataLoader(Report(12), batch_size=6, num_workers=3)

   [[0, 0], [0, 3], [0, 6], [0, 9]]
   [[1, 1], [1, 4], [1, 7], [1, 10]]
   [[2, 2], [2, 5], [2, 8], [2, 11]]

Batch size 6 was requested. Every batch has four elements, and the first column shows why: each worker holds four items, so each produces one short batch. With an iterable dataset and multiple workers, several worker replicas can each end with an incomplete batch. That means you may see several short batches rather than the one short final batch you might expect from a map-style dataset.

There is a second consequence: with drop_last=True, PyTorch drops the last incomplete batch of each worker’s iterable-dataset replica. More than one partial batch can therefore disappear.

This is also why len(loader) is only an estimate for some multi-worker IterableDataset configurations: PyTorch cannot generally know how user-defined worker sharding will break the stream into batches. Code that assumes a fixed batch size — a reshape, a hard-coded dimension, a metric divided by a constant — will meet this eventually.

Worker randomness, and a piece of folklore worth re-checking

Random transforms inside workers raise an obvious question: do all the workers produce the same randomness?

   idx wid  torch.randint  random.randrange
   [0, 0, 253906, 223507]
   [1, 0, 636010, 838017]
   [2, 0, 320804, 289877]
   [3, 0, 768013,  59890]
   [4, 1, 827724, 810954]
   [5, 1, 885002, 852360]
   [6, 1, 105724, 171360]
   [7, 1,  32613, 673038]

In the PyTorch 2.13 implementation tested for this chapter, the worker loop seeds PyTorch and Python’s random from worker-specific state and also initializes NumPy’s legacy global RNG when NumPy is available. That explains the different values observed above without a custom worker_init_fn.

Treat that as an observation about the implementation we actually tested, not as a reason to stop thinking about worker randomness. PyTorch’s public reproducibility guidance still recommends using the worker’s PyTorch seed to initialize the random-number generators your pipeline depends on.

The durable rule is:

Know which random generators your pipeline actually uses, and make their worker-specific initialization explicit when reproducibility matters. Re-running the same loader with such a function changes the random column but does not fix a duplication bug, because there was not one.

Do not read that as “seeding is handled.” Read it as: the set of generators that PyTorch seeds is specific, and any generator outside that set is yours to manage. An image library with its own internal RNG, a NumPy Generator object, a C extension with independent seeding, or a third-party augmentation package may sit outside the random generators you have controlled.

The resulting failure is not always identical augmentations. It may be duplicated streams, unexpected correlations, irreproducible runs or simply randomness whose seed you cannot account for.

Inspect and seed the generator that actually produces the randomness.

The check is direct: emit whatever your transform’s randomness produced alongside the sample, and compare across worker_id. If two workers report the same values for different indices, that generator is not being seeded per worker.

Reproducibility across runs is a separate axis, controlled by passing an explicit generator:

g = torch.Generator()
g.manual_seed(1234)
loader = DataLoader(dataset, batch_size=64, num_workers=4, generator=g)

For a pipeline whose relevant randomness is controlled by this generator and correctly initialized worker-local generators, successive epochs can advance through different random states while repeated runs begin from the same reproducible setup.

That does not automatically make the entire training run deterministic. External random generators, input ordering, nondeterministic kernels and other sources of nondeterminism still need to be controlled separately.

There is a performance reason to care, too. Comparing loader configurations while the augmentation pipeline is drawing from a different random stream in each run means comparing two different workloads, and the timing difference you attribute to prefetch_factor may be the difference between a run that happened to draw many expensive augmentations and one that did not.

Host to device, and timing it honestly

Everything so far runs on any machine. This section is about CUDA, and it is written so that the reasoning survives even if you have no GPU to hand.

Once a batch exists in the training process, it still has to reach the accelerator. pin_memory=True asks DataLoader to place recognized CPU tensors in page-locked host memory, which can make host-to-CUDA transfers faster.

Keep two kinds of asynchrony separate.

non_blocking=True can avoid host-side synchronization after a transfer is issued even when the source began in ordinary pageable memory.

Actually overlapping a host-to-device copy with GPU kernel execution is a stronger claim. Current PyTorch guidance identifies three requirements for that form of overlap: pinned source memory, hardware with available DMA capability, and scheduling the transfer on a separate non-default CUDA stream.

The claim to avoid is that this makes training faster. Three separate questions decide whether it can:

  1. Is the host-to-device copy a meaningful fraction of the step? If reading and decoding cost 50 ms and the copy costs 2 ms, then a perfect copy optimization buys 4%, and the answer to “why didn’t pin_memory help?” is that there was nothing there to win.
  2. Are the batches in a form that benefits? Pinning applies to the tensors the loader produces. A batch containing Python objects, or tensors constructed after collation in the main process, may not be pinned at all.
  3. Does the schedule permit overlap? This is the one people skip.

On that third point: non_blocking=True allows the copy to be issued without blocking the host, but the copy still has to finish before anything that reads the tensor can run. If the very next thing you do is a forward pass on the copied batch, the dependency forces a wait, and you have an asynchronous copy that overlaps nothing. Useful overlap requires there to be independent work available while the transfer is in flight — the classic arrangement being to issue the copy for batch N+1 before computing on batch N.

The batch inspector answers question 2 directly:

def inspect_batch(batch, prefix="batch"):
    if torch.is_tensor(batch):
        pinned = batch.is_pinned() if batch.device.type == "cpu" else None
        print(f"  {prefix:12s} shape={tuple(batch.shape)} dtype={batch.dtype} "
              f"device={batch.device} pinned={pinned}")
    elif isinstance(batch, dict):
        for k, v in batch.items():
            inspect_batch(v, f"{prefix}.{k}")
    elif isinstance(batch, (list, tuple)):
        for i, v in enumerate(batch):
            inspect_batch(v, f"{prefix}[{i}]")
    else:
        print(f"  {prefix:12s} {type(batch).__name__} (not a tensor)")

  batch.x      shape=(4, 3) dtype=torch.float32 device=cpu pinned=False
  batch.id     shape=(4,) dtype=torch.int64 device=cpu pinned=False
  batch.name[0] str (not a tensor)

That output is from the reference machine, which has no accelerator, and it demonstrates the other half of the point. Asking for pinned memory on a machine without a device gets you a warning and unpinned tensors:

UserWarning: 'pin_memory' argument is set as true but no accelerator is found,
then device pinned memory won't be used.

The name field is worth a glance too: the default collate turned a list of strings into a list of strings rather than a tensor. It is not pinned, it is not transferable, and it will be quietly carried along by every batch.

Wall-clock time around CUDA code can lie

This matters more than any loader setting, because it invalidates measurements rather than degrading them.

CUDA kernel launches are often asynchronous with respect to the host. A naive wall-clock interval around CUDA code therefore does not necessarily measure completion of the device work. It may measure mostly host launch time, plus any synchronization that happened to occur inside the measured region.

start = perf_counter()
output = model(x)            # returns as soon as the kernels are queued
print(perf_counter() - start)

There are two correct approaches, and they measure different things.

To measure a host-visible interval that includes device work, synchronize at both ends:

torch.cuda.synchronize()
start = perf_counter()
output = model(x)
torch.cuda.synchronize()
elapsed = perf_counter() - start

To measure elapsed time between two points in a CUDA stream, use CUDA events. The events are recorded into the stream and timestamped by the device:

def time_cuda(fn):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    fn()
    end.record()
    torch.cuda.synchronize()      # wait until `end` has actually occurred
    return start.elapsed_time(end)   # milliseconds

The difference is not pedantry. perf_counter with synchronization tells you elapsed wall-clock time including host-side Python, launch overhead and any queueing. CUDA events tell you the device-side duration between two points in the stream. When the synchronized wall-clock interval and CUDA-event interval disagree substantially, the host and device timelines are telling different stories.

The difference can reflect host launch overhead, scheduling, synchronization, queueing or overlap. Do not interpret the numerical gap itself as a direct measurement of one single category of work.

For the measure() harness, this means one thing: if the consumer runs on CUDA, synchronize before the t2 timestamp, or the “work” column will be a measurement of how fast Python can queue kernels, and every batch will appear to compute in under a millisecond while the loop mysteriously fails to speed up.

Profiling comes after the simple measurements

The PyTorch profiler is genuinely useful, and it is the wrong first tool. A trace of an unexamined training loop is a large amount of information about a question you have not yet asked. The sequence that works is:

simple timing  →  identify the suspicious boundary  →  decompose that stage
               →  profiler when you need finer attribution

When you do reach for it, name the ranges after the mental model, so the timeline becomes evidence about the system you already understand:

from torch.profiler import profile, record_function, ProfilerActivity

activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
    activities.append(ProfilerActivity.CUDA)

with profile(activities=activities) as prof:
    it = iter(loader)
    for step in range(20):
        with record_function("data_wait"):
            x, y = next(it)
        with record_function("host_to_device"):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
        with record_function("forward"):
            logits = model(x)
            loss = loss_fn(logits, y)
        with record_function("backward"):
            loss.backward()
        with record_function("optimizer"):
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=20))

Those five names are the stages from the diagram at the start of the chapter. A trace labelled that way answers “which stage” directly, which is what you came for.

Using AI on an input pipeline

Ask an assistant for a fast DataLoader configuration and you will get one. It will probably look like this, and it may well be excellent:

DataLoader(dataset, batch_size=128, num_workers=8, pin_memory=True,
           prefetch_factor=4, persistent_workers=True)

The worker sweep already shows why such a recommendation cannot be judged from its appearance. On this one-core CPU-bound producer, increasing worker count leaves end-to-end throughput around 450 examples per second while reducing visible batch wait and substantially increasing post-batch step time. The waits also become much more bursty.

That is enough evidence to reject “more workers must be better” for this pipeline. The chapter has not established that every other option in the suggested configuration — larger batches, deeper prefetching, pinning and persistent workers — would produce the same effects, so those remain separate hypotheses to test. Nothing about the suggestion was wrong. It was an answer to a question about typical pipelines, offered to someone who had not established which stage was limiting theirs.

The move that works is the one this book keeps using: make the assistant produce a falsifiable model of your program before it produces a configuration.

Here are measurements from a PyTorch training run. Do not recommend a
configuration yet.

  batch size / num_workers / prefetch_factor / persistent_workers / pin_memory
  mean, median, p95 and max next(loader) wait
  mean post-batch step time
  end-to-end examples/sec
  CPU utilization, core count, and whether workers call threaded libraries
  RAM and shared-memory usage
  device, and host-to-device time if measured
  whether these are steady-state numbers or include startup

1. Which stage is the strongest bottleneck hypothesis, and which specific
   numbers support it?
2. Give two alternative explanations that would produce similar symptoms.
3. Propose the smallest experiment that distinguishes the leading explanation
   from those alternatives. It should change one mechanism.
4. State what each measurement should do if the leading hypothesis is correct,
   and what it should do if each alternative is correct instead.

I will run the experiment and report the numbers. Only then propose a change.

Point 4 is the one that converts the exchange from advice into science, because it is what makes the assistant’s model wrong in a way you can detect. “Raise num_workers and see” is unfalsifiable — any outcome can be rationalized afterwards. “If production is latency-bound, doubling workers should roughly halve the mean wait while leaving step time unchanged; if it is CPU-bound, the mean wait will fall while step time rises and examples/sec stays flat” is a prediction that the two tables at the top of this chapter would have decided in ninety seconds.

Ask for a falsifiable model of the program before asking for a rewritten configuration.

There is a specific trap here worth naming. Assistants are good at the knobs and much weaker at whether your IterableDataset shards correctly, because the duplication bug does not look like a bug in the source. Ask about the knobs and you will get knobs. Ask “how many unique samples should one epoch of this dataset contain, and what code guarantees that with four workers?” and the question has a checkable answer.

What you should now be able to answer

Here is a pipeline of the kind you will be handed. It runs, it produces a falling loss, and someone has already tuned it. Work through it before reading on.

class Shards(IterableDataset):
    def __init__(self, paths):
        self.paths = paths
        self.seen = 0

    def __iter__(self):
        for path in self.paths:
            for record in read_records(path):     # ~0.5 ms, mostly disk wait
                self.seen += 1
                yield augment(decode(record))     # ~3 ms of CPU

loader = DataLoader(Shards(paths), batch_size=64, num_workers=8,
                    prefetch_factor=4, persistent_workers=True,
                    pin_memory=True)

for epoch in range(20):
    for x, y in loader:
        train_step(x.to("cuda", non_blocking=True))
    print(f"epoch {epoch}: {loader.dataset.seen} records seen")

How many unique examples does one epoch contain? Not the number you would get from the file contents. __iter__ has no worker awareness, so each of the eight workers iterates every shard, and the epoch contains eight copies of the dataset. The loss will fall, the epoch will take a plausible amount of time, and nothing will raise.

What does the seen counter print? Zero, every epoch. self.seen is incremented in the worker replicas, and loader.dataset in the main process is the original object, which no worker ever touched. This is the same effect the persistence experiment showed: the parent’s .calls stayed at zero while workers counted to twelve.

Would persistent_workers=True change that counter? Not in the main process. It would make the workers’ own counters accumulate across epochs rather than resetting, which is invisible from here and is exactly the kind of worker-local state that makes persistence a semantic decision rather than a performance one.

Is num_workers=8 justified? Unknown from the code, and the per-sample costs suggest it needs testing rather than assuming. Production is roughly 0.5 ms of waiting and 3 ms of CPU per sample. The CPU-heavy portion means available CPU capacity will eventually limit scaling, while the 0.5 ms waiting component means some additional concurrency may still be useful.

The point at which workers begin competing destructively with the training process cannot be read from the code alone. That is exactly what the worker sweep is for. The evidence that would settle it is a sweep recording examples/sec, wait percentiles and post-batch step time together.

What does prefetch_factor=4 buy? With eight workers and prefetch_factor=4, PyTorch allows up to 32 batches to be prefetched across the workers before consumption.

That is a statement about DataLoader’s configured prefetch depth, not an exact count of every batch object, intermediate allocation or byte simultaneously alive elsewhere in the pipeline. If production latency is steady, the measurement in this chapter suggests that depth buys nothing over 2 and costs memory. If shard boundaries cause periodic stalls — a new file being opened, a decompression starting — then the deep queue is covering real variance and is earning its memory. A large p95-to-median gap is evidence that batch availability is bursty. It does not by itself prove that a deeper queue will remove the cause; that remains an experiment.

Is pin_memory=True helping? Three checks, and none of them is “it is enabled.” Confirm the batch is actually pinned with is_pinned(). Confirm the host-to-device copy is a meaningful fraction of the step, measured with events or with synchronization at both ends. Confirm what benefit you are testing. The forward computation for this batch must wait until this batch’s transfer is usable, so this loop does not demonstrate pipelined GPU copy/compute overlap.

non_blocking=True may still reduce host-side synchronization overhead. To test genuine copy/compute overlap, you would need an explicitly pipelined schedule — typically transferring a future batch on an appropriate separate CUDA stream while independent computation proceeds on the current batch — and then measure the end-to-end result.

The loop is slow. What do you measure first? Split the step. Record the wait for next() and the time after it, with warmup discarded, and keep mean, median, p95 and max. That first measurement tells you which side of the batch-arrival boundary currently dominates the visible step: waiting for the next batch, or work after the batch arrives.

If the second side dominates, it still needs decomposing — host-to-device transfer, CPU-side training code, forward, backward, optimizer and synchronization are different possible limits., and everything else in the investigation depends on which of those is true.

Suppose the wait is 40 ms and the step is 12 ms. What is the next experiment? Not eight more workers. Reproduce with num_workers=0 to get a clean measurement of unconcurrent production cost, then decompose it: read, decode, augment, collate. Three milliseconds of CPU per sample against 0.5 ms of waiting says most of the cost is in decode and augment, and whether decode is deterministic determines whether precomputing or caching its output is even a candidate. Whether doing so is worthwhile then depends on storage cost, cache size, invalidation and the rest of the pipeline.

A worker dies with exited unexpectedly on epoch 3 and never on epoch 1. What does that suggest? Accumulating worker-local state becomes a strong hypothesis, especially with persistent_workers=True, but it is not the only explanation. A particular late sample, external service failure or nondeterministic native crash could produce similar timing.

Test the accumulation hypothesis by disabling persistence and observing whether the failure moves or disappears. With persistent_workers=True the same processes have been alive for three epochs, so worker-local state — a cache that grows, handles that are never closed, memory that is never released — is a strong candidate. Test by setting persistent_workers=False: if the failure moves or disappears, the state was carrying across epochs. Also reduce num_workers, prefetch_factor and batch_size one at a time, since all three drive memory in flight, and check /dev/shm sizing if this is a container.

What evidence would prove a repair worked? End-to-end examples per second, measured at steady state, over a fixed number of batches, with the wait distribution and post-batch step time recorded alongside it. A fall in mean wait is not sufficient — the spin table halved its wait fraction while accomplishing nothing. And for the sharding repair specifically, the evidence is a count: unique samples per epoch, equal to the dataset size, at more than one worker count.

Exercises

These are written to be run, and they map onto the notebook that accompanies this chapter.

  1. Calibrate, then reproduce both curves. Measure ITERS_PER_MS on your machine, then run the worker sweep for the sleep and spin producers. Before running, predict where each curve stops improving, using your core count and the arithmetic in this chapter. Report where your predictions were wrong and why.

  2. Find the crossover. Build a producer that spends a fraction f of its per-sample cost sleeping and the rest spinning. Sweep f from 0 to 1 and, for each, find the worker count beyond which throughput stops improving. Plot that count against f and explain the shape.

  3. Make the wait ratio lie. Construct a configuration change that improves wait_fraction by at least 30% while leaving examples per second unchanged or worse. Then write the two-line summary you would send a colleague who reported the improvement as a win.

  4. Startup versus steady state. For worker counts 0 through 8, measure time to first batch and steady-state batch time separately. Then compute total wall-clock time for a job of 20 batches, 200 batches and 20,000 batches, and identify the worker count that wins each. Explain why the answer changes.

  5. Duplicate an epoch. Write an IterableDataset over a text file and confirm the duplication at 2, 3 and 4 workers by counting unique samples. Then repair it with get_worker_info(). Then repair it a second way, with contiguous shards instead of striding, and state which of the two you would use for sequential data and why.

  6. Short batches. Using the sharded iterable dataset, request batch_size=64 with 4 workers over 100 items and record the actual size of every batch. Explain the distribution. Then find one line in a training loop you have written that would break on it.

  7. Two crashes, two procedures. Build one dataset that raises a Python exception on a specific index, and one whose worker dies without raising. For each, run at num_workers=4 and then at 0, and record what you learn from each of the four runs. Write down the decision procedure you would follow next time.

  8. Persistent state. Give a dataset a worker-local cache that grows with every sample. Run five epochs with persistent_workers both False and True, recording per-epoch time and process memory. Identify the epoch at which persistence stops being a saving.

  9. Profile the wrong thing first. Take a pipeline with a cheap __getitem__ and an expensive collate_fn. Profile __getitem__ in isolation and record the conclusion you would have drawn. Then measure the full batch production time and account for the difference.

  10. Audit a pipeline you did not write. Take a DataLoader setup from a repository. Without changing it, answer: what is the wait distribution, what is the post-batch step time, does the sample count per epoch match the dataset size at more than one worker count, and which of its settings has any measured evidence behind it?

Next: what is inside the batch?

The pipeline is now a system rather than a function call. Something produces samples, something groups them into batches, worker processes may prepare future work concurrently, prepared batches wait in a queue, and the training loop consumes them. Every argument that looked like a tuning knob is a property of one of those stages, and each stage can be measured separately.

More usefully, the failure modes now have addresses. A slow pipeline is a stage that cannot keep up, and a worker sweep plus a wait distribution tells you which. A crashing pipeline is either a bad sample or a bad boundary, and num_workers=0 distinguishes them in one run. A pipeline that got faster without getting better is a measurement that moved rather than a system that changed, and end-to-end examples per second is what settles that. And a pipeline that changed meaning when concurrency was added is a partitioning bug, which you find by counting unique samples rather than by watching a loss curve.

There is one assumption left standing, and this chapter has leaned on it constantly without examining it. Every measurement here treated the batch as a unit of delivery: something with a size, a cost and an arrival time. The Producer returned torch.randn(128) and a random label, because the contents were irrelevant to the question of when it arrived.

Real samples are not random numbers. They are files that were decoded, values that were converted, arrays that were resized, tensors that were normalized against statistics computed from somewhere, and images that may have been augmented differently on this epoch than the last. Every one of those steps is a decision about what the numbers mean, and every one of them can be made wrong in a way that produces a tensor of exactly the right shape, dtype and device.

Chapter 6 asked why the batch was late. The next chapter asks a harder question about the same batch: what does the tensor inside it actually represent, and is it what the model was built to receive?

What the batch means.