Pinned Memory and DMA Transfers in PyTorch

The step is short. The bus isn’t.

You profile a training step. Compute is short, but the GPU waits between batches while the CPU prepares and transfers the next tensor. pin_memory=True changes how those host bytes reach the device and can make overlap possible.

Two numbers

1. Path · staging plus H2D, or direct H2D

The pageable path stages data in pinned host memory before the host-to-device transfer. The diagram counts one CPU staging copy plus one H2D transfer. An already pinned tensor can take the direct H2D path; this count excludes the earlier work needed to prepare or pin it.

2. Overlap · the win is density, not µs
non_blocking=True lets the host proceed without waiting for each transfer to finish. Overlapping a copy with a GPU kernel also requires pinned source memory, a separate CUDA stream and an available DMA engine. See PyTorch's transfer guide.

Compare both paths, then step through the authored overlap schedule. These figures illustrate dependencies; they do not report hardware timings.

Path: two copies vs one

Pageable is the default. Staging is invisible until you watch the path. Pinned skips the middle box.

batch = next(loader)          # pageablebatch = batch.to('cuda')      # sync path# hidden: pageable → staging → GPU

Overlap: why the flag exists

An input pipeline can hide some transfer time when loading, H2D and compute overlap. The tape assumes the stream, DMA and synchronization conditions above are satisfied; pinning alone does not establish that overlap.

DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=2,
)

images = batch["images"].to("cuda", non_blocking=True)

Dictionary: same bytes, different contract

Pageable pages can leave RAM. Pinned pages cannot. That is the whole OS-side trade.

When pinning hurts

Pinned memory cannot be swapped. Footprint roughly:

pinned ≈ num_workers × prefetch_factor × batch_bytes

On a machine with headroom, locking a few hundred MB is free performance. On a 16 GB box already full of model + optimizer + Python, those locked pages push other processes into swap — worse than the staging copy you were trying to avoid.

# if Swap used climbs during training, shrink the pin footprint
watch -n 1 'free -h | grep Swap'
# 1) lower prefetch_factor or num_workers
# 2) smaller batch
# 3) pin_memory=False last

On multi-socket hosts, pinned allocations follow the NUMA node of the allocating thread. If the GPU lives on the other socket, DMA crosses the interconnect. Bind with numactl when that shows up in profiles — DataLoader will not fix topology for you.

What to do

1. Enable pin_memory=True for GPU training unless RAM is already critical.
2. Always pair with .to(device, non_blocking=True) — pinning without async leaves half the win on the table.
3. Keep workers + prefetch so the pin thread has batches ready; tune those before disabling pin.
4. Watch swap if the box is small; locked pages are not free.
5. Confirm in the profiler — look for cudaMemcpyAsync from host pinned, not a double host path.

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=2,
)

for batch in train_loader:
    images = batch['images'].to('cuda', non_blocking=True)
    labels = batch['labels'].to('cuda', non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

Further Reading

• PyTorch DataLoader — pin_memory and multi-process loading
• CUDA C Programming Guide: Page-Locked Host Memory — cudaHostAlloc and DMA rules
• PyTorch Performance Tuning Guide — data pipeline knobs

Related concepts