Profiling & Benchmarking¶

Overview¶

You cannot optimize what you cannot measure. PyTorch ships torch.profiler (per-op CPU+GPU timing, memory traces, kernel flamegraphs) and integrates with NVIDIA tools (nsys, ncu). The discipline: profile first, hypothesize second, optimize third, re-measure fourth.

torch.profiler.profile(...) → table of op times, kernel names, memory.
ProfilerActivity.CUDA/CPU, schedule, on_trace_ready for trace export.
torch.cuda.memory._record_memory_history for allocator traces.
External: nsys profile (timeline/streams), ncu (per-kernel metrics).

💡 Wall-clock training is the ground truth. Profiles explain why — both matter.

The 30-Second Profile¶

import torch, torch.nn as nn
from torch.profiler import profile, ProfilerActivity, record_function

m = nn.Sequential(nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 16)).cuda()
x = torch.randn(128, 256, device='cuda')

def step():
    m(x).sum().backward()

# warm-up
for _ in range(3): step()
torch.cuda.synchronize()

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(10):
        with record_function("step"):
            step()
    torch.cuda.synchronize()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

Structuring a Profile¶

from torch.profiler import profile, ProfilerActivity, schedule

def trace_handler(p):
    p.export_chrome_trace("trace.json")      # open in chrome://tracing / Perfetto
    print(p.key_averages().table(sort_by="cuda_time_total", row_limit=20))

prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=2, warmup=2, active=5, repeat=1),
    on_trace_ready=trace_handler,
)
for step in range(30):
    train_step()
    prof.step()          # advance schedule
prof.stop()

Reading the Table¶

Column	Meaning
`Self CPU time`	CPU cost (launch + python)
`CUDA total`	device kernels under this op
`CUDA Mem`	allocations (device)
`# of Calls`	per-op call count

Look for: - Ops with high CUDA total but trivial math → algorithmic problem. - High CPU time with low CUDA → launch-bound (fix: compile/fuse/graphs). - copy_/to/contiguous() entries → layout conversions (Ch 04-03).

Memory Profiling¶

import torch
torch.cuda.memory._record_memory_history()

# run training steps...
snap = torch.cuda.memory._snapshot()
# visualize: python -m torch._C._profiler ... or use torch.profiler memory

# simple peak check
torch.cuda.reset_peak_memory_stats()
step()
print("peak MiB:", torch.cuda.max_memory_allocated() // 2**20)

Kernel-Level Tools (when you need more)¶

# Timeline + streams (system-wide view)
nsys profile -o run python train.py

# Per-kernel analyses (occupancy, memory, stalls)
ncu --set full python train.py

# torch profiler can also emit kernel lists
from torch.profiler import ProfilerActivity
with profile(activities=[ProfilerActivity.CUDA]) as prof:
    train_step()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=25))

Benchmarking Framework (thermal-safe & honest)¶

def benchmark(fn, n_iter=50, warmup=10):
    for _ in range(warmup): fn()
    torch.cuda.synchronize()
    times = []
    for _ in range(n_iter):
        t0 = torch.cuda.Event(True); t1 = torch.cuda.Event(True)
        t0.record(); fn(); t1.record()
        torch.cuda.synchronize()
        times.append(t0.elapsed_time(t1))
    times.sort()
    return {"median_ms": times[len(times)//2],
            "p95_ms": times[int(len(times)*0.95)],
            "min_ms": times[0]}

⚠️ Hot GPU throttles → take the median/min of repeated runs, not first.

Debugging Common Bottlenecks¶

Symptom	Likely cause	Next step
High CPU, low CUDA	python/launch bound	`torch.compile`, cudagraphs, fewer ops
High CUDA on a single op	kernel inefficiency	`ncu`, change layout/algorithm
High `copy_` count	format mismatch / `.cpu()` calls	fix layouts, remove syncs
Memory OOM at constant step	activations retained	checkpointing, gradient accumulation
DataLoader stalls visible in timeline	I/O bound	`num_workers`, prefetch, pin_memory

Key Takeaways¶

Profile before optimizing; 5 minutes of profiling saves hours of guessing.
Use torch.profiler for op-level CPU/GPU and memory; nsys/ncu for kernel depth.
Distinguish launch-bound (fix with fusion) vs compute-bound (fix with algorithm/layout).
Benchmark with warmup + median (thermal throttling skews means).
Trace export (chrome trace) is invaluable for stream/sync analysis.