What actually forks

The chapter states that a worker has to be able to unpickle the dataset it was handed. That is true here and it is not true everywhere, and the difference is one platform default.

On this machine multiprocessing.get_start_method() returns spawn, which is the Python default on macOS and Windows. A spawned worker is a fresh interpreter. It imports your module again, unpickles the arguments it was sent, and starts from nothing, so an unpicklable dataset fails before a sample loads and a script without an if __name__ guard re-runs its own top level.

Hand the same loader a fork context and the failure goes away. A forked worker inherits the parent's memory, so the dataset is never serialized at all; it is simply already there. The run below is the same lambda-carrying dataset twice, once under each start method, and the two results are the honest scope of the rule: what breaks is the start method, not the DataLoader.

Linux defaults to fork, which is why a dataset that runs fine on a training box can fail on a laptop, and why the reverse trap exists too. A forked worker inherits open file handles, database connections and library state that were never designed to be shared, and those bugs disappear the moment someone switches to spawn.

import torch, torch.multiprocessing as mp
from torch.utils.data import DataLoader, Dataset

class Lam(Dataset):
    def __init__(self, f): self.f = f
    def __len__(self): return 4
    def __getitem__(self, i): return self.f(torch.tensor([float(i)]))

if __name__ == "__main__":
    print("default start method:", mp.get_start_method())
    ds = Lam(lambda t: t * 2)
    try:
        print(next(iter(DataLoader(ds, batch_size=2, num_workers=1))))
    except Exception as e:
        print(type(e).__name__ + ":", str(e).strip().split("\n")[-1])
    ctx = mp.get_context("fork")
    dl = DataLoader(ds, batch_size=2, num_workers=1, multiprocessing_context=ctx)
    print("fork + lambda ->", next(iter(dl)).flatten().tolist())

# ---- stdout, the two runs joined ----
# default start method: spawn
# PicklingError: Can't pickle <function <lambda> at 0x10acc0fe0>: attribute lookup <lambda> on __main__ failed
# fork + lambda -> [0.0, 2.0]

Overlap is worth paying for, and the payment has a number. Time the first batch out of a loader over a tiny in-memory dataset, so that nothing but process startup is being measured, and the spread across start methods is three orders of magnitude.

A spawned worker pays for a full torch import, which is most of that three seconds. A forked worker pays for a page-table copy. The single-process loader pays for a list comprehension.

Two things follow. Short runs on a spawn platform can spend more time starting workers than loading data, and persistent_workers=True is how you stop paying the bill once per epoch instead of once per job. On this machine the same two worker pids served both epochs with the flag on, and four different pids appeared without it.

One caveat about the table, since the course publishes provenance rather than adjectives. This is a laptop that was doing other work while the numbers were taken, and the spread was wide, up to 17.6 seconds for one spawn-with-four-workers run. The minimums are quoted because the ratio between rows is the durable part; the absolute values are not.

The topology is asymmetric, and the asymmetry is the scheduling policy. Each worker gets its own index queue; all workers share one result queue. Indices go out addressed to a specific worker, and results come back into a common pile.

_try_put_index is where a batch is assigned. It pulls the next index list from the sampler in the main process, walks a cycle over worker ids until it finds an active one, and puts (send_idx, index) on that worker's queue. Plain round robin, so with all workers alive, batch i goes to worker i % num_workers.

Two consequences are visible from the run underneath. Batches 0 and 2 and 4 came from one pid and 1, 3, 5 from the other, exactly as the modulo predicts. And every worker reported torch.get_num_threads() of 1, because _worker_loop calls torch.set_num_threads(1) before it touches your dataset. The third lesson takes that apart.

Round robin also explains a stall that looks like a bug. Assignment happens before anyone knows how long a sample will take, so one slow index does not get rerouted to an idle worker. It sits in the queue of the worker it was addressed to, and the batches behind it in that queue wait.

    def _try_put_index(self):
        assert self._tasks_outstanding < self._prefetch_factor * self._num_workers

        try:
            index = self._next_index()
        except StopIteration:
            return
        for _ in range(self._num_workers):  # find the next active worker, if any
            worker_queue_idx = next(self._worker_queue_idx_cycle)
            if self._workers_status[worker_queue_idx]:
                break
        else:
            # not found (i.e., didn't break)
            return

        self._index_queues[worker_queue_idx].put((self._send_idx, index))
        self._task_info[self._send_idx] = (worker_queue_idx,)
        self._tasks_outstanding += 1
        self._send_idx += 1

# batch_size=4, num_workers=2, a dataset reporting its own pid; stdout:
# batch 0: idx [0, 1, 2, 3] pid 79276 worker 0 threads 1 shared True
# batch 1: idx [4, 5, 6, 7] pid 79277 worker 1 threads 1 shared True
# batch 2: idx [8, 9, 10, 11] pid 79276 worker 0 threads 1 shared True
# batch 3: idx [12, 13, 14, 15] pid 79277 worker 1 threads 1 shared True

prefetch_factor is not a queue size in batches of memory and not a lookahead in time. It is a cap on outstanding tasks, and the loader primes itself to that cap the moment you call iter(), before you have asked for anything.

