Performance: What Is the Machine Waiting For?

Turn PyTorch performance into a measured comparison by defining the workload and metric, correcting CUDA timing, separating training phases, tracing allocated and reserved memory, profiling bottlenecks, and testing AMP, checkpointing, SDPA, and torch.compile as evidence-backed interventions.

Here is a benchmark. A model, a batch, a loop, a timer. Nothing exotic.

model.train()
t0 = time.perf_counter()
for _ in range(20):
    optimizer.zero_grad(set_to_none=True)
    loss = F.cross_entropy(model(x), y)          # MLP, batch 256
    loss.backward()
    optimizer.step()
ms_per_step = (time.perf_counter() - t0) / 20 * 1000

naive loop of   20:   9.491 ms/step reported   + 132.0 ms still queued at the final sync   -> really 16.093 ms/step
naive loop of  100:  14.692 ms/step reported   + 138.8 ms still queued at the final sync   -> really 16.080 ms/step
naive loop of  400:  15.741 ms/step reported   + 132.1 ms still queued at the final sync   -> really 16.071 ms/step

synchronized + warmed, median of 400 :  16.430 ms/step   <- the real number
one step, host-visible return only   :   4.866 ms/step
first training step ever             :  20.890 ms

Read the first line again. The loop ran twenty iterations and reported 9.491 ms per step. Then a single torch.cuda.synchronize() immediately afterward blocked for 132 more milliseconds — roughly eight more steps’ worth of work that the GPU had not finished when the timer stopped. Loop over a hundred iterations instead and the same code reports 14.7. Over four hundred, 15.7. Time one step in isolation and it looks like 4.9.

Nothing raised. Nothing is illegal. Every one of those numbers is a real measurement of something. None of them is the step time, which is 16.4 ms, and you only get that by warming the loop and synchronizing the device around the region you are timing.

Before optimizing a program, prove that the number you are optimizing measures the work you think it measures.

The reason the loop lies is that CUDA execution is asynchronous with respect to the host. model(x) and loss.backward() can return after launching device work rather than waiting for that work to complete. A short unsynchronized loop therefore measures mostly host-side submission. A longer loop can include backpressure and implicit waits as outstanding work and allocations accumulate, so its apparent time may drift toward the completed-work rate. Neither result is a trustworthy device-completion time unless the measurement boundary forces the queued work to finish.

And even the honest number is only honest about one question. Here is the same MLP and the same batch, timed four legitimate ways:

forward, inference_mode      median=    2.968 ms
forward, grad enabled        median=    2.967 ms
forward + backward           median=    8.094 ms
full training step           median=   16.450 ms

3 ms and 16 ms are both true. They answer different questions. “How fast is the model?” is not one question.

Where we are

Chapter 11 finished the learning chain. Given a fixed batch, it proved every link from the task contract to parameter movement:

TASK → OBJECTIVE → DEPENDENCY → GRADIENT → OPTIMIZER OWNERSHIP → PARAMETER UPDATE → CAPABILITY

and closed with one line: the model learns.

That is this chapter’s starting condition. We do not reopen the question of whether the computation is correct. We ask the other question:

At what cost?

A correct training loop can waste most of a GPU. It can fit in memory and spend half its time in the optimizer. It can run out of memory from a list that looks harmless. It can be slower after torch.compile than before it — and faster on one workload and neutral on another using the identical call. None of Chapter 11’s instruments would notice, because all of them measure whether the right thing happened and none of them measures what it cost.

The central question of this chapter:

When a correct PyTorch workload is slow or uses too much memory, which resource is limiting it — and what measurement proves that diagnosis?

The method, in the same slot where Chapter 6 put “split the step and measure each boundary” and Chapter 11 put “prove the chain one link at a time”:

Measure → localize → intervene → remeasure.

Each of those four words is load-bearing. Measure a number that means something. Localize the limiting resource before touching it. Intervene on one mechanism. Remeasure the same contract, and report what moved in every dimension, not just the one you hoped for.

The deepest principle

4.2 ms is not a fact about a program. It is the output of a function with six inputs:

workload + hardware + runtime + execution mode + metric + measurement boundary
= performance result

Change any input and the number changes, often by more than any optimization you are considering. The opening failure held workload, hardware and runtime fixed and moved only the measurement boundary — synchronized or not, warmed or not, loop of 20 or 400 — and got answers from 4.9 to 16.4.

So every measurement in this chapter is written down before it is taken. The recurring device is a measurement contract:

WORKLOAD     training step
MODEL        MLP, dim 1536, depth 12
INPUT        [256, 1536]
DEVICE       cuda
DTYPE        float32
MODE         train
EXECUTION    eager
METRIC       examples / second
WINDOW       steady state after warmup

A controlled comparison holds every non-target part of the contract constant. If the experiment deliberately changes batch size, dtype, execution mode, or another factor, name that factor and keep the rest fixed. A ratio becomes uninterpretable when several contract fields drift at once and you can no longer say which change caused it.

The environment these numbers came from

Every number, table and trace in this chapter was produced by running the code shown, with fixed seeds, on one machine, in a single back-to-back session:

Python        3.11.4
PyTorch       2.6.0+cu118        CUDA runtime 11.8, cuDNN 9.1.0
OS            Windows 11
CPU           24 logical cores; torch.get_num_threads() == 1 (OMP_NUM_THREADS=1 in the environment)
GPU           NVIDIA GeForce RTX 2060, compute capability 7.5, 12288 MiB, 34 SMs
torch.compile triton-windows 3.2.0, MSVC toolchain

This is a desktop GPU with a display attached. Before any experiment ran, nvidia-smi already reported about 2.6 GiB of the 12 in use by the compositor and other applications, and the card sat at 40–50% utilization at idle. That has consequences the chapter does not hide: run-to-run variation on a warmed benchmark is a few percent, and it is one-sided — a step is never faster than its floor but can be much slower when the desktop demands the GPU. The full-step baseline drifted between roughly 16.4 and 19.5 ms across sessions as that load varied. Within a session, held-constant comparisons were stable to about one percent, and where a difference was smaller than the spread this chapter says so rather than reporting a ratio.

The compute capability matters more than the raw speed. 7.5 is Turing: it has fp16 tensor cores but no native bfloat16 tensor-core path, and in this PyTorch 2.6 build the SDPA experiments below did not use the Flash-Attention backend. Several results are therefore specific to this hardware/software combination. On an A100 or an RTX 4090 the underlying mechanisms still matter, but even the qualitative shape of a timing or memory curve can change when a different kernel becomes available.

Absolute numbers here are not portable. The reusable artifact is the measurement procedure, not the milliseconds it produced on this card.

One controlled workload

Debugging performance without a fixed workload is the same mistake as debugging training without a fixed batch. This chapter uses one synthetic training workload for baseline timing, phase decomposition, the batch sweep, memory analysis, mixed precision, checkpointing, accumulation and compilation. It is a deliberately GEMM-bound classifier — a stack of square linear layers — chosen so one training step is tens of milliseconds (above timer noise) and repeats dozens of times per second (cheap to measure).

class MLP(nn.Module):
    def __init__(self, dim=1536, depth=12, num_classes=1000):
        super().__init__()
        layers = []
        for _ in range(depth):
            layers += [nn.Linear(dim, dim), nn.GELU()]
        self.body = nn.Sequential(*layers)
        self.head = nn.Linear(dim, num_classes)

    def forward(self, x):
        return self.head(self.body(x))

parameters: 29,866,984  (113.9 MiB in fp32)

A narrow attention microbenchmark, derived from Chapter 10’s manual implementation, appears once — for the sequence-length experiment, where the point is precisely about a tensor that the MLP never allocates. The GPT capstone is Chapter 15’s; it is not built early here.

