Same order, different numbers

Pass a torch.Generator to DataLoader and two different things read from it. dl.sampler.generator is g comes back true, so the RandomSampler holds the object you handed over rather than a copy. The iterator draws a single 64-bit integer off that same object at construction.

Chapter four's arc already puts the first of those two to work. Its third lesson, 'The resume that matches', saves the generator's state at an epoch boundary to replay the next epoch's order, and names the one piece of a resume that has no state to save at all. None of that is repeated here. The question this lesson needs answered is a different one: where the draw runs.

It runs in the main process. self._next_index() sits inside the _try_put_index you read in the last lesson, and only the resulting list of indices goes onto a worker's queue, so no worker ever runs the sampler. Shuffle eight samples with zero, two and three workers and the order comes back [5, 3, 1, 7, 2, 6, 4, 0] all three times.

The second reader of that generator is one line of _BaseDataLoaderIter.__init__, and it is what makes the other half of reproducibility hard. self._base_seed = torch.empty((), dtype=torch.int64).random_(generator=loader.generator).item() is where the single integer comes from, and every worker's random state descends from it.

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

class Noisy(Dataset):
    def __len__(self): return 8
    def __getitem__(self, i):
        return torch.tensor([float(i), torch.rand(()).item()])

def run(nw, seed=0):
    g = torch.Generator(); g.manual_seed(seed)
    torch.manual_seed(seed)
    dl = DataLoader(Noisy(), batch_size=2, shuffle=True, num_workers=nw, generator=g)
    return dl, g, [[round(v, 4) for v in row] for b in dl for row in b.tolist()]

if __name__ == "__main__":
    dl, g, _ = run(0)
    print("sampler:", type(dl.sampler).__name__, "| sampler.generator is g:", dl.sampler.generator is g)
    for nw in (0, 2, 3):
        _, _, r = run(nw)
        print(f"num_workers={nw}")
        print("  order    ", [int(x[0]) for x in r])
        print("  torch.rand", [x[1] for x in r])

# ---- stdout ----
# sampler: RandomSampler | sampler.generator is g: True
# num_workers=0
#   order     [5, 3, 1, 7, 2, 6, 4, 0]
#   torch.rand [0.4963, 0.7682, 0.0885, 0.132, 0.3074, 0.6341, 0.4901, 0.8964]
# num_workers=2
#   order     [5, 3, 1, 7, 2, 6, 4, 0]
#   torch.rand [0.7821, 0.0536, 0.6938, 0.298, 0.9888, 0.1949, 0.1669, 0.2847]
# num_workers=3
#   order     [5, 3, 1, 7, 2, 6, 4, 0]
#   torch.rand [0.7821, 0.0536, 0.6938, 0.298, 0.654, 0.2994, 0.9888, 0.1949]

Three lines at the top of _worker_loop set the whole random state of a worker, and the arithmetic is as simple as it looks. The worker adds its own id to the base seed, then seeds Python's random, then seeds torch. NumPy gets a derived state instead, computed by a local reimplementation of SeedSequence, because seeding two generators with the same integer is a known way to correlate them.

Reading WorkerInfo.seed back out of three workers confirms the addition. With the loader generator seeded at 0, the three workers reported seeds ending 455, 456 and 457, and torch.initial_seed() inside each worker matched its own.

Now put that beside round robin from the last lesson and the divergence in the run above stops being mysterious. Batch i goes to worker i % num_workers, and each worker draws from a stream seeded by base_seed + worker_id, so changing the worker count changes which stream a given batch draws from. The last four values of the num_workers=3 row are the proof: 0.9888, 0.1949 appear there as the fourth batch and appear in the num_workers=2 row as the third, because in both cases that is worker 0's second batch.

The practical rule is short. Sample order is reproducible across worker counts; anything random inside __getitem__ is reproducible only for a fixed num_workers. If a result has to survive a change in worker count, the randomness has to be derived from the sample index rather than drawn from ambient state, and worker_init_fn plus get_worker_info().seed is where you would put that derivation.

        torch.set_num_threads(1)
        seed = base_seed + worker_id
        random.seed(seed)
        torch.manual_seed(seed)
        if HAS_NUMPY:
            np_seed = _generate_state(base_seed, worker_id)
            import numpy as np
            np.random.seed(np_seed)

# an IterableDataset reporting its own WorkerInfo, generator seeded at 0, stdout:
# worker id, WorkerInfo.seed mod 1000, torch.initial_seed() mod 1000:
#   [[0, 455, 455], [1, 456, 456], [2, 457, 457]]

The first line of that excerpt is not about randomness at all, and it is the one with performance consequences. torch.set_num_threads(1) runs before your dataset is touched, so every worker does its tensor work single-threaded.

The reason is defensive. Without it, N workers each opening a thread pool sized to the machine would oversubscribe every core and fight the training step for them. The cost is that a transform torch would have parallelized in the main process no longer is.

That cost is measurable and the measurement needs an honest caveat. A dataset whose __getitem__ does one 768 by 768 matmul reported torch.get_num_threads() of 8 when fetched in the main process and 1 inside a worker, every run. The per-sample time for that matmul came out at 6.00, 24.11 and 23.75 ms in the main process across three runs, against 29.26, 79.20 and 92.49 ms in a worker. The absolute numbers move with whatever else this laptop is doing and should not be read as anything; the ratio, roughly five to eight times slower per sample, held across all three.

