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.

-

Precision & Numerics
[Profiling & Benchmarking](/06-pytorch/04-performance-and-compilation/(04-profiling-benchmarking/)
[Activation Checkpointing](/06-pytorch/08-advanced-training-techniques/(03-activation-checkpointing-memory-restructuring/)