The shared instruments live in perflab.py: fixed-seed model and batch builders, a timing helper that returns raw per-call samples with warmup excluded, a median/IQR summariser, the CUDA memory counters, and environment_report(). Every table in this chapter is reproducible from those.

The technique in one function: benchmark hygiene

The opening failure had three causes: no warmup, no synchronization, and a summary (the mean of one loop) that cannot survive an outlier. Fix all three.

def time_calls(fn, *, warmup=20, repeats=100, sync=True):
    for _ in range(warmup):
        fn()
    if sync:
        torch.cuda.synchronize()
    samples = []
    for _ in range(repeats):
        if sync:
            torch.cuda.synchronize()
        t0 = time.perf_counter()
        fn()
        if sync:
            torch.cuda.synchronize()
        samples.append((time.perf_counter() - t0) * 1000)
    return samples          # milliseconds, warmup already discarded

Warmup removes first-call costs that belong to the process, not the step: CUDA context creation, kernel loading, cuDNN algorithm selection, allocator growth, and — later — compilation. Measured on the controlled workload:

effect of warmup (repeats=100, synchronized):
  warmup= 0  median= 16.390  mean= 17.338  IQR= 0.159  min= 16.258  max= 96.243
  warmup= 1  median= 16.474  mean= 16.549  IQR= 0.166  min= 16.326  max= 17.777
  warmup=30  median= 16.480  mean= 16.608  IQR= 0.202  min= 16.347  max= 18.859

One 96.2 ms step — the first — pulls the mean of a hundred steps up by nearly a millisecond at warmup=0. The median barely moves, because a median ignores one bad sample. That is the argument for the median in one line.

Spread matters as much as the centre. Over four hundred warmed steps:

median 16.449   IQR 0.189   (p25..p75 = 16.383..16.572)
mean   16.525   std 0.242   max 18.030
steps above median+3*IQR: 21/400

Twenty-one steps out of four hundred sit well above the bulk. The distribution has a one-sided tail — the display GPU being borrowed — and the mean chases it. The interquartile range is 0.19 ms; the standard deviation is 0.24 ms. Report the median and the IQR. A difference smaller than the within-run spread is a warning that the effect may not be resolved; the IQR is not a confidence interval for the difference. For small effects, use repeated or interleaved A/B measurements and judge the difference against the paired/block-to-block variation rather than declaring significance from one ratio.

A benchmark utility buys one more thing: fewer hidden differences between measurements you intend to compare. torch.utils.benchmark.Timer handles its own warmup, synchronization and adaptive iteration count:

t = torch.utils.benchmark.Timer(stmt="step()", globals={"step": step})
m = t.blocked_autorange(min_run_time=2.0)   # -> Measurement with .median, .iqr, .mean, .times

Timer.blocked_autorange: median 16.169 ms   IQR 0.036 ms   n_per_run 10   runs 13
our helper (warmup=30, repeats=400): median 16.449 ms   IQR 0.189 ms

The Timer median is slightly lower and much tighter: blocked_autorange() chooses an inner block size that amortizes timer and accelerator-synchronization overhead, then reports per-call cost. Both measurements are valid; they measure marginally different boundaries. Timer can execute state-mutating statements, but automatic repetition means each timed call may see a different model or optimizer state unless you deliberately control/reset it. For the mutable full-training-step experiments in this chapter, the explicit helper keeps that state progression visible instead of hiding it inside an adaptive benchmark loop.

With warmup, synchronization and a robust summary, the controlled workload has a baseline: a full training step is about 16.4 ms, IQR about 0.2 ms within a session. Every intervention below is compared against that, remeasured under the same contract.

Localize before intervening: decompose the step

The step is 16 ms. Where does it go? Bracket each phase with a deliberate synchronize() and time it.

def phase_step(model, opt, get_batch):
    m0 = time.perf_counter()
    x_cpu, y_cpu = get_batch();                 torch.cuda.synchronize(); m1 = time.perf_counter()
    x = x_cpu.to("cuda", non_blocking=True)
    y = y_cpu.to("cuda", non_blocking=True);    torch.cuda.synchronize(); m2 = time.perf_counter()
    opt.zero_grad(set_to_none=True)
    logits = model(x);                          torch.cuda.synchronize(); m3 = time.perf_counter()
    loss = F.cross_entropy(logits, y);          torch.cuda.synchronize(); m4 = time.perf_counter()
    loss.backward();                            torch.cuda.synchronize(); m5 = time.perf_counter()
    opt.step();                                 torch.cuda.synchronize(); m6 = time.perf_counter()
    float(loss.detach());                                                 m7 = time.perf_counter()
    return {"input_wait": m1-m0, "host_to_device": m2-m1, "forward": m3-m2,
            "loss": m4-m3, "backward": m5-m4, "optimizer": m6-m5, "logging": m7-m6}

The synchronize() calls are not decoration. A wall-clock interval around asynchronous work does not measure the work; it measures submission. To attribute time to a phase you must wait for that phase’s device work to finish before reading the clock. That waiting is a real cost the instrument adds, and the experiment measures it too.

phase             median ms     IQR ms    share
input_wait           0.030      0.007     0.2%
host_to_device       0.311      0.020     1.8%
forward              3.047      0.077    17.4%
loss                 0.102      0.037     0.6%
backward             4.993      0.104    28.4%
optimizer            9.015      0.090    51.3%
logging              0.064      0.012     0.4%
sum of phases       17.560
measured total      16.758      0.260   (uninstrumented, one sync at the end)

Two checks first. The sum of phase medians (17.56 ms) is close to the uninstrumented step median measured separately (16.76 ms). The roughly 0.8 ms gap is evidence that the extra synchronization boundaries perturb the measurement; because medians from separate phase distributions are not algebraically additive, do not treat that difference as an exact synchronization bill to subtract. It is small enough here that the decomposition remains useful for localization. And input_wait is 0.03 ms because the batch is resident on the host; in a real pipeline this is where Chapter 6 lives, and if it dominated, the correct move would be to stop here and go back to Chapter 6, not to optimize anything in this chapter.

Now the result that is worth the whole experiment: the optimizer is 51% of the step. Not forward. Not backward. AdamW.step() on this model takes longer than the forward and loss combined. That is not a number anyone predicts, and it is exactly why the method says localize before intervene — a day spent making the forward pass faster would have moved 17% of the step.

Run the same brackets without the per-phase synchronize and the lie from the opening returns in a new costume:

same brackets, NO per-phase synchronize:
  forward              1.571 ms
  backward             2.173 ms
  optimizer            0.824 ms
  logging             11.683 ms

Forward looks nearly free, the optimizer looks free, and logging — the float(loss) call — appears to take 11.7 ms. It does not. float(loss) is the first operation that forces the host to wait for the device, so it absorbs the entire backlog of forward, backward and optimizer work that was still in flight. The measurement boundary moved; the time did not.

Two instruments, two questions

The synchronized wall clock is one instrument. CUDA events are another, and they measure something slightly different — the device-side duration between two points recorded into the stream, without the host-side launch and synchronization overhead:

start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
start.record(); fn(); end.record()
torch.cuda.synchronize()
device_ms = start.elapsed_time(end)

controlled workload (dim1536 depth12 batch256):
  synchronized wall clock : 16.281 ms  (IQR 0.181)
  CUDA events (device)    : 15.945 ms  (IQR 0.076)
  host submission only     : 15.155 ms/step
  wall - events            : +0.336 ms

The two agree to within 0.34 ms here because the step is device-bound: the host is blocked on the queue whether or not you ask it to synchronize, so elapsed host time and device time are nearly the same. On a workload where host submission is a real fraction of the step — many tiny kernels, heavy Python between them — the synchronized wall clock includes that host time and CUDA events do not, and the gap between them is itself a measurement: how much of the step is the host rather than the device. Neither instrument is “right.” Use the wall clock for “how long does my step take,” and CUDA events for “how long does the device spend,” and know which question you asked.

