Profiling with torch.profiler
torch.profiler records what every operation actually cost, on CPU and GPU, so you optimise the real bottleneck instead of the one you guessed.
- 8 min read
- 3 reading levels
- Published
Read these first
On this page 5
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
A profiler is a stopwatch attached to every single operation, so you learn where the time truly goes.
A restaurant is serving food late. The owner is sure the cook is slow, scolds him, buys a faster stove. Nothing improves. Then someone stands in the kitchen with a stopwatch and times every station. The cooking takes four minutes; the billing counter takes eleven. The bottleneck was never the kitchen.
Performance work without measurement is scolding the wrong cook.
Why it exists
Training feels like one blurry "step" that takes some amount of time. Inside it, dozens of very different activities happen: loading data, moving it to the GPU, the forward pass, the backward pass, the weight update, logging.
Intuition about which part is slow is famously bad. People rewrite their model to be faster when the real problem was the data loader keeping the GPU waiting. The profiler replaces guessing with a table: every operation, how many times it ran, and what it cost.
How it works
your code: load -> forward -> backward -> update
| | | |
profiler: [ stopwatch on every operation, both CPU and GPU ]
|
v
a table: name | times called | total time
a timeline picture you can open in a browserYou wrap the suspicious region, run it a few times, and print the table. The top rows are your real bottleneck.
A real example you have seen
Phone battery settings do this for energy. You assume the game ate your battery; the screen says the browser did. Same lesson: measured truth beats confident feeling. The profiler is that settings page for your training step.
Remember this
- Measure first, optimise second — intuition picks the wrong bottleneck often.
- The profiler times every operation, on both CPU and GPU.
- Optimising anything that is not near the top of the table is wasted work.
What to learn next
- Timing GPU code correctly — the manual stopwatch, done right.
- Is the GPU waiting for data? — the most common trace diagnosis, fixed.
- torch.compile — the standard cure for launch-bound traces.
Developer — Code and libraries.
Setup
pip install torchProfiling works on CPU alone, and adds GPU timing when CUDA is present. Outputs captured with torch 2.5.1 (CPU table on a Linux machine, GPU table on an NVIDIA RTX A6000). Your times will differ; the ranking is what you read. Named CUDA kernel rows such as ampere_sgemm_128x64_tn come from CUPTI, which ships with the Linux builds; Windows wheels report the aten:: operators only, so that row is missing there.
Profile a forward pass on CPU
import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(256, 1024), nn.ReLU(), nn.Linear(1024, 10))
x = torch.randn(64, 256)
with profile(activities=[ProfilerActivity.CPU]) as prof:
with record_function("my_forward"):
for _ in range(10):
model(x)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=5))---------------------- ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg # of Calls
---------------------- ------------ ------------ ------------ ------------ ------------ ------------
my_forward 15.19% 3.694ms 100.00% 24.309ms 24.309ms 1
aten::linear 3.53% 859.183us 72.57% 17.643ms 882.128us 20
aten::addmm 44.05% 10.708ms 64.70% 15.727ms 786.353us 20
aten::copy_ 20.44% 4.969ms 20.44% 4.969ms 248.464us 20
aten::relu 2.49% 605.255us 12.23% 2.973ms 297.315us 10
---------------------- ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 24.309msTwo columns matter. Self CPU is time spent in that operation itself; CPU total includes everything it called. aten::linear has small self time but large total — it is a wrapper whose real work happens in aten::addmm, the matrix multiply. record_function("my_forward") plants your own label in the table, which is how you profile "my attention block" rather than reading raw operator soup.
Add the GPU
import torch
import torch.nn as nn
from torch.profiler import profile, ProfilerActivity
if not torch.cuda.is_available():
raise SystemExit("needs a GPU for the CUDA timeline")
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(1024, 4096), nn.ReLU(), nn.Linear(4096, 1024)).cuda()
x = torch.randn(512, 1024, device="cuda")
model(x) # warm-up outside the profile
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for _ in range(10):
model(x)
torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=5))------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls
------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
aten::linear 3.58% 242.419us 36.40% 2.465ms 123.256us 0.000us 0.00% 4.814ms 240.700us 20
aten::addmm 9.46% 640.696us 31.30% 2.120ms 106.007us 4.582ms 94.19% 4.814ms 240.700us 20
ampere_sgemm_128x64_tn 0.00% 0.000us 0.00% 0.000us 0.000us 4.569ms 93.91% 4.569ms 228.434us 20
aten::relu 0.59% 39.976us 2.44% 165.189us 16.519us 0.000us 0.00% 282.785us 28.279us 10
aten::clamp_min 1.02% 69.209us 1.85% 125.213us 12.521us 282.785us 5.81% 282.785us 28.279us 10
------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 6.773ms
Self CUDA time total: 4.865ms(The table is wide — scroll it sideways.) The new columns time the GPU itself. Note the row ampere_sgemm_128x64_tn: that is the actual CUDA kernel the matrix multiply launched, taking 94% of GPU time. Note also that CPU time (6.8 ms) and CUDA time (4.9 ms) overlap rather than add — the CPU queues work and moves on, as explained in the timing lesson.
For a picture instead of a table, add prof.export_chrome_trace("trace.json") and open the file at chrome://tracing or ui.perfetto.dev — gaps in the GPU row are the smoking gun for input-bound training.
Common mistakes
Profiling the first iterations. The first steps pay one-time costs: allocator warm-up, cuDNN picking algorithms, torch.compile compiling. Warm up first, or use torch.profiler.schedule(wait=1, warmup=2, active=3) to skip them automatically.
Reading Self CPU for GPU work. A kernel launch costs microseconds of CPU while the kernel runs milliseconds on the GPU. Sort by cuda_time_total when the question is about GPU speed.
Leaving the profiler on. It adds overhead and memory (it records everything). Profile short representative windows, never whole training runs.
Profiling with tiny toy inputs. Kernel choice depends on shapes. Profile at your real batch size and sequence length, or the ranking you see is not the ranking you have.
Try it yourself
Add profile_memory=True to the CPU profile and sort the table by self_cpu_memory_usage. Find which operation allocates the most. Then wrap only the ReLU in its own record_function label and find it in the table.
What to learn next
- Timing GPU code correctly — the manual stopwatch, done right.
- Is the GPU waiting for data? — the most common trace diagnosis, fixed.
- torch.compile — the standard cure for launch-bound traces.
Researcher — Mathematics and papers.
What is actually recorded
The profiler is built on Kineto. CPU-side events come from ATen operator callbacks (RecordFunction), capturing op name, input shapes (with record_shapes=True), FLOP estimates (with_flops=True) and Python stacks (with_stack=True). GPU-side timing uses CUPTI activity records — actual kernel start/stop timestamps from the device, not host-side wrappers — which is why profiler CUDA times are trustworthy where naive time.time() is not.
Correlation IDs link each launch to its kernel, letting the trace viewer draw launch→execution arrows. The Chrome trace also exposes stream assignment, memcpy engines, and (with experimental_config) CUDA graphs and NCCL collectives — the tool of first resort for debugging distributed hangs that are really stragglers.
Reading a trace like an engineer
Three canonical signatures:
- Input-bound: periodic gaps on the GPU stream aligned with
DataLoaderactivity on CPU threads. Fix in the data pipeline, not the model. - Launch-bound: thousands of microsecond-scale kernels, CPU row saturated with launches, GPU idle between them. Fix with fusion (torch.compile) or CUDA graphs.
- Compute-bound: one or two kernels dominate
Self CUDA. Fix with better kernels (tensor cores, AMP, SDPA) or less compute.
Amdahl governs the payoff: accelerating a fraction $f$ of runtime by factor $s$ yields overall $\left((1-f) + f/s\right)^{-1}$. Profiling exists to find the largest $f$ before you spend effort on any $s$.
Overhead caveat: instrumentation inflates CPU-side times of very short ops (sub-10µs) and can perturb the launch pipeline; treat microsecond-level CPU numbers as indicative, and confirm end-to-end wins with proper wall-clock timing.
References
- PyTorch Kineto (github.com/pytorch/kineto) — the collection library, with CUPTI integration details.
- NVIDIA CUPTI documentation — the activity API semantics behind device timestamps.
- Amdahl (1967), Validity of the single processor approach to achieving large scale computing capabilities — the ceiling on every optimisation.
What to learn next
- Timing GPU code correctly — the manual stopwatch, done right.
- Is the GPU waiting for data? — the most common trace diagnosis, fixed.
- torch.compile — the standard cure for launch-bound traces.