PyTorch DevLog

Pinned memory: what it is for, and why nobody gives it back

Edward Yang (@ezyang) · August 9, 2026 · 10 min read
eagercudamemorypinned-memorycuda-graphs

Disclosure. This post was drafted by Claude (Anthropic’s coding assistant) with editing from ezyang.

For quite some time, I used to think of pinned memory as something you sprinkled around your code to let you do async transfers to GPU. In fact, the old Caffe2 used to pin every single CPU tensor, even if it never actually participated in GPU compute. Under this regime, you might imagine that you need an allocator for pinned memory that lets you allocate and free pinned memory as necessary.

However, if you think carefully about the implications of async transfers on pinned memory lifetime, as well as the implications for CUDA graphs, it turns out that you don’t… really want to ever free pinned memory. The tl;dr:

  1. CUDA graphs forbid it: a captured graph bakes the host buffer’s address into its nodes, so any pinned buffer that participates in a graphed copy must stay alive at that address for as long as the graph is replayed.

  2. Re-pinning is expensive: cudaHostAlloc costs milliseconds, so freeing in steady state just means paying that again on the next iteration.

  3. The workloads that use pinned memory have a fixed working set: same shapes every step, so releasing between steps buys nothing.

The rest of this post works through the mechanics.

By the way, as a shortcut to believing this claim, it’s worth looking at how Megatron-LM and vLLM actually use pinned memory, because neither of them ever releases it. Both have their own pinned pools. Both allocate up front and hold the memory until the process exits. One of them has a config flag whose entire documented purpose is to stop pinned buffers from being freed.

What pinned memory is actually for

Pinned (page-locked) host memory is memory the OS has promised not to swap or move, which lets the GPU’s DMA engines read and write it directly by physical address. This buys you four things:

The complication: you don’t own the buffer when you think you do

When you issue

gpu_tensor.copy_(cpu_pinned, non_blocking=True)

the call returns immediately and the DMA is still in flight. The buffer is not safe to reuse when the Python statement finishes, nor when the last reference drops; it is safe to reuse when the copy actually completes on the stream, and only the GPU knows when that is. If you recycle the buffer early, you get silent data corruption.

PyTorch’s pinned memory allocator helps you avoid this use-after-free problem. Every copy_(non_blocking=True) involving pinned memory records a CUDA event against the host block (aten/src/ATen/native/cuda/Copy.cu:481-486):

const auto& host_tensor = (dst_device == kCPU ? dst_tensor : src_tensor);
auto* ptr = (dst_device == kCPU ? dst : src);
auto* ctx = host_tensor.storage().data_ptr().get_context();
at::getHostAllocator(at::kCUDA)->record_event(ptr, ctx, stream.unwrap());

Note that it keys on both the data pointer and the storage context, so a slice of a pinned tensor, or one built with from_blob over allocator memory, still attributes to the right block. On free, the block is parked in an event queue and is not returned to the free list until every recorded event has retired (aten/src/ATen/core/CachingHostAllocator.h:428-500). Setting PYTORCH_ALLOC_CONF=pinned_use_background_threads:true moves that polling to a helper thread so it stays off the allocation fast path.

Because frees in Python are always implicit (via refcount), my personal opinion is that making sure freed buffers don’t get reused too early is a very important anti-footgun that PyTorch’s pinned memory allocator provides.

CUDA graphs turn buffers into permanent fixtures

A CUDA graph bakes device and host addresses into the captured nodes. If a D2H copy is part of the graph, it writes to one fixed host address on every replay, forever. Two consequences:

This is why graph-friendly code works the way it does. If you want to read something back from the device (loss, grad norm, a NaN check, sampled token ids, KV cache lengths) and you want the whole step to stay inside one graph, you allocate a pinned staging buffer once at startup, capture the D2H copy into it, and read it on the host after the replay. The alternative is a device-to-host sync in the middle of your step, which un-graphs everything.

Megatron has a config flag for exactly this (megatron/core/model_parallel_config.py:480):

cpu_offloading_retain_pinned_cpu_buffers: bool = False
"""If True, the pinned CPU buffers are retained after offloading and reused for the
   next iteration. It is useful for cuda graphs capture.
"""

A flag whose entire purpose is: do not give the pinned memory back, because CUDA graphs need stable addresses.

PyTorch’s host allocator encodes the same invariant internally. Blocks allocated during stream capture are never recycled (aten/src/ATen/core/CachingHostAllocator.h:481-493):

// If the block was ever used, block->event_count_ will be above
// 0 and thus can never be recycled by
// process_events_for_specific_size. Thus, this block will never
// be returned again. "Leaking" memory like this is intentional
// to avoid subtle cuda graph problems described here: ...

Some examples of pinned memory usage

You don’t have to use PyTorch’s pinned memory allocator; it’s relatively simple to call cudaHostRegister yourself.

vLLM allocates its host-side input buffers once at worker init, sized to max_num_reqs, and reuses them every step for the duration of the process (vllm/v1/worker/gpu_model_runner.py, some twenty pin_memory=PIN_MEMORY allocations, mostly at init). This is the graph-friendly input-prep pattern: fill the pinned buffer on the host, one async H2D, replay the graph.