Intervene, remeasure

The optimizer phase is the bottleneck. Is it AdamW specifically, or optimization in general? Time .step() in isolation, with the same gradients present, under four optimizers:

optimizer phase, isolated (synchronized, warmed, median ms):
  SGD                1.407 ms   IQR 0.057
  SGD momentum       3.605 ms   IQR 0.084
  AdamW foreach      8.969 ms   IQR 0.115      <- the default
  AdamW fused        5.490 ms   IQR 0.051

AdamW with the foreach implementation executes a sequence of multi-tensor elementwise operations — multiply, add, addcmul, addcdiv, sqrt, lerp — over roughly 30 million parameters in 25 tensors, and on this GPU that costs about 9 ms. The fused implementation applies more aggressive horizontal and vertical fusion and measures 5.5 ms here; do not assume it literally becomes one kernel on every device/build. SGD takes 1.4 ms, but it is a different optimization algorithm, not a drop-in performance-only replacement.

This is a real intervention with a real tradeoff, and it carries a correctness column. fused=True implements the same AdamW update semantics through a different execution strategy, but floating-point evaluation need not be bit-identical: from identical state and identical gradients, one fused step matches one foreach step to 3.7e-9 per element (parameter scale ~2.6e-2) in this experiment. Switching to SGD changes the optimization dynamics entirely and belongs to Chapter 11’s analysis, not to a performance-only decision. The remeasurement: with fused=True the optimizer phase drops from 9.0 to 5.5 ms and the step falls with it, to roughly 13 ms. That is a defensible optimization for this contract. “Make the forward pass faster” was not.

Latency, throughput and peak memory are different questions

“Which batch size is fastest?” has no answer until you say fastest at what. The contract for this experiment holds everything except the one thing being swept:

WORKLOAD   training step        MODE       train
MODEL      MLP dim1536 depth12  EXECUTION  eager
INPUT      [B, 1536], B swept   METRIC     ms/step AND examples/s AND peak MiB
DEVICE     cuda / float32       WINDOW     steady state after warmup

Sweep the batch size, holding the model and its initialization fixed, and measure three things:

 batch    ms/step     IQR   examples/s  peak_alloc  peak_resv  ms/example
     32     13.150   0.566         2434       586.3      668.0      0.4109
     64     13.123   0.484         4877       814.3      932.0      0.2050
    128     13.226   0.240         9678       814.7      948.0      0.1033
    256     16.488   0.136        15526       815.9      926.0      0.0644
    512     23.392   0.172        21888       817.1      946.0      0.0457
   1024     37.466   0.407        27331       820.1      942.0      0.0366
   2048     66.080   0.264        30993       927.7     1130.0      0.0323

Read the columns separately, because they disagree.

Latency (ms/step) is flat at about 13.1 ms from batch 32 to 128, then climbs. Below 128, the step is dominated by fixed costs — kernel launches, the optimizer over 30M parameters — and the batch is almost free. Above 128, the GEMMs saturate the SMs and latency scales with the batch.

Throughput (examples/s) rises the whole way, but with sharply diminishing returns: doubling the batch from 1024 to 2048 buys 13% more throughput for 76% more latency and 13% more peak memory.

Peak memory barely moves from batch 64 to 1024 — 814 to 820 MiB — then jumps at 2048. For this shallow-wide model the activations are a small fraction; the peak is set by persistent state, which does not depend on the batch. (The lower value at batch 32 is an allocator artifact: a smaller workspace was reserved.)

The lowest-latency batch is 64 (13.1 ms). The highest-throughput batch is 2048 (30,993 ex/s). They are different batch sizes, and neither is “fastest.” A training run cares about throughput; an interactive inference service cares about latency; a memory-constrained fine-tune cares about the peak. One sweep, three answers.

Never collapse latency, throughput, peak memory and cold-start cost into “faster.” An intervention routinely improves one and worsens another.

Where memory actually goes

The peak is ~816 MiB at batch 256. Read torch.cuda.memory_allocated() after each phase, on the first step and on a steady-state step:

allocated MiB after each phase:
step 1 (cold)    start= 115.9  zero_grad= 115.9  forward= 164.1  loss= 165.0  backward= 247.6  optimizer= 474.6   peak= 588.6
step 3 (steady)  start= 474.1  zero_grad= 360.2  forward= 405.7  loss= 406.7  backward= 475.1  optimizer= 474.1   peak= 589.0

Several distinct pressures are now separable.

Persistent state for this configuration is about 456 MiB: parameters, one gradient allocation per parameter, and two AdamW moment buffers per parameter. It is 96% of the resident allocated footprint at this point. Batch size does not change those tensor counts, although dtype, optimizer choice, gradient-clearing policy and other training choices can.

Live forward state adds about 45 MiB between zero_grad and the return from forward. That is small here because the model is shallow and wide; the counter at a boundary is not a complete inventory of every temporary allocation that occurred inside the phase.

The first step is not the steady state. On step 1, the optimizer phase jumps from 248 to 475 MiB — consistent with the 227 MiB of AdamW moment state that appears on the first .step(). Every later step already has it. A memory measurement taken before the first optimizer step can therefore underreport steady-state training memory substantially.

The peak is not at any phase boundary. Every checkpoint above reads at most 475 MiB, but the peak is 589. That proves an additional ~114 MiB was live transiently inside one of the measured regions and freed before the next boundary. Given where the peak occurs in this workload, backward temporaries and overlapping saved activations/gradient allocations are plausible contributors, but the aggregate counter alone does not uniquely decompose those bytes. max_memory_allocated() reveals the hidden peak; a memory snapshot is the escalation tool if ownership of the peak matters.

OOM is an ownership-and-scaling question

When CUDA runs out of memory, the useful response is not a list of tricks. It is a decomposition:

PERSISTENT      parameters, gradients, optimizer state
SHAPE-DEPENDENT activations, temporary buffers, workspaces
ACCUMULATING    retained graph and tensor references across iterations
ALLOCATOR       reserved blocks, fragmentation, workspace behavior

and then a sequence of measurements that assign the failure to one of them:

The intervention has to follow the diagnosis. If persistent state dominates — as it does in the controlled workload, at 96% — then halving the batch size barely helps, because the batch was never the pressure. If allocated memory climbs across identical iterations, the fix is to find the retained reference, and shrinking the batch only delays the same failure. The next three experiments are that decomposition, one branch at a time.

Run it on a real report. Suppose a training job dies with CUDA out of memory and the evidence collected is:

Halving the batch from 32 to 8 moved the peak by 0.8 GiB — it barely responds, because 9.1 of the 9.9 GiB is still present at batch 8. The measured parameter bytes and first-step optimizer allocation already account for a large persistent component, while the flat repeated-step trace argues against iteration-to-iteration accumulation. The next intervention should therefore target persistent footprint or the remaining fixed workspaces before assuming batch-shaped activations are dominant. Reducing the batch may still help at the margin, but the measurement says it is not the main lever.

Parameter bytes are easy: sum(p.numel() * p.element_size() for p in model.parameters()) — 113.9 MiB for the controlled model. Training memory is not easy. It depends on the gradient dtype, the optimizer, whether master weights are kept, the optimizer-state dtype, the activation shapes, which tensors autograd saved, and workspace behavior. Do not assume gradients always add exactly one parameter-sized allocation, or that optimizer state is exactly 2×. Measure it: the counters above are the instrument.

There is no generally useful “typical training-memory breakdown” pie chart. The breakdown for this workload is roughly 24% / 24% / 48% / 4% across parameters, gradients, AdamW state and everything else — and that is a property of a 30M-parameter fp32 MLP under AdamW, not of training in general. A 7-billion-parameter model in bf16 with a fused optimizer and gradient checkpointing has a completely different profile. Measure your workload.

