Profiling — Finding the Real Bottleneck

Don't optimize what you haven't measured

The PyTorch profiler captures op-level timing on CPU and GPU, exports a Chrome-trace JSON, and (with profile_memory=True) tracks tensor allocation. It's the right tool to answer "where is my training time actually going" and "why am I OOMing".

Two profilers, two purposes

torch.profiler.profile — modern, comprehensive. CPU + CUDA timing, memory tracking, Chrome trace export. The right default.
torch.utils.bottleneck — older, lighter-weight Python wrapper. Useful for quick "what's slow" answers without setting up the full profiler.

The Chrome trace

The profiler can export a JSON that loads in Chrome's chrome://tracing viewer (or perfetto.dev). You see a timeline of every CPU and GPU op, with their durations and call relationships. Spotting "I have a huge cudaStreamSynchronize" or "this tiny op dispatches 10,000 times" becomes obvious visually.

Reproducibility — orthogonal but worth knowing

If you want bit-for-bit deterministic runs (for debugging or paper-style ablations), you have to set every RNG seed AND tell PyTorch to use deterministic algorithms. Not free — some ops have no deterministic implementation, and others are slower in deterministic mode. Use it for debugging, not for production training.

Code

torch.profiler — basic CPU+GPU profile·python

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

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    with record_function("model_inference"):
        out = model(input_data)

# Top ops by GPU time
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

# Export Chrome trace — open in chrome://tracing or perfetto.dev
prof.export_chrome_trace("/tmp/trace.json")

Stepped profiler for training loops·python

import torch
from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity

# schedule(wait, warmup, active, repeat) — only profile some steps
sched = schedule(wait=2, warmup=2, active=4, repeat=1)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=sched,
    on_trace_ready=tensorboard_trace_handler('/tmp/profile'),
) as prof:
    for step, (x, y) in enumerate(loader):
        if step >= 10: break
        out = model(x.cuda())
        loss = criterion(out, y.cuda())
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        prof.step()                 # tells the profiler we crossed a step boundary

Memory snapshots — find the leak·python

import torch

# Begin recording allocator history
torch.cuda.memory._record_memory_history(max_entries=100_000)

# ... run your training for a while ...

# Dump a snapshot — visualize at https://pytorch.org/memory_viz
torch.cuda.memory._dump_snapshot('/tmp/memory.pickle')

# Quick numeric summary
print(torch.cuda.memory_summary(abbreviated=True))
print(f"allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
print(f"reserved : {torch.cuda.memory_reserved()/1e9:.2f} GB")

Reproducibility — turn the dials all the way·python

import os, random, torch
import numpy as np

def seed_everything(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Force deterministic behavior even at a perf cost
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.use_deterministic_algorithms(True, warn_only=True)
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'

seed_everything(42)

Exercise

Run torch.profiler on one training step of any model. Open the resulting Chrome trace in perfetto.dev. Find the longest GPU op and the longest CPU op. Try one optimization (compile? bigger batch? num_workers?) and re-profile. Document the before/after numbers — that habit is the difference between optimization and superstition.