Aggregate throughput is the number that would actually decide a configuration, and this machine could not produce a trustworthy one. Repeat runs of the same epoch swung between 18 and 42 samples per second at num_workers=2, so no table of it appears here. Chapter nine's harness discipline is what a claim like that would have to pass first.

Map-style datasets are safe under workers because the main process owns the indices and each worker is told which ones to fetch. An iterable-style dataset has no indices to hand out. Each worker gets its own copy of the object and calls iter() on it, which means each worker produces the entire stream.

The result is duplication with no error and no warning. An eight-item stream read with two workers yielded sixteen items, every sample twice, in an order that interleaves the two copies.

The fix is to shard inside __iter__, and get_worker_info() is the only information you need to do it. Stride the stream by num_workers starting at id, treat a None return as the single-process case, and the same dataset yields each sample exactly once at any worker count.

This is also the reason _InfiniteConstantSampler shows up on an iterable-style loader. There is nothing to sample, so the loader emits a constant to drive the fetch loop and lets the dataset decide when the epoch ends.

import torch
from torch.utils.data import DataLoader, IterableDataset, get_worker_info

class Stream(IterableDataset):
    def __iter__(self):
        for i in range(8): yield torch.tensor([float(i)])

class Sharded(IterableDataset):
    def __iter__(self):
        wi = get_worker_info()
        lo, step = (0, 1) if wi is None else (wi.id, wi.num_workers)
        for i in range(lo, 8, step): yield torch.tensor([float(i)])

if __name__ == "__main__":
    for nw in (0, 2):
        got = [int(v) for b in DataLoader(Stream(), batch_size=2, num_workers=nw) for v in b.flatten()]
        print(f"IterableDataset, num_workers={nw}: {got}")
    print("sampler:", type(DataLoader(Stream(), batch_size=2, num_workers=2).sampler).__name__)
    for nw in (0, 2, 3):
        got = [int(v) for b in DataLoader(Sharded(), batch_size=2, num_workers=nw) for v in b.flatten()]
        print(f"sharded, num_workers={nw}: {sorted(got)} (n={len(got)})")

# ---- stdout, the two runs joined ----
# IterableDataset, num_workers=0: [0, 1, 2, 3, 4, 5, 6, 7]
# IterableDataset, num_workers=2: [0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7, 6, 7]
# sampler: _InfiniteConstantSampler
# sharded, num_workers=0: [0, 1, 2, 3, 4, 5, 6, 7] (n=8)
# sharded, num_workers=2: [0, 1, 2, 3, 4, 5, 6, 7] (n=8)
# sharded, num_workers=3: [0, 1, 2, 3, 4, 5, 6, 7] (n=8)

Page-locked host memory is what pin_memory=True is asking for, and it only pays off against a device that can DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → out of it while the CPU does something else. This course was written on a machine with no such device, and the loader's behaviour there is worth stating exactly rather than skipping.

The flag is accepted and then dropped. DataLoader.pin_memory still reads True, while the iterator's _pin_memory reads False, because the iterator recomputes it as loader.pin_memory and torch.cuda.is_available(). No pin-memory thread starts, batches come back with is_pinned() false, and torch 2.2.2 issues no warning at all: catching warnings around a full iteration returned an empty list.

Current torch fixed the silence. In pytorch at the v2.9.0 tag the same two lines warn first, with the text 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used., and the availability check is torch.accelerator.is_available() rather than a CUDA-specific one. Same outcome, louder.

What the flag buys where it does apply is a thread, not a process. _pin_memory_loop runs in the main process, pulling finished batches off the worker result queue, calling .pin_memory() on every tensor in them, and putting them on a second queue for your loop to read. Only then can a device copy be issued with non_blocking=True and actually overlap with compute, which is the pairing the CUDA notes describe and the reason the flag is nearly always set together with that argument.

Worth knowing what the failure looks like if you call it by hand instead. torch.zeros(4).pin_memory() on this machine raises NotImplementedError: Could not run 'aten::_pin_memory' with arguments from the 'CUDA' backend, which names CUDA on a machine that has none because the pinning op dispatches to the accelerator backend by default.

import warnings, torch
from torch.utils.data import DataLoader, TensorDataset

ds = TensorDataset(torch.randn(64, 4))
print("cuda available:", torch.cuda.is_available())
dl = DataLoader(ds, batch_size=8, pin_memory=True)
it = iter(dl)
print("loader.pin_memory:", dl.pin_memory, "| iterator _pin_memory:", it._pin_memory)
b, = next(it)
print("batch is_pinned:", b.is_pinned())

with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    next(iter(DataLoader(ds, batch_size=8, pin_memory=True)))
    print("warnings raised:", [str(x.message) for x in w])

# ---- stdout ----
# cuda available: False
# loader.pin_memory: True | iterator _pin_memory: False
# batch is_pinned: False
# warnings raised: []

# and the line that drops it, verbatim from dataloader.py:588-590 at 2.2.2:
#        if (len(loader.pin_memory_device) == 0):
#            self._pin_memory = loader.pin_memory and torch.cuda.is_available()
#            self._pin_memory_device = None