Read alongside
PyTorch: torch.utils.data (Dataset and DataLoader): Single- and Multi-process Data Loading
How workers, batching, collation and memory pinning behave in the standard data-loading API.
Objectives
- Diagnose whether a training step is input-bound by comparing it with a synthetic-data baseline and reading a profiler timeline.
- Configure a DataLoader (workers, batch size, prefetching, persistent workers, collation) and explain the cost of each setting.
- Explain pinned memory, non-blocking copies and the conditions under which a host-to-device copy overlaps with kernels.
- Identify and remove hidden host-device synchronization points in a training loop.
- Explain how gradient communication and parameter gathering are overlapped with computation and verify it in a trace.
- Design an input pipeline for a dataset that does not fit in memory, including sharding, shuffling and deterministic resumption.
Why it matters
Accelerators do useful work only while data and communication keep up with them. When input loading, host-to-device copies or collectives are not overlapped with compute, the most expensive component waits. "Training is slower than expected" is often an input-pipeline or synchronization problem, and interviewers want to see a measured diagnosis rather than a list of tricks.
Core ideas
A pipeline view. Training is a chain: storage, decode and augment (or tokenize and pack), batch, host memory, device, compute, communication. The slowest stage sets throughput; buffers and prefetching hide variability but cannot fix a producer that is slower on average. The quickest test is to replace the dataset with synthetic tensors already on the GPU: if the step gets much faster, you are input-bound.
DataLoader mechanics. Worker processes load and transform samples in parallel, outside the main Python thread. Each worker prefetches a few batches; persistent workers avoid restarting processes every epoch. More workers help until CPU cores, host memory or storage bandwidth saturate, and each worker holds its own copy of dataset state, which can inflate host memory. For sequences, padding to the longest item in a batch wastes compute; bucketing by length or packing several sequences into one row reduces the waste.
Pinned memory and copies. The GPU's copy engines read host memory directly only when it is page-locked (pinned); a copy from ordinary pageable memory first passes through a pinned staging buffer. With non_blocking=True the host thread does not wait for the copy, but the copy overlaps with kernels on the GPU only if the source is pinned, the copy runs on a non-default stream and a copy engine is free. Pinning a tensor right before copying it adds work and can be slower than a plain copy; the benefit comes from producing batches directly in pinned memory (DataLoader pin_memory=True pins in a background thread) or reusing pinned buffers. Non-blocking device-to-host copies need an explicit synchronization before the CPU reads the result, or it may read incomplete data.
Streams and synchronization. Kernels on one stream run in order; work on different streams may overlap. Hidden synchronization points serialize the pipeline: .item(), .cpu(), printing tensor values, operations whose output shape depends on data (nonzero, boolean masking) and explicit synchronize calls. Logging the loss value every step is a classic stall.
Communication overlap. DDP launches bucketed gradient all-reduces during backward; FSDP prefetches the next layer's parameter all-gather. Overlap helps only when enough independent compute exists to cover the transfer and when communication and compute do not compete for the same resources.
What to look for in a timeline. GPU idle gaps at the start of each step (waiting for data), host-to-device copy durations, CPU time inside the DataLoader's next call, and collectives not covered by compute.
Common misconceptions
- "More DataLoader workers are always faster." Beyond the available cores, memory and I/O bandwidth, workers contend and startup cost grows.
- "non_blocking=True makes every copy asynchronous and overlapped." It frees the host thread; overlap with kernels needs pinned memory and a separate stream, and unsynchronized device-to-host results can be wrong.
- "High GPU utilization means the input pipeline is fine." Utilization counts any running kernel; compare step time with a synthetic-data baseline.
- "Overlap happens automatically." It needs independent work, separate streams and no synchronization points in between.
- "Data loading only matters for images." Tokenization, packing, decompression and remote object storage bottleneck text pipelines too.
How to approach it
Profile first, then change one thing at a time: workers, pinned memory, batch layout, synchronization points, communication overlap. In interviews, describe the synthetic-data test, the timeline signals you would look for, and the fix for each signal. Senior discussions add dataset sharding across many hosts, deterministic shuffling with resumable state, and storage throughput planning.
Self-check before you finish
- What does the synthetic-data test tell you, and what can it miss?
- Which conditions must hold for a host-to-device copy to overlap with a kernel?
- Name four operations that force host-device synchronization.
- How do length bucketing and sequence packing change wasted compute?
- Why can adding DataLoader workers make a job slower?