PyTorch DataLoader Interview Questions: Workers, Seeding, and Host-to-Device Bottlenecks

A fast GPU can still wait for a slow or incorrect input pipeline. In a PyTorch DataLoader interview, “increase num_workers” is rarely a sufficient answer: it may hide a loader stall, duplicate an iterable dataset, or change the sample stream you thought you were training on. Diagnose sample identity first, then randomness, then the time from storage to the device.

This article is a worked, original exercise with ten invented sample IDs and a synthetic timing trace. It is not a measured GPU benchmark. The behavior of map-style and iterable-style datasets, worker seeds, and pinned memory follows PyTorch's data-loading reference.

Interview prompt: correct the samples before tuning the speed

A team has an IterableDataset representing IDs 0 through 9. With num_workers=0, an epoch yields ten IDs. With num_workers=2, the training loop sees twenty: every ID appears twice. Meanwhile, GPU utilization is low and each step has a long gap before compute. The team proposes four workers, pin_memory=True, and a global seed. What would you check, what code would you change, and which measurements would justify the performance settings?

The first answer is not a performance flag. An iterable dataset is replicated into each worker process. If each replica independently iterates over all ten IDs, two workers each yield 0..9, producing twenty records. The bug is in partitioning the stream, not in the GPU. The official DataLoader guide calls out duplicate data as a consequence of naive multi-process iterable loading. More workers would multiply the error.

An interview-ready invariant is: for this toy single-process, single-epoch source, the collected IDs must have count ten, unique count ten, minimum zero, maximum nine, and no missing IDs. The exact arrival order may change under worker scheduling, so check a multiset or sorted IDs rather than assuming the two worker streams will interleave in one fixed pattern.

Exercise 1: shard an iterable dataset across workers

Here is the intentionally incorrect version. The method runs in each worker's replica when multiprocessing is enabled:

from torch.utils.data import IterableDataset, DataLoader

class IDs(IterableDataset):
    def __iter__(self):
        return iter(range(10))

loader = DataLoader(IDs(), batch_size=None, num_workers=2)

The corrected version uses worker identity to give each replica a disjoint slice of the same finite ID set:

from torch.utils.data import get_worker_info

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

With two workers, worker 0 yields {0,2,4,6,8} and worker 1 yields {1,3,5,7,9}. The union is {0..9}, the intersection is empty, and the total count is ten. Those set properties are the useful proof; the displayed batch order is incidental. The same stride idea extends to an integer ID stream, but real streaming sources often need partition-aware cursors or upstream shard assignment rather than reading all records and discarding most of them.

ids = list(DataLoader(ShardedIDs(), batch_size=None, num_workers=2))
assert len(ids) == 10
assert len(set(ids)) == 10
assert sorted(ids) == list(range(10))

This small test should fail for the naive loader and pass for the sharded loader. On macOS or Windows, multiprocessing commonly uses spawn, so keep custom dataset classes and worker functions importable at module top level and put the driver under if __name__ == "__main__":. The code above is a conceptual fixture; the assertions have been checked by ordinary Python set arithmetic, while no PyTorch runtime was available locally. In production, run the exact DataLoader test on the deployment OS and process start method.

Distinguish this from a map-style dataset. For map-style data, the main process normally generates sample indices via a sampler and sends those indices to workers; it does not simply replicate the full index sequence in each worker. A DistributedSampler or a custom sampler may still introduce repeats or omissions by design, especially around padding and incomplete batches, so inspect the actual sampled IDs instead of assuming all duplicates imply an IterableDataset bug. Explain which sampling layer owns the order.

There is another subtlety at epoch boundaries: for multi-worker iterable datasets, drop_last=True drops an incomplete last batch per worker replica, not necessarily one global tail. A small shard may lose proportionally many samples. A reported len(dataloader) is not a definitive proof of how many unique IDs reached training in an iterable multi-worker run. Count IDs in a diagnostic pass, and make the intended epoch boundary explicit.

For a real streaming source, the ten-ID test is only the start. Define whether a sample may appear once per epoch, once per rank, or at least once over a restart window. With multiple distributed ranks, each rank's worker IDs start locally; a worker-aware stride alone does not partition across ranks. Combine rank and worker ownership, or use source partitions that are exclusive across all consumers, and test the global set of IDs. At-least-once storage delivery may legitimately replay a record after a crash; if training cannot tolerate that, checkpoint the source offset and sample identity at a well-defined commit boundary. “No duplicate in one local iterator” is weaker than “no duplicate in the distributed training run.”

Exercise 2: seed the right random generators

Suppose the duplicate-ID fix works, but randomized augmentation still differs across runs. Ask which random library chooses the transform, where the randomness executes, and what “reproducible” means. A fixed top-level torch.manual_seed(0) alone may not control Python random, NumPy, a third-party augmentation library, sampler order, or every nondeterministic device operation. Conversely, a seed can reproduce the same incorrect duplicate stream; seeding is not sharding.

PyTorch sets each DataLoader worker's PyTorch seed from a base seed plus its worker ID. Its reproducibility notes show a generator plus worker_init_fn for other libraries. For a classic NumPy/Python-random setup, the pattern is:

import random
import numpy as np
import torch

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(0)
loader = DataLoader(
    ShardedIDs(), num_workers=2, batch_size=None,
    worker_init_fn=seed_worker, generator=g,
)

The function must be defined where workers can import it under spawn. If your augmentation code uses a separate NumPy Generator object, seeding NumPy's legacy global RNG may not control that object; wire an explicit generator into the component that uses it. Record the seed, dataset snapshot, worker count, sampler configuration, and software environment when comparing runs. PyTorch's reproducibility guidance also cautions that complete reproducibility is not guaranteed across releases, platforms, or CPU and GPU executions.

persistent_workers=True keeps worker processes and dataset instances alive across DataLoader iterations. That may save startup cost, but it means stateful iterable datasets, file handles, and worker-side RNG streams need deliberate epoch handling. Do not assume worker_init_fn reruns each epoch when workers persist. If the stream needs per-epoch reshuffling, define and test an epoch protocol; if it needs exact resume, persist cursor and RNG state at the appropriate level rather than merely reusing a global seed.

Exercise 3: locate the host-to-device bottleneck

Now assume the ID and seed checks pass, but training is slow. The team shows a synthetic single-step timeline: 80 ms waiting for the next decoded/augmented batch, 15 ms for host-to-device transfer, and 30 ms of GPU compute. These numbers are invented interview inputs, not observed performance. They suggest an input path longer than compute, but one step does not identify the exact cause. It may be disk or object-store latency, decoding, Python augmentation, collation, queue starvation, host memory pressure, or synchronous copying.

Measure at least four intervals under representative load: time spent waiting for next(loader), CPU decode/augmentation and collate time, copy time and overlap, and GPU forward/backward/optimizer time. Warm the pipeline before measuring steady state. CUDA calls can be asynchronous, so a host timer around a .to(device) or a kernel launch does not automatically measure completed device work. Use a profiler with CPU and CUDA activities or carefully placed synchronization/events when comparing device-stage timings. Check end-to-end examples per second and tail step time alongside individual spans.

Run one bounded experiment at a time. Compare num_workers=0,2,4 while holding dataset, transforms, batch size, and model fixed. If two workers shrink the input wait and improve end-to-end throughput, that supports parallel CPU loading; if four workers add contention, memory use, or more variable tail latency, stop there. Measure startup separately from steady-state epochs so persistent_workers does not disguise a slow first epoch or a stuck shutdown. Tune prefetch_factor only after confirming queue depth is relevant; it controls batches prepared per worker and can increase memory use.

For CUDA training, pin_memory=True asks DataLoader to return fetched tensors in page-locked host memory, which can improve host-to-device transfer. Combine it with a deliberate batch.to(device, non_blocking=True) path and verify actual overlap and throughput; PyTorch's pinning/non-blocking guide warns against assuming manual pin_memory() in the main thread is free. A custom batch object may need a pin_memory() method for automatic pinning to recognize its tensors. Pinned memory consumes a limited host resource, and indiscriminate pinning or deep prefetch queues can make pressure worse.

Do not move CUDA tensors out of DataLoader worker processes merely to “skip” the copy. PyTorch's multi-process data guide discourages returning CUDA tensors from workers because CUDA multiprocessing and tensor sharing have subtleties. Keep the worker path CPU-side unless you have a deliberately engineered exception. If the profiler shows the GPU is busy and the input queue never empties, loader flags are not the current bottleneck; inspect model kernels, synchronization, and distributed communication instead.

A decision table for the interviewer

Observed symptom	First discriminating check	Likely next move
Every ID appears once per worker	Log (worker_id, sample_id) for one epoch	Shard IterableDataset replicas; recheck union/intersection
Augmentations change unexpectedly	Identify each RNG and its process	Seed worker-side libraries and record the generator/environment
Long idle gap before GPU work	Compare next(loader), CPU stage, copy, and GPU spans	Fix the longest measured stage, then remeasure throughput
More workers reduce throughput	Inspect CPU saturation, memory, queue wait and startup	Back off worker/prefetch settings; examine I/O contention
Copy dominates after batch is ready	Check pinned batches and asynchronous copy behavior	Test pinning/non-blocking path against an end-to-end baseline

The table is a diagnostic sequence, not a list of settings to enable together. Each intervention must preserve sample count, labels, ordering requirements, and model-quality checks. If you change batch size or drop policy while tuning, you have changed the experiment as well as the input pipeline.

In a strong interview answer, lead with the 0..9 proof: two naive iterable replicas produce twenty records; sharding produces two disjoint five-ID sets. Then separate worker seeding from data partitioning. Finally, use a trace to justify whichever of worker count, prefetching, pinned memory, or transfer overlap actually improves full-step throughput. Low GPU utilization is a clue, not a diagnosis.

Sources and Further Reading

• PyTorch: torch.utils.data and DataLoader behavior
• PyTorch: Reproducibility and DataLoader worker seeding
• PyTorch: Pin memory and non-blocking transfers
• PyTorch: Performance tuning guide