The accumulating branch: a retained reference

Here is a monitoring loop. It only reads the model — no training, no backward().

model.eval()                        # eval() is not no_grad()
kept = []
for i in range(60):
    out = model(x)
    kept.append(out.softmax(-1).amax(-1).mean())   # mean top-1 probability

allocated MiB, every 10 iterations:
kept.append(metric)            # forgot no_grad(); metric carries the graph
  i00=  164.1  i10=  542.8  i20=  922.1  i30= 1300.9  i40= 1680.2  i50= 2059.0   -> final 2400.3 MiB
  after empty_cache(): allocated 2400.3 MiB

with torch.no_grad(): ...       # graph never built
  i00=  124.1  i10=  124.1  i20=  124.1  i30=  124.1  i40=  124.1  i50=  124.1   -> final 124.1 MiB

kept.append(metric.item())     # keep a float, not a tensor
  i00=  164.1  i10=  164.1  i20=  164.1  ...                                     -> final 162.0 MiB

The metric is 0.001034 in all three variants — the computation is identical. What differs is what Python keeps alive. In the first variant, metric is a tensor still connected to the autograd graph of that iteration’s forward pass, and every element of that graph — the saved activations of all twelve layers, about 38 MiB — stays reachable as long as the list holds the scalar. Sixty iterations, sixty graphs, a linear climb to 2.4 GiB.

This is Chapter 3’s lesson at loop scale. It is not a leak: every byte is reachable from a live Python name. model.eval() sets a flag on the module tree; it does not disable autograd, so the graph is built whether you want it or not. torch.no_grad() is what stops the graph from being built. .item() is what keeps a number instead of a tensor.

And torch.cuda.empty_cache() does nothing here — 2400.3 MiB before, 2400.3 MiB after. empty_cache() returns unused cached blocks to the driver. These blocks are not unused; they hold live activations. No allocator call can free a tensor that Python still references.

When memory_allocated() climbs across otherwise-identical iterations, the set of live PyTorch allocations is growing. In this experiment the cause is a retained autograd graph; in an unfamiliar system it could also be intentional state or another growing owner. Trace ownership before blaming the caching allocator or merely shrinking the batch.

Allocated, reserved, and the third number

Three counters, three stories. Allocate a 400 MiB tensor, drop the Python reference, then call empty_cache(), reading everything at each step:

                                    allocated   reserved   max_reserved   nvidia-smi (this process)
steady state                           474.1      664.0        664.0        [N/A]
after allocating a 400 MiB tensor       874.1     1064.0       1064.0        [N/A]
after 'del big' (Python ref gone)       474.1     1064.0       1064.0        [N/A]
after torch.cuda.empty_cache()          474.1      542.0       1064.0        [N/A]

del big: allocated drops by 400 MiB (the tensor is freed), reserved stays at 1064 (the caching allocator keeps the block, so the next allocation of that size does not have to call into the driver).
empty_cache(): allocated is unchanged (no live tensor was freed), reserved drops to 542 (the now-unused block is returned to the driver). max_reserved still reads 1064 — it is a high-water mark, not a current value.
nvidia-smi could not attribute per-process memory at all on this WDDM Windows driver — it returned [N/A]. The overall display showed 2596 / 12288 MiB in use before the process started, from the compositor and other applications. That is a fourth number, and PyTorch does not control it.

reserved > allocated is not a leak. It is the caching allocator doing its job. nvidia-smi reporting a large number is not a leak either; it includes the CUDA context, the compositor, other processes, and reserved-but-unused blocks. No single counter equals “the memory my model uses.”

torch.cuda.empty_cache() releases unoccupied cached blocks held by PyTorch so that other GPU users can see/use that memory. It does not free live tensors and does not increase the amount of memory PyTorch could use for live allocations; current PyTorch documentation also notes that it may reduce fragmentation in some cases. That makes it an allocator-management tool, not a general leak repair. Repeatedly calling it in a hot training loop usually throws away useful cache reuse and changes the workload being benchmarked.

This chapter uses it deliberately between some isolated benchmark configurations to reduce carry-over from unused cached blocks. Even then, do not call the resulting allocator state “clean”: live tensors, CUDA context state, library workspaces and allocation history can still differ.

For deeper allocator forensics, torch.cuda.memory_summary() prints a detailed block report, and the underscore-prefixed torch.cuda.memory._record_memory_history() / _dump_snapshot() APIs capture allocation stack traces for offline inspection. Treat those as an escalation tool and as advanced, less-stable interfaces — aggregate counters first, snapshots only when the counters stop being enough.

The synchronization tax

loss.item() copies one scalar from the device to the host, and to do that it must wait for the work that produces loss to finish. Called every step, does it serialize the loop? Measure throughput, with the logged value checked unchanged, using an interleaved A/B so the effect is resolvable below the run-to-run drift:

device-bound step (dim1536 depth12 batch256, ~16 ms/step), interleaved A/B:
  none        median 16.201 ms
  item_every  median 16.633 ms
  difference  +0.432 ms/step (+2.7%)

small fast step (dim256 depth4 batch64, ~1 ms/step):
  no host read of any device scalar                    28292 ex/s   2.262 ms/step   rel 100.0%
  loss.item() every step                               27660 ex/s   2.314 ms/step   rel  97.8%

On the device-bound step the tax is +0.43 ms — about 2.7%, and only resolvable with an interleaved comparison because it is near the run-to-run noise floor. On the small fast step it is about 2%. The tax is small in both cases here because both steps are device-bound: the host was going to wait on the queue within a few iterations regardless. The tax is large only when the host could otherwise run far ahead of the device and a per-step sync stops it.

The pattern that removes the per-step sync without losing the number keeps the running total on the device and reads it only when it is time to log:

running = torch.zeros((), device="cuda")
for i in range(steps):
    loss = train_step()
    running += loss.detach()                 # stays on the device, no sync
    if (i + 1) % 50 == 0:
        print((running / 50).item())          # one sync every 50 steps
        running.zero_()

The logged running-mean loss is identical to five decimal places whether you read it every step, every fifty steps, or accumulate it on the device and read it once. So the lesson is not “.item() is slow” and not “never log.” It is:

Know where host/device synchronization occurs, and whether its frequency is part of the workload’s contract. Logging every step is a choice; logging every fifty steps is a different choice; measure which one you are making.

The same applies to print(), .cpu(), .tolist(), if loss > threshold, and anything else that needs a device value on the host.

The profiler localizes; it does not measure

When the coarse decomposition is not enough — you know backward is expensive but not which operator — the profiler attributes device time to individual kernels. Use a schedule so it records a few representative steps rather than hundreds, and add record_function ranges for your own boundaries:

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

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=5, warmup=5, active=10, repeat=1),
    profile_memory=True,
) as prof:
    for _ in range(20):
        with record_function("forward"):   logits = model(x)
        with record_function("loss"):      loss = F.cross_entropy(logits, y)
        with record_function("backward"):  loss.backward()
        with record_function("optimizer"): opt.step()
        opt.zero_grad(set_to_none=True)
        prof.step()

record_function phase ranges (total device time over 10 steps):
  forward      device_time_total   33.591 ms
  loss         device_time_total    0.870 ms
  backward     device_time_total   78.494 ms
  optimizer    device_time_total  114.727 ms

top operators by self CUDA time:
  backward (autograd)        78.3 ms
  aten::mm                   46.5 ms
  aten::addmm                27.1 ms
  aten::_foreach_mul_        22.5 ms      \
  aten::_foreach_addcdiv_    18.6 ms       |  the AdamW step,
  aten::_foreach_addcmul_    14.6 ms       |  death by a thousand
  aten::_foreach_add_        14.0 ms       |  elementwise kernels
  aten::_foreach_sqrt        11.5 ms      /