The priming loop runs prefetch_factor * num_workers times. With the default factor of 2 and two workers, four batches are already assigned and in flight before the first next(). Reading the iterator's own counters right after iter() shows exactly that: send_idx at 4, rcvd_idx at 0, four entries in the task table.

The cap is also a memory statement nobody writes down. Four batches in flight means four batches of tensors alive at once, plus whatever is sitting in the result queue, and raising prefetch_factor to smooth a jittery loader multiplies that. The assert at the top of _try_put_index is the invariant being enforced.

As batches are consumed, _process_data calls _try_put_index again, one task out for one task in, so the pipeline stays exactly that full until the sampler is exhausted.

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

class Slow(Dataset):
    def __len__(self): return 12
    def __getitem__(self, i):
        time.sleep(0.30 if i < 4 else 0.01)
        return torch.tensor([float(i), float(os.getpid())])

if __name__ == "__main__":
    dl = DataLoader(Slow(), batch_size=2, num_workers=2)
    it = iter(dl)
    time.sleep(1.0)
    print("prefetch_factor", dl.prefetch_factor, "| tasks outstanding before first next():", it._tasks_outstanding)
    print("send_idx", it._send_idx, "rcvd_idx", it._rcvd_idx, "| task_info keys", sorted(it._task_info))

# ---- stdout ----
# prefetch_factor 2 | tasks outstanding before first next(): 4
# send_idx 4 rcvd_idx 0 | task_info keys [0, 1, 2, 3]

Workers finish whenever they finish, and the result queue is shared, so results arrive in whatever order they completed. Your loop still sees batch 0, then batch 1, then batch 2, and the machinery that guarantees it is four lines long.

Every result carries the send_idx it was dispatched with. _next_data wants _rcvd_idx and nothing else. A result whose index does not match gets appended to its entry in _task_info and left there; a result that matches is returned, and _rcvd_idx advances.

That reorder buffer is measurable. Make the first batch slow and the rest fast, iterate with four workers, and the first next() blocks while three completed batches pile up in the table behind it. The three that follow then return in microseconds, because they were finished before the first one was.

Which means a single slow sample delays every batch after it, even though the work for those batches is already done. Current torch offers a way out that 2.2.2 does not have: an in_order argument, defaulting to true, whose docstring says the loader will not enforce first-in-first-out order when it is false.

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

class Uneven(Dataset):
    def __len__(self): return 8
    def __getitem__(self, i):
        time.sleep(1.0 if i < 2 else 0.01)   # only the first batch is slow
        return torch.tensor([float(i)])

if __name__ == "__main__":
    dl = DataLoader(Uneven(), batch_size=2, num_workers=4)
    it = iter(dl)
    t0 = time.perf_counter(); first = next(it); t1 = time.perf_counter()
    print(f"batch 0 = {[int(v) for v in first.flatten()]} after {t1-t0:.2f}s")
    print("already finished and buffered:", sorted(k for k, v in it._task_info.items() if len(v) == 2))
    for n in range(3):
        t = time.perf_counter(); b = next(it)
        print(f"batch {n+1} = {[int(v) for v in b.flatten()]} after {time.perf_counter()-t:.4f}s")

# ---- stdout (the 5.21 s includes spawning four workers) ----
# batch 0 = [0, 1] after 5.21s
# already finished and buffered: [1, 2, 3]
# batch 1 = [2, 3] after 0.0002s
# batch 2 = [4, 5] after 0.0001s
# batch 3 = [6, 7] after 0.0001s

# ---- dataloader.py:1341-1346, the branch that holds them ----
#            if idx != self._rcvd_idx:
#                # store out-of-order samples
#                self._task_info[idx] += (data,)
#            else:
#                del self._task_info[idx]
#                return self._process_data(data)

Passing a whole batch of tensors through a multiprocessing queue would mean pickling every byte, and collate quietly avoids it. collate_tensor_fn asks whether it is running inside a worker, and if it is, allocates the output in shared memory before stacking into it.

So the queue carries a handle, not the data. Check is_shared() on a batch and the answer tracks the process boundary exactly: true with workers, false without. The comment in the source states the intent in one line, and it is about avoiding an extra copy.

How that handle travels is a platform question again. torch.multiprocessing supports two sharing strategies, and on this machine get_all_sharing_strategies() returns only file_system. Linux defaults to file_descriptor, which passes an open fd per tensor over the socket, and that is the origin of the Too many open files failure that a large num_workers with a small ulimit -n produces. A long comment block in dataloader.py documents the failure and includes a standalone script to reproduce it outside torch.

    if torch.utils.data.get_worker_info() is not None:
        # If we're in a background process, concatenate directly into a
        # shared memory tensor to avoid an extra copy
        numel = sum(x.numel() for x in batch)
        storage = elem._typed_storage()._new_shared(numel, device=elem.device)
        out = elem.new(storage).resize_(len(batch), *list(elem.size()))
    return torch.stack(batch, 0, out=out)

# stdout on this machine:
# num_workers=2: shared True
# num_workers=0: shared False
# sharing strategy: file_system | available: {'file_system'}