Skip to content

Device & Memory Management

Overview

PyTorch splits work across CPU (RAM + pinned host memory) and GPU (device memory). Getting tensors to the right device — and overlapping transfers with compute — is often the biggest, cheapest speedup available, especially with the caching allocator masking the truth about memory usage.

  • device='cpu'|'cuda[:i]'|'mps'|'meta'
  • .to(device) allocates + copies; .to(device, non_blocking=True) overlaps with pinned host memory.
  • The CUDA caching allocator hoards freed blocks, so nvidia-smi shows more than Python statistics report.
  • pin_memory + DataLoader(..., num_workers) hides data-loading I/O.

💡 Device transfers are async; profile with .cuda.synchronize() or torch.profiler to see the true cost.

Device Types & Discovery

import torch

print("cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("count  :", torch.cuda.device_count())
    print("name   :", torch.cuda.get_device_name(0))
    print("current:", torch.cuda.current_device())
    print("max mem (MiB):", torch.cuda.get_device_properties(0).total_memory // 2**20)

The meta device — allocations without memory

Perfect for measuring model sizes and running shape inference.

import torch.nn as nn
m = nn.Linear(1024, 1024)
meta = nn.Linear(1024, 1024).to('meta')
print("meta forward is shape-only:", (meta(torch.randn(8, 1024, device='meta'))).shape)

Moving Data: .to, cuda, cpu, clone

x = torch.randn(4, 4)
x_gpu = x.to('cuda')          # copy to GPU
x_back = x_gpu.cpu()          # back to CPU (returns copy)
x_clone = x.clone()           # same device, new storage
x_shared = x_gpu.detach()     # same device, same storage, no grad

dtype/device shorthand

y = x.to(device='cuda', dtype=torch.float16)
z = x.half().cuda()

⚠️ .to() preserves shared storage when only dtype/device match but the tensor is contiguous; otherwise it copies. Prefer explicit .detach() when you need isolation.

Pinned Memory & non_blocking Transfers

Pageable (default) host memory requires a staging copy through pinned buffers during transfer. Pinned host memory can transfer directly to device — enabling real async overlap.

# Pin host tensors (future data stays on device-ready pages)
x = torch.randn(4096, 4096)
x_pinned = x.pin_memory()

# non_blocking -> transfer overlaps with compute already queued
x_gpu = x_pinned.to('cuda', non_blocking=True)

# Process / compute on other tensors here -> runs concurrently
y = x_pinned * 2
torch.cuda.synchronize()

DataLoader with pin_memory + workers

from torch.utils.data import DataLoader, TensorDataset

ds = TensorDataset(torch.randn(64, 8), torch.randint(0, 2, (64,)))
loader = DataLoader(
    ds,
    batch_size=16,
    num_workers=4,        # parallel CPU loading
    pin_memory=True,      # host allocations are pinned -> fast async H2D
    persistent_workers=True,
)
for xb, yb in loader:
    xb = xb.to('cuda', non_blocking=True)  # fuse into training stream
    break

The CUDA Caching Allocator — "Leaky" Is a Lie

import torch

def mem():
    torch.cuda.synchronize()
    a = torch.cuda.memory_allocated()
    r = torch.cuda.memory_reserved()
    return a // 2**20, r // 2**20

print("start (alloc, reserved) MiB:", mem())
x = torch.randn(10_000, 10_000, device='cuda')   # ~800 MiB
print("after alloc  (alloc, reserved):", mem())
del x
print("after del    (alloc, reserved):", mem())
torch.cuda.empty_cache()
print("after empty_cache (reserved)  :", mem())
# reserved (driver allocations) stays high until empty_cache; that's normal.

✅ The allocator reuses freed blocks to avoid driver round-trips. High "reserved" ≠ leak. Empty only before measuring peak or OOM triage.

Peak memory accounting

torch.cuda.reset_peak_memory_stats()
model(torch.randn(64, 64, device='cuda'))
peak = torch.cuda.max_memory_allocated() // 2**20
print("peak activation+param MiB:", peak)

Async Pitfall: Synchronization Is Implicit & Expensive

GPU ops are queued; many "gotchas" come from hidden sync points (.item(), .numpy(), .tolist(), comparisons in eager Python, len()).

x = torch.randn(10000, 10000, device='cuda')
s = x.sum()
print(s.item())   # forces a device->host sync (stalls pipeline)

⚠️ In a hot loop, hoist .item()/.numpy() out. Use torch.cuda.Event for timing, not Python time.time() around each op.

def time_gpu(fn, n=20):
    torch.cuda.synchronize()
    s, e = torch.cuda.Event(True), torch.cuda.Event(True)
    s.record()
    for _ in range(n):
        fn()
    e.record(); torch.cuda.synchronize()
    return s.elapsed_time(e) / n   # ms

val = time_gpu(lambda: x @ x)
print(f"{val:.3f} ms/matmul")

Device Movement Best Practices

Scenario Pattern
Training loop Keep model on GPU; pin data, non_blocking=True
Checkpoint save Move to CPU (state_dict() stores CPU copies anyway)
Large cross-device transfer Split into chunks + non_blocking=True
MPS/Apple Silicon device='mps', same API, watch dtype support
Shape-only logic device='meta' to skip memory

Strategic device single-line pattern

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = x.to(DEVICE, non_blocking=DEVICE.type != 'cpu')
model = model.to(DEVICE)

Key Takeaways

  • The caching allocator hides the truth; use memory_allocated/reserved and max_memory_allocated.
  • Pinned memory + non_blocking=True unlocks async H2D copies that overlap compute.
  • Async GPU queues + implicit syncs (item()) are the classic perf killers — learn torch.cuda.Event.
  • Prefer meta device for shape/size inspection and peak-memory prediction.
  • Set a single DEVICE variable and pass non_blocking conditionally to stay portable.