The profiler independently supports the phase decomposition’s surprise: the optimizer is the largest recorded phase, and its trace is dominated by a long sequence of multi-tensor elementwise operations rather than one GEMM-like kernel. That is localization — it makes a more fused AdamW implementation worth testing. The per-phase totals are themselves perturbed by profiling (11.5 ms/step for the optimizer against 9.0 ms unprofiled), and profiling can distort proportions as well as absolute times. Use the trace to form a mechanism hypothesis, then validate that hypothesis with an unprofiled intervention.

But look at the cost:

profiled step:   median 26.121 ms
unprofiled step: median 16.532 ms
profiler overhead: +9.589 ms (+58.0%)

The profiled step is 58% slower than the real one. The profiler’s job is to tell you where the time is, in proportion — not how much time the step takes. The workflow is fixed:

unprofiled benchmark  ->  profile a few representative steps  ->  hypothesis
->  change one mechanism  ->  unprofiled benchmark again

Never compare a profiled timing against an unprofiled baseline and call the difference a speedup. The profiler is an instrument for localization, and every instrument perturbs what it measures.

Mixed precision: an experiment, not advice

torch.autocast runs selected operations in a lower-precision dtype. Whether that helps is a question about this hardware and this workload, answered with three separate measurements — time, memory, and numerics — never merged into one verdict.

scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
with torch.autocast(device_type="cuda", dtype=dtype):   # float16 or bfloat16
    loss = F.cross_entropy(model(x), y)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()

The one-step numerics, against fp32:

fp32   loss 6.90711   grad-norm 0.06892
fp16   loss 6.90710  (Δ 1.05e-05)   grad-norm rel Δ 8.33e-05
bf16   loss 6.90735  (Δ 2.34e-04)   grad-norm rel Δ 1.70e-03

Both are within any reasonable training tolerance for one step; bf16’s larger deviation is its shorter mantissa. Over 100 steps the loss stays finite in all three. Now time and memory, at two batch sizes:

batch 256:   config             ms/step   examples/s   peak MiB
             fp32                16.469       15544       588.1     time 100.0%  mem 100.0%
             autocast fp16       15.722       16283       588.1     time  95.5%  mem 100.0%
             autocast bf16       22.301       11480       588.1     time 135.4%  mem 100.0%

batch 2048:  config             ms/step   examples/s   peak MiB
             fp32                66.262       30907       699.7     time 100.0%  mem 100.0%
             autocast fp16       26.590       77020       603.9     time  40.1%  mem  86.3%
             autocast bf16       94.186       21744       603.9     time 142.1%  mem  86.3%

The same autocast(dtype=torch.float16) call is nearly neutral at batch 256 (95.5% — a few percent, at the edge of the noise) and transformative at batch 2048 (40% of the fp32 time, 2.5× throughput). At batch 256 the step is optimizer-bound and the fp32 optimizer is untouched by autocast, so speeding up the GEMMs moves almost nothing. At batch 2048 the GEMMs dominate the step and Turing’s fp16 tensor cores halve them.

bfloat16 is 35–42% slower at both sizes, because compute capability 7.5 has no bf16 tensor-core path — the API accepts it and the kernels fall back to a slower route.

Peak memory only drops (86%) at batch 2048, where activations are a real fraction; at batch 256 the peak is persistent fp32 state, which autocast does not change. This is the same finding as the memory-by-phase experiment, arriving from a different direction.

“AMP should make this faster” is a hypothesis. The answer is three numbers — time, memory, numerics — on this hardware, at this batch size.

Activation checkpointing: earn it with a measurement first

Activation checkpointing recomputes activations during backward instead of storing them. It trades compute for memory. Before adding it, establish whether activation memory is a problem — the memory-by-phase experiment already showed it is ~45 MiB of an 816 MiB peak at batch 256, so the answer there is no. Measure it anyway, at two batch sizes:

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    for block in self.blocks:
        x = checkpoint(block, x, use_reentrant=False) if self.training else block(x)
    return self.head(x)

Same instrument, opposite verdicts. At batch 256, checkpointing costs 33% more time and saves nothing — the activations it drops were not what set the peak. At batch 4096, it costs 31% more time and saves 288 MiB (28%), which is the difference between fitting and not fitting. The verdict is not a property of checkpointing; it is a property of whether activations dominate the peak, which is a measurement you take first.

For this deterministic MLP, checkpointing reproduced the measured loss and gradient norm exactly. Do not promote that observation into a universal invariant: checkpointing re-executes a forward region during backward, so randomness, mutable global/module state, device movement, or a forward/recompute mismatch can change behavior. PyTorch preserves relevant RNG state by default to make common stochastic cases agree, at a performance cost, and currently recommends passing use_reentrant=False explicitly because that implementation supports more autograd functionality.

Gradient accumulation: what it does and does not move

Splitting a batch of B into K micro-batches, scaling each loss by 1/K, and stepping once approximates the full-batch gradient:

opt.zero_grad(set_to_none=True)
for k in range(K):
    xs, ys = micro_batch(k)
    (F.cross_entropy(model(xs), ys) / K).backward()
opt.step()

gradient equivalence (accumulated K micro-batches vs one full batch of 4096):
  K=2: max abs diff 4.64e-10   (gradient scale ~1.70e-03)
  K=4: max abs diff 2.02e-09
  K=8: max abs diff 2.01e-09

  K   micro-batch   ms/step   peak MiB
  1         4096    121.115    1249.1     peak 100%   time 100%
  2         2048    126.096    1054.2     peak  84%   time 104%
  4         1024    129.921     892.7     peak  72%   time 107%
  8          512    138.977     838.5     peak  67%   time 115%

The accumulated gradient matches the full-batch gradient to floating-point roundoff (~2e-9 against a gradient scale of ~2e-3). Peak memory falls to 67% at K=8 — because the live activation set is now one micro-batch, not the whole batch — for a 15% time cost from the extra kernel launches and the smaller, less efficient GEMMs.

What does not change with K is the shape and storage requirement of the model parameters, gradient buffers and optimizer state. The accumulated gradient values are only numerically equivalent here — the table above already measured roundoff differences up to 2e-9, so they are not byte-for-byte identical. Accumulation reduces the live micro-batch-dependent footprint; it does not remove the model or optimizer state. The ~839 MiB observed at K=8 is therefore a practical floor for this experiment, not a proof that every remaining byte is persistent state: one micro-batch and temporary workspace still exist.

One correctness caveat: the equivalence assumes every operation is linear in the batch or independent across examples. BatchNorm, some contrastive losses, and anything else whose result depends on the other examples in the micro-batch will not reproduce the full-batch result. Verify the equivalence numerically for your model rather than assuming it.

Sequence length and what the kernel physically allocates

Attention’s cost is often described as “quadratic in sequence length.” Be precise about what that means. The logical relation is Tq × Tk — every query position is scored against every key position — and that is true of every attention implementation. What varies is what the kernel physically allocates.

Chapter 10’s manual implementation materializes the score matrix explicitly:

scores = (q @ k.transpose(-2, -1)) / math.sqrt(Dh)   # [B, H, T, T], stored for backward
weights = scores.softmax(dim=-1)                       # another [B, H, T, T]
out = weights @ v

F.scaled_dot_product_attention is a dispatcher: on CUDA it may select a fused/memory-efficient backend, or fall back to the math implementation depending on inputs and availability. In this PyTorch 2.6 / Turing experiment it selected a memory-efficient backend that did not materialize the same full score/weight tensors as the manual implementation. Compare that selected path with the manual version, forward and backward, over sequence length, with B=4, H=8, Dh=64:

     T  manual ms  manual MiB   SDPA ms   SDPA MiB   max|Δ out|
   256      0.951        58.3      1.349       38.3    5.66e-07
   512      3.111       168.3      4.840       62.4    7.75e-07
  1024     11.023       576.3     16.095      108.6    7.75e-07
  2048     42.616      2160.3     60.200      201.0    8.49e-07
  4096        OOM         OOM    234.921      353.8           -

