PyTorch on GPU

This page covers GPU-specific optimization for PyTorch. For install/tensors/dataloaders/checkpoint basics, see PyTorch basics.

mixed precision / AMP

with torch.autocast(device_type='cuda', dtype=torch.bfloat16): out = model(batch)| https://pytorch.org/docs/stable/amp.html | runs matmul/conv ops in bf16 while keeping master weights/reductions in fp32 — bf16 needs no loss-scaling, unlike fp16 autocast |'ptg_amp1'

scaler = torch.cuda.amp.GradScaler() # only needed for fp16, not bf16| https://pytorch.org/docs/stable/notes/amp_examples.html | fp16's narrow dynamic range can underflow small gradients to zero; GradScaler rescales the loss before backward to keep gradients representable |'ptg_amp2'

torch.compile

model = torch.compile(model)| https://pytorch.org/docs/stable/generated/torch.compile.html | traces the model into an FX graph, fuses eligible ops via Inductor/Triton, and caches compiled kernels per input shape |'ptg_comp1'

model = torch.compile(model, mode='max-autotune')| https://pytorch.org/docs/stable/generated/torch.compile.html | spends extra compile time benchmarking multiple kernel variants per shape before picking one — slower first call, faster steady state |'ptg_comp2'

Expect a slow first iteration (or first per new input shape) while it traces and compiles; variable/dynamic input shapes without dynamic=True can trigger a recompile every time, silently erasing the speedup.

profiling

with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]) as prof: model(batch)| https://pytorch.org/docs/stable/profiler.html | captures a per-op CPU+GPU timeline, exportable as a Chrome trace to see gaps, launch overhead, and which ops dominate wall time |'ptg_prof1'

print(prof.key_averages().table(sort_by='cuda_time_total'))| | quick tabular summary ranked by GPU time — usually the first thing to check before reaching for Nsight/rocprof |'ptg_prof2'

memory: the caching allocator, pinned memory, non_blocking

PyTorch's CUDA/HIP caching allocator holds onto freed memory instead of returning it to the driver, so repeated alloc/free of similarly-sized tensors is fast after warmup — but it means nvidia-smi/rocm-smi memory usage reflects the allocator's reserved pool, not just what's "live."

torch.cuda.memory_summary()| https://pytorch.org/docs/stable/generated/torch.cuda.memory_summary.html | breaks down allocated vs reserved memory per pool — the tool for diagnosing OOMs that don't match the model's expected footprint | 'ptg_mem1'

x = tensor.pin_memory(); x_gpu = x.to(device, non_blocking=True)| https://pytorch.org/docs/stable/notes/cuda.html#use-pinned-memory-buffers | pinned (page-locked) host memory allows async H2D copies that overlap with compute; non_blocking is a no-op without pinned memory | 'ptg_mem2'

data loading bottlenecks

DataLoader(dataset, num_workers=8, pin_memory=True, persistent_workers=True, prefetch_factor=4)| https://pytorch.org/docs/stable/data.html | the standard combination to keep the GPU fed: parallel CPU workers, pinned staging buffers, workers that survive across epochs, and lookahead prefetch |'ptg_data1'

If GPU utilization (check via nvidia-smi dmon or rocm-smi) sits well below 100% during training, the data pipeline — not the model — is usually the bottleneck.

multi-GPU: DDP basics

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])| https://pytorch.org/docs/stable/notes/ddp.html | each process owns one GPU and a full model replica; DDP overlaps gradient all-reduce (via NCCL/RCCL) with backward computation instead of waiting for it to finish first |'ptg_ddp1'

DDP scales well because the only cross-GPU communication is gradient all-reduce, which is bandwidth-bound and overlappable — contrast with naive data-parallel approaches that gather full outputs to one GPU and serialize.

why PyTorch can be slow

Two separate costs stack on top of raw kernel time: Python/dispatcher overhead (every op call walks through Python, autograd, and dispatch machinery before a single GPU instruction issues) and kernel launch overhead (each of those tiny ops is a separate launch, and a model with thousands of small ops can be launch-bound rather than compute-bound). Fusion (kernel fusion, torch.compile) attacks exactly this by turning many small launches into fewer, larger ones.

reference

PyTorch Performance Tuning Guide
torch.compile tutorial
PyTorch Profiler recipe