Performance

Time a PyTorch op the naive way, wrap it in time.perf_counter() and stop the clock a line later, and CPU tensors mostly tell the truth: CPU dispatch is synchronous, so the line doesn't return until the op has actually finished. CUDA breaks that assumption completely. A CUDA call returns to Python the moment it's queued, not the moment it runs, so an unsynchronized clock measures how fast you can enqueue work, not how fast the device does it. torch.cuda.synchronize() is what makes the clock wait for the device to actually catch up.

The TPU bridges chapter 10 covers add a third trap on top of the first two: those tensors are lazy, recording operations instead of running them until a sync point forces materialization, so a naive timer there measures almost nothing at all. Three backends, three different reasons the same lazy mistake produces a fast-looking, wrong number. One discipline covers all three: never trust a clock you haven't forced to wait.

This site's rule for every number it publishes carries over here without modification: state the chip, the dtype, the shapes, and the method, every time a latency claim gets made. A latency number with no provenance attached isn't a measurement. State the chip, the dtype, the shapes, and the method, or the number doesn't count.

The profiler is the referee, on every device.

The harness below is built to survive exactly the traps the last section named. A warmup call runs first and forces whatever compilation or first-time cost exists to happen once, off the clock. The timed loop then calls the function n times in a row, back to back, and only synchronizes once, at the very end, so the wait happens after every call has been dispatched rather than after each individual one.

Dividing the total by n gives you the per-call time with dispatch overhead spread thin across every iteration instead of dominating a single one. On this machine, on the CPU backend, that number for a 512-by-512 linear layer comes out to 0.142 milliseconds. On CUDA the same shape of harness needs torch.cuda.synchronize() added exactly where the comment marks it, or the number it reports is a lie about dispatch, not a fact about compute.

import time
import torch

def bench(f, *args, n=50):
    out = f(*args)                     # warmup: capture + compile if any
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    if out.device.type == "cuda":
        torch.cuda.synchronize()       # async dispatch settles before the clock
    return (time.perf_counter() - t0) / n

lin = torch.nn.Linear(512, 512)
print(f"{bench(lin, torch.randn(64, 512)) * 1e3:.3f} ms")   # 0.142 ms (this machine, CPU)

Wrapping a benchmark loop tells you the average cost of a call. It doesn't tell you where inside that call the time actually went, and for that torch.profiler is the tool: it records both host-side Python overhead and device-side kernel execution on the same timeline, so a single trace shows you the CPU launching work and the accelerator running it, side by side.

Profiling every single step of a long run would be wasteful and would skew the very thing you're measuring, so the profiler takes a schedule: a wait phase where nothing gets recorded, a warmup phase that primes caches and compilation without keeping the trace, and an active phase that's the only part actually captured. What lands in the trace is steady-state behavior, not the noisy first few steps every training loop pays once and never again.

Automatic mixed precision changes the dtype decision from something you hardcode per op into something torch.autocast picks for you, per operation, based on which dtype each op tends to be numerically safe and fast in. Matmuls and convolutions typically autocast to a lower-precision dtype; reductions and other numerically sensitive ops stay in float32 underneath the same context manager, without you naming either choice by hand.

fp16's narrow dynamic range means small gradients can underflow to exactly zero before they ever reach the optimizer, and GradScaler exists to prevent that: it scales the loss up before backward, so gradients land in a range fp16 can represent, then unscales them back down before the optimizer step. bf16 doesn't need any of that machinery. Its exponent range matches float32's, just with less mantissa precision, so no scaler is required, and bf16 is also the dtype the kernel path measures with natively on TPU, the hardware this whole track eventually lands on.