The manual version’s peak memory quadruples rapidly as T doubles because it explicitly stores B·H·T·T score and probability tensors; B·H·T·T·4 bytes is 512 MiB at T=2048 for one fp32 matrix alone. At T=4096 the manual path OOMs. The SDPA backend selected on this machine grows far more slowly in measured peak allocation and runs T=4096 in 354 MiB. The outputs agree to about 8e-7.

What that selected backend does not remove is the dense pairwise computation: exact dense attention still has Tq × Tk interactions, so runtime can remain quadratic in sequence length even when memory is not. On this specific PyTorch 2.6 / Turing setup, the memory-efficient SDPA backend was selected rather than Flash Attention. At small T the manual matmul/softmax path is actually faster, which is another reminder that backend selection and fixed overhead are part of the workload-specific result.

The pairwise relation Tq × Tk is structural. What the selected kernel allocates for it is an implementation choice. Measure both, and do not assume a modern attention kernel materializes the score matrix.

torch.compile as an optimization candidate

With its default Inductor backend, torch.compile attempts to capture compilable regions of the Python/PyTorch program, optimize them, and generate or call efficient kernels. It may fuse pointwise work, call external libraries such as cuBLAS for operations that are already highly optimized, or leave graph breaks between compiled regions. Whether any of that helps is measured like any other intervention — with the contract held constant on both sides:

                    eager                 compiled
MODEL               same state            same state (deepcopy)
INPUT               [256, ...]            [256, ...]  same shape
DTYPE / DEVICE      float32 / cuda        float32 / cuda
MODE               eval, no_grad          eval, no_grad
SYNC BOUNDARY       synchronize both ends  synchronize both ends
MEASURED ITERS      N                     N   (never cold-compiled vs warm-eager)

and with an interleaved A/B so that machine drift on a shared GPU cannot invent a speedup that is not there.

def interleaved_ab(step_eager, step_compiled, blocks=12, inner=30):
    for _ in range(20): step_eager(); step_compiled()   # warm both
    torch.cuda.synchronize()
    a, b = [], []
    for _ in range(blocks):
        # time a block of eager, then a block of compiled, and repeat
        a.append(timed_block(step_eager, inner))
        b.append(timed_block(step_compiled, inner))
    return median(a), median(b)

Run it on two workloads: the controlled GEMM-bound MLP, and a pointwise-only chain (forty layers of tanh, scale, shift, mean-subtract, sigmoid — no large matmul).

[GEMM-bound MLP  (dim1536 d12, batch256, forward)]
  eager steady (interleaved median)   :    2.924 ms   (block spread 2.90..3.21)
  compiled steady (interleaved median):    2.959 ms   (block spread 2.94..2.98)
  speedup                             :    0.988x
  correctness max|Δ| vs eager         : 1.30e-08
  peak MiB  eager 130.0  compiled 130.0

[pointwise chain (dim4096 depth40, batch256)]
  eager steady (interleaved median)   :    9.009 ms
  compiled steady (interleaved median):    0.727 ms
  speedup                             :   12.399x
  correctness max|Δ| vs eager         : 4.35e-04
  peak MiB  eager 144.1  compiled 140.1

The identical torch.compile(model) call is neutral on the GEMM-bound MLP (0.99×, inside the block-to-block spread) and a 12.4× speedup on the pointwise chain. The mechanism explains the split: the MLP’s time is in large cuBLAS GEMM kernels that are already near-optimal and that Inductor calls into unchanged; the pointwise chain is two hundred tiny elementwise kernels whose cost is launch overhead and memory round-trips, and Inductor fuses them into essentially one kernel. Same method, two workloads, opposite verdicts, and a mechanism that predicts which is which.

The correctness columns are what make those performance numbers usable. The MLP output differs from eager by at most 1.3e-8; the pointwise chain differs by 4.3e-4. A compiler is allowed to choose numerically different floating-point schedules while implementing the same intended function, so neither difference should be labelled acceptable merely from the dtype. Whether 4.3e-4 is acceptable belongs to the computation’s tolerance or task-level contract; the chapter records the number instead of assuming the answer.

Compile cost is not one number. Run a script that compiles the pointwise chain in a fresh process, once with the on-disk Inductor cache wiped and once reusing it:

MODE=cold  first_compiled_call_s=10.945   steady_state_ms=0.892
MODE=warm  first_compiled_call_s= 2.542   steady_state_ms=0.885

On a cold Inductor cache the first compiled call is 10.9 s. A fresh process that reuses the on-disk cache pays 2.5 s. Steady state is the same either way. So the break-even — how many calls it takes for the steady-state gain to repay the compile cost — depends on which compile cost:

For a training run of a million steps, either is trivial. For a script that compiles, runs fifty inferences and exits, compilation is the entire runtime.

If the compiled result is disappointing — slower steady state, a new shape pausing to compile again, logs mentioning recompilation, graph breaks fragmenting the program — this chapter’s job is to classify it (compile cost dominates / steady state slower / new shapes / graph breaks) and stop. Chapter 13 owns graph breaks, guards, recompilation, dynamic shapes and backend selection. Chapter 12 answers only “did compilation improve this workload, and which cost did I measure?”

The performance ledger

Chapter 7 has a stage report, Chapter 8 a shape ledger, Chapter 10 an attention ledger, Chapter 11 a learning ledger. This chapter’s instrument, assembled now that the mechanisms are earned, records what a performance result must expose so that the reader can ask which resource boundary is dominating.

PERFORMANCE LEDGER   MLP dim1536 depth12   AdamW   CUDA

workload    train step | batch [256, 1536] | float32 | cuda | eager
benchmark   warmup 30 | repeats 400 | step latency 16.4 ms (IQR 0.19) | 15500 examples/s
phases      input wait 0.2% | h2d 1.8% | forward 17% | backward 28% | optimizer 51%
memory      peak allocated 589 MiB | peak reserved 664 MiB | persistent 456 MiB | activations 45 MiB
correctness loss finite | one fused-optimizer step within 4e-9/element of foreach
measurement synchronized at both ends | uninstrumented | first step excluded

The value is not the specific numbers. It is that a performance claim should say what workload was measured, where time accumulated, where memory peaked, and how the measurement was made — so that the next question (“which resource is the limit?”) has evidence behind it instead of a reflex.

Do not invent rows that cannot be measured honestly. On a machine with no CUDA device, the phase row still splits into forward/backward/optimizer and the memory row uses VmHWM from a subprocess, but the row says CPU and makes no fp16 or CUDA-events claim.

A worked pass: measure, localize, intervene, remeasure

The four words are a procedure. Run them once, end to end, on the controlled workload.

Measure. The benchmark is warmed and synchronized. Median step latency is 16.4 ms, IQR 0.19 ms, over four hundred steps. Contract recorded: train step, MLP dim 1536 depth 12, batch [256, 1536], float32, cuda, eager, metric ms/step, window steady state.

Localize. The phase decomposition says: forward 17%, backward 28%, optimizer 51%. The profiler, on ten representative steps, agrees on the ordering and adds detail: the optimizer phase contains a long sequence of _foreach_* elementwise kernels, none individually dominant, together 115 ms of recorded device time over the ten profiled steps. That trace supports a mechanism hypothesis — repeated memory traffic plus launch overhead across AdamW’s elementwise updates — but timing alone is not a hardware-bandwidth proof.

Hypothesis. If kernel/memory-traffic overhead in the foreach implementation is the reason the optimizer phase is large, a more aggressively fused AdamW implementation should reduce that phase while preserving the optimizer’s intended update semantics. torch.optim.AdamW(..., fused=True) is the intervention to test.

