The schedule decides what the table contains
Profiling every step of a run would record the noisy first steps along with the steady ones, so torch.profiler.schedule divides the steps into phases and the profiler only keeps one of them. Ask for wait=2, warmup=2, active=5, repeat=1 and the schedule is a plain function of the step index that you can call yourself before you profile anything.
Printing it for ten steps shows the shape directly: two NONE, two WARMUP, four RECORD, one RECORD_AND_SAVE, then NONE for anything past the single cycle. The last active step carries the save because that is where on_trace_ready fires.
The counterpart in your loop is prof.step(), one call per iteration, which is what advances that index. Forget it and the schedule never moves off its first phase. The proof that the schedule worked is in the table's call counts: ten steps ran, the train_step row says 5.
import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile, record_function, schedule
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 10))
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x, y = torch.randn(64, 512), torch.randint(0, 10, (64,))
loss_fn = nn.CrossEntropyLoss()
def step():
opt.zero_grad(set_to_none=True)
loss_fn(model(x), y).backward()
opt.step()
sched = schedule(wait=2, warmup=2, active=5, repeat=1)
print("phase per step:", [sched(i).name for i in range(10)])
with profile(activities=[ProfilerActivity.CPU], schedule=sched,
record_shapes=True, profile_memory=True) as prof:
for _ in range(10):
with record_function("train_step"):
step()
prof.step()
rows = {e.key: e.count for e in prof.key_averages()}
print("train_step rows kept:", rows["train_step"], "of 10 steps run")
# ---- stdout ----
# phase per step: ['NONE', 'NONE', 'WARMUP', 'WARMUP', 'RECORD', 'RECORD', 'RECORD', 'RECORD', 'RECORD_AND_SAVE', 'NONE']
# train_step rows kept: 5 of 10 steps runThe largest row in the table is not an op
Sort that capture by self CPU time and train_step is at the top with 31.20 percent, 3.094 milliseconds across five steps. Nothing in that row is an operator. Self time is what a region spent outside all of its children, so the annotation's self time is exactly the part of the step that ran between ATen calls: Python, module __call__, the autograd engine deciding what to run next, the optimizer's bookkeeping before it dispatches anything.
That is the host-side gap the chapter asks you to name, and here it is 0.619 milliseconds of a step whose whole recorded cost is 1.942. Across five runs of the same file the share stayed between 28.1 and 31.6 percent while the milliseconds moved by half. On an accelerator that same third of the step is time the device spends idle unless the host runs far enough ahead, which is why it is worth extracting rather than eyeballing.
The total column reads the opposite way and the two must not be confused. The row for autograd::engine::evaluate_function: AddmmBackward0 shows 1.56 percent self against 21.00 percent total, because the total includes the backward node and the aten::mm calls under it. Self answers where the time went; total answers what a subtree cost. Sorting by the wrong one is how a wrapper ends up looking like the bottleneck.
Read the counts as a check on your own model of the step. aten::addmm fires 10 times over 5 steps, which is the two forward linears. aten::mm fires 15, three per step, and the shape-grouped view names them: the second linear needs a gradient for its input and one for its weight, the first linear needs only the weight one, because the batch it was handed is a leaf that requires no gradient.
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
train_step 31.20% 3.094ms 97.92% 9.710ms 1.942ms 1.02 Mb -652.54 Kb 5
aten::addmm 16.13% 1.599ms 17.44% 1.729ms 172.900us 652.50 Kb 652.50 Kb 10
aten::mm 14.70% 1.458ms 14.70% 1.458ms 97.200us 5.72 Mb 5.72 Mb 15
Optimizer.step#SGD.step 7.95% 788.000us 11.88% 1.178ms 235.600us 0 b 0 b 5
aten::add_ 3.93% 390.000us 3.93% 390.000us 19.500us 0 b 0 b 20
Optimizer.zero_grad#SGD.zero_grad 2.63% 261.000us 2.63% 261.000us 52.200us -4.09 Mb -4.09 Mb 5
ProfilerStep* 2.08% 206.000us 100.00% 9.916ms 1.983ms 1.02 Mb 0 b 5
aten::t 1.94% 192.000us 3.52% 349.000us 7.756us 0 b 0 b 45
aten::sum 1.63% 162.000us 1.78% 177.000us 17.700us 10.20 Kb 10.20 Kb 10
autograd::engine::evaluate_function: AddmmBackward0 1.56% 155.000us 21.00% 2.082ms 208.200us 5.10 Mb -652.50 Kb 10
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 9.916msWhat each flag costs, and one that will not run here
Recording is not free, and the honest way to size the cost is with the instrument from lesson one rather than with the profiler's own totals. Timing the same step through Timer, inside and outside a profile() block, put the ratio at 1.29, 1.18 and 1.21 across three rounds of one script. Call it a fifth to a quarter, on this step, on this machine.
That is small enough to trust the proportions in the table and too large to quote the absolute milliseconds anywhere. The five profiled captures reported 1.638 to 2.377 milliseconds per step, while Timer put the same unwatched function at 1.420 to 1.619. Take shares of time from the profiler and take the time itself from a benchmark.
record_shapes and profile_memory earned their place in this capture, and neither one moved the ratio measurably. with_stack is the flag that would attribute every op to the Python line that made it, and on this pair of versions it refuses. A single profiled optimizer step raises an internal assert from the Python stack replay, reproducibly, at n equal to 1 as well as 50, while the same flag over a plain forward and backward with no optimizer succeeds. Treat it as a version fact about torch 2.2.2 with CPython 3.12, not as a statement about the flag.
$ python3 observer.py # Timer around the same step, three rounds, one script
run 1 plain 1.420 ms under profile() 1.838 ms ratio 1.29x
run 2 plain 1.427 ms under profile() 1.682 ms ratio 1.18x
run 3 plain 1.619 ms under profile() 1.959 ms ratio 1.21x
$ python3 with_stack.py # one optimizer step under with_stack=True
python 3.12.0 | torch 2.2.2
n=1: RuntimeError: !stack.empty() INTERNAL ASSERT FAILED at ".../torch/csrc/autograd/profiler_python.cpp":969, please report a bug to PyTorch. Python replay stack is empty.
n=5: RuntimeError: !stack.empty() INTERNAL ASSERT FAILED at ".../torch/csrc/autograd/profiler_python.cpp":969, please report a bug to PyTorch. Python replay stack is empty.
n=50: RuntimeError: !stack.empty() INTERNAL ASSERT FAILED at ".../torch/csrc/autograd/profiler_python.cpp":969, please report a bug to PyTorch. Python replay stack is empty.The biggest allocation in the step has a shape
With record_shapes=True the same capture can be grouped by input shape, and with profile_memory=True it carries bytes. Together they answer a question the flat table cannot: which single call allocated the most, and for which operand.
The winner in this step is an aten::mm on [[512, 64], [64, 512]], 5.00 Mb over five steps, one megabyte per step. Read the shapes and it names itself. That is the transposed activation times the incoming gradient, which is the weight gradient of the first linear, and it is larger than the activations because the weight is 512 by 512 while the batch is only 64 rows.
Two rows carry negative memory and they are not errors. Optimizer.zero_grad#SGD.zero_grad shows -4.09 Mb self, because set_to_none=True drops the gradient tensors and the profiler accounts the release to whoever released it. The train_step annotation nets out at -652.54 Kb for the same reason: over a whole step, this loop frees slightly more than it keeps.
None of these bytes are device bytes. The columns say CPU Mem, they come from the CPU allocator, and the third lesson is about why the CUDA equivalents on this machine answer zero rather than refusing.
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------------------------------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls Input Shapes
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------------------------------------
aten::mm 12.08% 1.198ms 12.08% 1.198ms 239.600us 5.00 Mb 5.00 Mb 5 [[512, 64], [64, 512]]
aten::addmm 13.89% 1.377ms 14.83% 1.471ms 294.200us 640.00 Kb 640.00 Kb 5 [[512], [64, 512], [512, 512], [], []]
aten::clamp_min 0.91% 90.000us 0.91% 90.000us 18.000us 640.00 Kb 640.00 Kb 5 [[64, 512], []]
aten::mm 1.04% 103.000us 1.04% 103.000us 20.600us 640.00 Kb 640.00 Kb 5 [[64, 10], [10, 512]]
aten::threshold_backward 1.11% 110.000us 1.11% 110.000us 22.000us 640.00 Kb 640.00 Kb 5 [[64, 512], [64, 512], []]
aten::mm 1.58% 157.000us 1.58% 157.000us 31.400us 100.00 Kb 100.00 Kb 5 [[10, 64], [64, 512]]
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------------------------------------
Self CPU time total: 9.916msEvery backward event carries a forward sequence number
Export the capture with export_chrome_trace and it is plain JSON, so you can check what landed without opening a viewer. This one holds 784 events over five steps: 540 operator events, 20 annotations, 146 memory events from profile_memory, and 70 events in a category called fwdbwd.
Those 70 are flow events, and they exist to answer the question a flat table cannot. Every operator event carries a Sequence number in its args, and the backward node that consumes it carries the same one. Sequence 29 holds aten::addmm with Fwd thread id 0 and AddmmBackward0 with Fwd thread id 1, which is the pairing drawn as an arrow when the trace is opened in Perfetto.
That integer is not new here. It is the same _sequence_nr() the autograd lesson at /pytorch/autograd/the-node-behind-grad-fn reads off a graph node, seen from the other end: the tape assigns it when the forward op runs, and the profiler stamps it on both events so a trace can be walked back into your own code.
A backward row you cannot attribute is a forward op you have not looked up yet.
# the tail of the capture script from the first section, same process
import collections
import json
prof.export_chrome_trace("trace.json")
ev = json.load(open("trace.json"))["traceEvents"]
print("trace events:", len(ev), "| by category:", dict(collections.Counter(e.get("cat") for e in ev)))
pairs = collections.defaultdict(list)
for e in ev:
if e.get("cat") == "cpu_op" and e.get("args", {}).get("Sequence number", -1) >= 0:
pairs[e["args"]["Sequence number"]].append((e["name"], e["args"]["Fwd thread id"]))
for seq in sorted(pairs)[:3]:
print("seq", seq, pairs[seq])
# ---- stdout ----
# trace events: 784 | by category: {'user_annotation': 20, 'cpu_op': 540, 'fwdbwd': 70, 'cpu_instant_event': 146, None: 7, 'Trace': 1}
# seq 28 [('aten::linear', 0), ('aten::t', 0), ('autograd::engine::evaluate_function: TBackward0', 1), ('TBackward0', 1)]
# seq 29 [('aten::addmm', 0), ('autograd::engine::evaluate_function: AddmmBackward0', 1), ('AddmmBackward0', 1)]
# seq 30 [('aten::relu', 0), ('autograd::engine::evaluate_function: ReluBackward0', 1), ('ReluBackward0', 1)]Check yourself
01 The top row of your profile is a record_function annotation, not an operator. What is in it?
Host work that ran between ATen calls: Python, module __call__, the autograd engine, optimizer bookkeeping. Self time excludes children, so that row is the host-side gap. In this capture it was 3.094 ms over five steps, 31 percent of the recorded step, and it stayed between 28 and 32 percent across five runs.
02 Why should you take shares of time from the profiler but not milliseconds?
Because recording costs a fifth to a quarter of this step: Timer put the ratio at 1.18 to 1.29, and the profiled captures reported 1.638 to 2.377 ms per step against an unwatched median of 1.420 to 1.619. Proportions survive that inflation, absolute numbers do not.
03 A backward row dominates your trace and you cannot tell which layer it came from. What do you look up?
The Sequence number in the event args. The forward op and its backward node share it, so sequence 29 pairs aten::addmm with AddmmBackward0 here. It is the same _sequence_nr the autograd lesson reads off a graph node, stamped onto both trace events.
Readings
- torch.profiler reference ↗ every argument this lesson set, and the ones it did not
- profiler.py at v2.2.2 ↗ schedule as a plain function of the step index, at lines 318 to 350, and the four actions above it
- Perfetto UI ↗ where the exported json opens; the fwdbwd events are the arrows between forward and backward