For KV offload, vLLM goes further and bypasses PyTorch’s allocator entirely (vllm/v1/simple_kv_offload/cuda_mem_ops.py:23-32):

def pin_tensor(tensor: torch.Tensor) -> None:
    """Pin a CPU tensor via cudaHostRegister.

    This bypasses PyTorch's CUDACachingHostAllocator which rounds
    every ``pin_memory=True`` allocation up to the next power of 2
    (e.g. 100 GB becomes 128 GB).
    """
    err = torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0)

Megatron has its own pinned pool, OffloadTensorPool, which it builds as an abstraction over PyTorch’s pinned memory pool:

All three implementations (PyTorch’s, vLLM’s, Megatron’s) have one thing in common: none of them ever return pinned memory to the OS during steady-state operation. They allocate a working set, reuse it, and hold it until the process exits.

So, does deleting the tensor unpin it? No

When the last reference to a pinned tensor drops, PyTorch records events on any streams that touched it, and once those retire, the block goes on a free list (maybe_cache_block, CachingHostAllocator.h:808). Nothing calls cudaFreeHost. The pages stay page-locked and registered with the driver.

import gc, torch

MiB = 1024 * 1024

def report(label):
    s = torch.cuda.host_memory_stats()
    print(f"{label:20} owned={s['allocated_bytes.current']//MiB:5} MiB"
          f"  checked_out={s['active_bytes.current']//MiB:5} MiB"
          f"  cudaFreeHost calls={s['num_host_free']}")

t = torch.empty(256 * MiB, dtype=torch.uint8, pin_memory=True)
report("after alloc")

del t
gc.collect()
torch.cuda.synchronize()
report("after del")          # checked_out -> 0, owned stays at 256 MiB

torch.accelerator.empty_host_cache()
report("after empty_cache")  # owned drops, num_host_free increments

active_bytes is memory checked out to callers; allocated_bytes is memory the allocator owns, active plus cached. Deleting a tensor moves bytes from the first to the second, and only empty_cache removes them from the second.

When reasoning about memory allocations in PyTorch, it’s important to distinguish three different things. Releasing the tensor is refcounted and deterministic. Making the block reusable is deferred, but only until the stream events retire, and that deferral is correctness, not laziness. Returning pages to the OS never happens on its own.

When you do want it back, and how

“Allocate once and hold” is right for a training job that owns the box. There are two cases where it is not:

For those cases, you can force returning pinned memory to the OS:

torch.accelerator.empty_host_cache()   # public, device-generic, 2.12+
torch._C._host_emptyCache()            # CUDA-only, private, since 2.5

torch.cuda.empty_cache() does not do this. It only calls _cuda_emptyCache for the device allocator, and there is deliberately no torch.cuda.empty_host_cache().

empty_cache only reclaims blocks sitting on a free list, so it is safe to call at any time; it just may free less than you hoped if there are still async transfers going on.

If a full flush is too blunt (every subsequent allocation of that size pays a fresh multi-millisecond cudaHostAlloc), bound the cache instead:

PYTORCH_ALLOC_CONF=pinned_max_cached_size_mb:512

Blocks above this size are freed as soon as their copy events retire rather than cached. There is also pinned_max_round_threshold_mb, which disables the power-of-two rounding above a given size; it postdates vLLM’s cudaHostRegister workaround above and would address the same 100 GB-becomes-128 GB problem.

One thing that is never released under any circumstances: if you enable pinned_reserve_segment_size_mb, that slab is explicitly skipped by free_block (aten/src/ATen/cuda/CachingHostAllocator.cpp:95) and lives until the process exits, by design.

Unpinning memory

If you want to think of pinned memory as something that happens to normal CPU memory, you are welcome to directly use cudaHostRegister / cudaHostUnregister on an ordinary CPU tensor, reachable from Python via the public torch.cuda.cudart(). PyTorch ships a four-line wrapper at torch/cuda/_pin_memory_utils.py (private, but the underlying call is not), and distributed checkpointing uses it for staging buffers. You get fully explicit lifetime, no caching, no rounding; but you have to make sure you get your events correct.

Can I use pinned memory without CUDA?

_pin_memory is an accelerator-specific API and will error with no accelerator present (e.g., a CPU-only build). This can be somewhat inconvenient if you sometimes like to test with CPU device. I don’t really have any advice besides irritating t.pin_memory() if torch.accelerator.is_available() else t tests.

Note that technically you can get page-locked host memory with no accelerator involved: e.g., mlock-ed buffers, or the page-aligned allocations that O_DIRECT / libaio NVMe offload wants. This has nothing to do with CUDA and PyTorch doesn’t provide APIs for it.

Conclusion

Hopefully, this answers some questions you might have about pinned memory and how to go about working with it. It is indeed quite easy to roll your own pinned memory allocator directly from Python, and you should feel free to do so if appropriate (though I doubt you’d need this post if you do.) But the built-in pinned memory allocator I think is also quite serviceable, and you should mostly consider not using it if you want exact-size fit instead of power-of-two buckets, or if you promise to take care of deallocations yourself.

References