Intervene. One change: foreach=True → fused=True. Nothing else.

Correctness check. The parameter update after one step, from the same state and the same gradients, differs from the foreach update by at most 3.7e-9 per element in this experiment. That is strong one-step evidence that the performance intervention preserved the intended AdamW update closely enough for this test; it is not bitwise identity and it is not a guarantee that two long floating-point trajectories will remain numerically locked. For a real task, use an explicit tolerance or task-level equivalence criterion.

Remeasure, same contract. The isolated optimizer phase falls from 9.0 to 5.5 ms. End to end, in a back-to-back pair, the step falls from 19.1 to 14.4 ms — 24% lower latency, 32% higher throughput. (The 19.1 ms foreach number here ran above the earlier baseline because the desktop was busier during this pair; the paired comparison is the relevant evidence, not a ratio against a measurement from another session.)

Record the measured tradeoff. Latency down ~24%. Throughput up ~32%. The measured peak allocation did not increase in this workload. First-step behavior was excluded by warmup. Within this stated contract no measured dimension regressed, which is much narrower — and more defensible — than calling the optimization “free.”

Contrast that with the reflex it replaces. “The model is slow, add torch.compile” would have targeted the model body, while the phase decomposition said the optimizer was the largest cost. In the compile experiment below, the GEMM-bound model body was neutral at 0.99×; compiling only that model therefore would not address the measured optimizer bottleneck. The measurement is what pointed at the right mechanism.

Forensic mode versus performance mode

This book has leaned hard on instrumentation: hooks, detect_anomaly, the profiler, memory-history recording, verbose assertions, record_function everywhere. Every one of those costs time, and this chapter measured several of them:

instrument                              measured cost on the controlled step
per-phase torch.cuda.synchronize (x7)   +0.8 ms   (+5%)
torch.profiler with a schedule           +9.6 ms   (+58%)
.item() every step                       +0.4 ms   (+3%)

Keep two modes and never confuse their measurements:

FORENSIC MODE     maximize evidence: hooks, anomaly detection, profiler,
                  memory-history, verbose logging, per-phase synchronization
PERFORMANCE MODE  remove every instrument that is not part of production behavior

Do not optimize the instrumented version and assume the production workload changed by the same amount. An optimization that halves the profiled step may do nothing to the real one, and an instrument you left enabled is a regression you shipped.

Without a CUDA device

Everything above assumed a GPU. On CPU the questions are the same and the instruments shrink:

The opening failure is not asynchrony — CPU work is synchronous — but a cold-versus-steady-state trap: an un-warmed benchmark of eager versus torch.compile can report the compiled version several times slower, purely because the first call includes compilation. A warmed, interleaved comparison shows the steady-state difference, which for a GEMM-bound step is usually inside the noise.
Phase decomposition, batch sweeps, profiling, torch.compile, CPU autocast, checkpointing and accumulation remain available, but device-specific timing and memory instruments must be replaced rather than copied unchanged.
CUDA’s memory_allocated() / max_memory_allocated() counters have no direct CPU equivalent. For process-level peak memory, isolate configurations in fresh subprocesses and use an OS-appropriate metric such as VmHWM on Linux; that measures process RSS high-water mark, not a PyTorch tensor-allocation ledger.
The CUDA-only mechanisms — memory_allocated versus reserved, empty_cache, CUDA events, the .item() synchronization tax, fp16 tensor cores — are verified against the installed PyTorch and given as runnable code, with the timing stated as not executed on that machine.

State the constraint once, early, and stop apologising for it.

Using AI on a slow workload

An assistant asked “how do I make this faster?” will produce a list — AMP, a bigger batch, torch.compile, more workers, empty_cache(), checkpointing. Every item is a real optimization for a real bottleneck. None of them is evidence about your bottleneck, and applying the wrong one moves nothing while costing a day.

The prompt below refuses the optimization and demands the reconstruction first. It is long on purpose; the structure is the value.

My PyTorch training step is slower (or uses more memory) than I expect.

Do not suggest optimizations yet. Do not mention AMP, torch.compile,
checkpointing, larger batches, empty_cache or more workers until a
measurement identifies a mechanism one of them could affect.

First, reconstruct the measurement contract from what I give you:
  workload (train / inference), model, input shape, dtype, device,
  execution mode (eager / compiled), metric (latency / throughput /
  peak memory), measurement window (cold / warm), synchronization boundary.
State which of these I have not specified.

Then classify the evidence I provide:
  - step latency, warmed and synchronized, with a robust spread
  - time by phase: input wait, host-to-device, forward, loss, backward,
    optimizer, logging  -- and whether the phase sum matches the total
  - memory: allocated, reserved and peak, phase by phase, and their trend
    across identical iterations
  - if compiled: eager steady state, compiled first call, compiled steady
    state, and whether new shapes trigger fresh compilation

Identify the FIRST measurement or resource boundary that is not yet
supported by evidence.
  - If the benchmark itself is invalid (no warmup, no sync, a metric that
    does not match the question), stop there. The benchmark is the bug.
  - Otherwise: state the bottleneck hypothesis, the smallest intervention
    that tests it, which metric should improve if the hypothesis is right,
    which correctness property must stay unchanged, and which competing
    explanation the experiment distinguishes.

Three clauses earn their place. The forbidden-interventions paragraph stops the assistant from reaching for the optimizations most discussed in its training data before any measurement justifies one. Asking which correctness property must stay unchanged is what makes the proposed experiment falsifiable rather than confirmatory — the mixed-precision, checkpointing and compilation experiments all carried a correctness column for exactly this reason. And asking which competing explanation the experiment distinguishes prevents “try AMP and see,” which can be rationalized whatever happens.

Ask AI to reconstruct the benchmark before asking it to optimize the code.

The performance debugging sequence

 1. Define the workload. Write the measurement contract: model, input shape,
    dtype, device, mode, execution, metric, window.

 2. Define the metric. Latency, throughput, peak memory and cold-start cost are
    different questions. Pick the one the job cares about.

 3. Prove the benchmark. Warm it. Synchronize around the region. Take repeated
    measurements and a robust summary (median, IQR). If the number moves with
    loop length or a stray sync, it is not measuring the work.

 4. Record the baseline. This is what every later change is compared against,
    under the same contract.

 5. Separate the phases. Input wait, host-to-device, forward, loss, backward,
    optimizer, logging. Synchronize deliberately at each boundary. Check the
    phase sum against the uninstrumented total.

 6. If input wait dominates, stop here and return to Chapter 6.

 7. Profile only after the coarse boundary is known. Use a schedule and
    record_function ranges. Measure the profiler's own overhead. Never compare
    a profiled timing to an unprofiled one.

  8. Distinguish allocated, reserved and peak -- phase by phase, and across
    repeated identical iterations. If allocated memory keeps climbing under a
    fixed workload, some set of live PyTorch allocations is growing; trace its
    owner. High reserved memory alone is not evidence of a live-tensor leak.

 9. Decompose any OOM: persistent state, shape-dependent memory, accumulating
    references, allocator behavior. Choose the intervention that matches the
    diagnosed source.

10. Form one mechanism hypothesis. Name the resource you believe is the limit.

11. Change one thing.

12. Verify correctness. A performance change that alters the computation is a
    regression. Use explicit tolerances for transformations expected to differ
    numerically (AMP, fused kernels, compilation).

13. Rerun the uninstrumented benchmark under the original contract.

14. Record the tradeoff across every dimension: latency, throughput, peak
    memory, cold-start cost. An improvement in one is not an improvement.

15. Only then test a larger intervention -- mixed precision, checkpointing,
    accumulation, compilation -- each as its own measured experiment.

16. If torch.compile itself becomes the unexplained problem -- seconds on the
    first call, a new shape pausing again, logs suggesting recompilation,
    graph breaks -- gather enough evidence to classify it and hand off to
    Chapter 13. Preserve the baseline and the contract.

What you should now be able to answer

Work through these before reading the resolutions.

“GPU utilization is low, so I should raise num_workers.” Utilization is a symptom with many causes: input wait, host-to-device transfer, host synchronization, launch and Python overhead, kernels too small to fill the SMs, bandwidth limits, gaps between kernels, or simply not enough parallel work. The measurement that comes first is the phase decomposition: what fraction of the step is actually input wait? If it is not the largest phase, workers are the wrong lever.

"nvidia-smi says memory is full, so PyTorch is leaking." nvidia-smi and PyTorch allocator counters answer different questions, and on this machine the display stack already consumed substantial GPU memory before the workload began. Inspect allocated, reserved and peak across a fixed workload. Flat allocated with high reserved can be normal caching-allocator behavior. If allocated keeps climbing across otherwise-identical iterations, some live PyTorch allocations are accumulating; then trace ownership instead of naming the cause from the counter alone.

"empty_cache() lowered the number, so it fixed the leak." It lowered reserved, by returning unused cached blocks to the driver. It did not change allocated, because it cannot free a live tensor. If your leak was a retained-graph list, empty_cache() did nothing to it — the monitoring loop climbed to 2.4 GiB and empty_cache() left it at 2.4 GiB.

“Batch 64 is slower than batch 32.” Its latency is about the same and its throughput is double and its peak memory is higher. “Slower” is not a property of a batch size until you name the metric.

"torch.compile is slower." Slower how? The first call is 2.5–10.9 seconds regardless of steady-state outcome. If you measured the first call, you measured compilation. If you measured warm steady state and it is still slower, that is a Chapter 13 question. The MLP compiled to 0.99× — neutral — and the pointwise chain to 12.4×, with the same call.

“AMP should make this faster.” On this Turing GPU, autocast(float16) was within a few percent at batch 256 and 2.5× at batch 2048; autocast(bfloat16) was 35–42% slower at both, because the hardware has no bf16 tensor cores. The answer is time, memory and numerics on your hardware at your batch size — three numbers, measured, not one assumption.

“The profiler says aten::mm is expensive.” It is 47 ms of self CUDA time over ten profiled steps — but the profiled step is 58% slower than the real one. The question is whether mm is a large enough fraction of the unprofiled metric you care about to be worth attacking. Here it is: the GEMMs are the second-largest real cost after the optimizer.

“My phase timings sum to more than the step.” Per-phase synchronization perturbs the workload, and the sum of separately measured phase medians is not mathematically required to equal the median of the whole step. Here the gap is about 0.8 ms, small enough that the decomposition still localizes the dominant phase. Do not simply subtract that gap as though it were an exact correction. If the discrepancy is large enough to change the diagnosis, use CUDA events, profiler ranges, or another less-perturbing attribution method and validate the intervention with an uninstrumented benchmark.

“I measured training memory before the first optimizer.step().” Then you measured it without the optimizer state, which AdamW allocates lazily on that first step — 227 MiB for the controlled model, 3.1 GiB in the constructed OOM. Always profile memory on at least the second step, and separate the cold step from the steady state.

“Doubling the sequence length quadrupled my memory.” That is the signature of an explicitly materialized [B, H, T, T] score matrix. It is a property of that implementation, not of attention. A fused or memory-efficient kernel keeps the same T×T logical relation while allocating far less; the compute stays O(T²) either way. Check which kernel is running before concluding anything.

“Compiled is 3% faster than eager.” On a shared machine, is 3% larger than the drift between the two measurements? Run an interleaved A/B — alternating blocks of each and pooling the samples — and look at the block-to-block spread. If 3% is inside it, the difference was not resolved, and reporting it as a speedup is reporting noise.

“How do I optimize this?” Show the measurement contract and the baseline first. Everything else is guessing.

Exercises

Reproduce the lying benchmark. Time the controlled training step in an unsynchronized loop of 10, 50 and 400 iterations, and after each loop call torch.cuda.synchronize() and record how long it blocks. Then time it warmed and synchronized. Explain why the short loop undershoots, why the long loop converges, and why both undershoot the synchronized number.
The four questions. Time one model and one batch as: forward under inference_mode, forward with grad enabled, forward plus backward, and a full training step. Explain what each number is for. Then add torch.compile and time all four again, separating first-call from steady state.
Phase decomposition and the optimizer. Decompose your own training step into seven phases with per-phase synchronization. Check the phase sum against an uninstrumented step. Identify the largest phase. If it is the optimizer, isolate .step() and compare SGD, SGD(momentum=0.9), AdamW(foreach=True) and AdamW(fused=True), and state which swap changes only performance and which changes the training dynamics.
Latency versus throughput. Sweep batch size across at least five values. Record ms/step, examples/s and peak allocated memory. Do not predict the curve shape in advance. Then answer three separate questions — lowest latency, highest throughput, largest batch that fits — and show they have three different answers.
Memory by phase, first step versus steady state. Read memory_allocated() after each phase, on step 1 and step 4, and also record max_memory_allocated(). Determine whether the peak is visible at a phase boundary or occurs transiently inside one. Quantify how much persistent state appears on the first optimizer step and remains on later steps.
Build the retained-reference leak. Write a monitoring loop that appends a graph-connected scalar to a list and watch allocated memory climb linearly. Repair it two ways — torch.no_grad() and .item() — and confirm the monitored quantity is unchanged. Then call empty_cache() on the leaking version and show it recovers nothing.
Allocated, reserved, nvidia-smi. Allocate a large transient, drop the reference, then call empty_cache(), reading all three counters at each step. Explain why del moves allocated but not reserved, why empty_cache() moves reserved but not allocated, and what nvidia-smi reports that neither PyTorch counter does.
Checkpointing at two batch sizes. Establish how much of your measured footprint appears batch/activation-dependent. Then compare normal against checkpointed at a smaller and a larger batch. Record peak memory, step time and a correctness check for all four. Do not assume the verdicts will differ; explain whatever tradeoff your measurements actually show.
Mixed precision, three numbers. Compare fp32 against autocast fp16 and bf16 (as your hardware supports). Report step time, throughput, peak memory and a one-step numerical check separately. Do the same at a small and a large batch size and explain any difference in verdict.
torch.compile, two workloads. Compile a GEMM-heavy model and a pointwise-heavy model. Use an interleaved A/B against eager. Record eager steady, compiled first call (cold and warm cache), compiled steady, correctness and peak memory. Derive the break-even for each. Explain the mechanism behind the different verdicts.

Next: compilation itself becomes the mystery

Chapter 12 leaves a model whose correctness Chapter 11 established, whose baseline is measured under a written contract, whose time is localized to a phase and an operator, whose memory is decomposed into persistent, shape-dependent, accumulating and allocator pressure, and against which candidate optimizations — mixed precision, checkpointing, accumulation, compilation — can be compared honestly, each with a correctness column.

The performance debugging move is complete:

define the workload  ->  prove the measurement  ->  localize the limiting resource
->  change one mechanism  ->  remeasure every dimension

But the last experiment opened a door it did not walk through. torch.compile was neutral on one workload and 12.4× on another. Its first call took between 2.5 and 10.9 seconds. Somewhere in there is a system with its own hidden structure — captured graphs, guards, recompilation — and when that system misbehaves, the performance method runs out of questions it can answer.

Chapter 12 asked: did compilation help?

Chapter 13 asks: what did the compiler capture, what assumption invalidated the graph, and why did it compile again?

That is a different investigation, and it starts where this one stops.

Performance is not a number you read from the machine. It is a comparison you construct. Define the workload, prove the measurement, locate the limiting resource, change one mechanism, and measure again.