Profiling

Problem

A training run reports 117.6 ms per step. Is that good? The GPU's spec sheet promises about 989 TFLOP/s; the model does a known number of FLOPs per token. Without a measurement you can trust, you cannot say whether the step is close to the hardware limit or wasting most of it, and without a breakdown you cannot say where the waste is.

Both are easy to get wrong. time.time() around a step measures how long Python took to queue the work, not how long the GPU took to do it. The first steps include one-off costs (CUDA context creation, cuDNN autotuning, allocator growth, compilation) that the average then carries forever. And a profiler that records every step slows the run and produces traces too large to open. Profiling is a set of habits that avoid these traps and point at the bottleneck.

Intuition

The CPU is a dispatcher handing work to the GPU through a queue. As long as the CPU stays ahead, the queue is never empty and the GPU is busy; a timer on the CPU stops long before the GPU finishes. If each piece of work is tiny, the GPU finishes it faster than the CPU can queue the next one, and the GPU idles between kernels no matter how fast it is. If the CPU stalls to fetch the next batch, the queue drains and the GPU idles for the length of the stall.

So the questions to ask a trace are visual: are there gaps in the GPU row, and what is the CPU doing during them? A gap under a data-loader call is an input stall; a stream of tiny kernels with gaps between them is host overhead; a solid GPU row means the step is compute-bound, and only then do kernel-level optimizations (lower precision, fused kernels) pay off.

Mechanics

Measuring correctly

CUDA calls return as soon as the work is queued. The PyTorch notes on asynchronous execution give the two correct ways to time GPU work:

• Synchronize, then read a host clock. torch.cuda.synchronize() blocks until every queued kernel on the device has finished. Call it before starting the timer and before stopping it. The benchmark utility's own timer() does exactly that: torch.accelerator.synchronize(), then timeit.default_timer().
• CUDA events. A torch.cuda.Event(enable_timing=True) is a marker placed in the stream by record(); elapsed_time() returns the GPU time between two of them in milliseconds. Events measure the device's own timeline and need a sync only before reading the result, so they disturb the pipeline less.

Setting CUDA_LAUNCH_BLOCKING=1 makes every launch synchronous. It is for debugging (errors are reported at the line that caused them) and never for timing, since it removes the overlap a real run has.

Then discard the first iterations. They pay for context creation, lazy module initialization, cudnn.benchmark autotuning, torch.compile compilation and the caching allocator requesting memory from the driver. Report the median of the rest, which ignores the occasional slow step from a garbage collection or a checkpoint write. torch.utils.benchmark.Timer does all of this: timeit runs max(n/100, 2) warmup calls first, and blocked_autorange picks a block size so each timed block is long enough that timer overhead is negligible, then reports the median per call.

The PyTorch profiler

torch.profiler.profile records every operator call on the CPU and, through NVIDIA's CUPTI via Kineto, every kernel and memory copy on the GPU, with the link between each launch and its kernel. Its main options:

• activities: ProfilerActivity.CPU and ProfilerActivity.CUDA.
• record_shapes (input shapes per operator), profile_memory (allocations and frees), with_stack (Python stack per event, expensive) and with_flops (estimated FLOPs for matmuls and convolutions, by formula from the shapes).
• schedule and on_trace_ready, below.

torch.profiler.record_function("name") adds your own labelled span, so "data", "forward", "backward" and "optimizer" show up as blocks in the trace. The profiler adds its own ProfilerStep#N span around each step.

The step schedule

Recording every step would slow training and fill the disk. torch.profiler.schedule(wait, warmup, active, repeat, skip_first) returns a function from step number to a ProfilerAction: NONE (off), WARMUP (tracing on but results discarded, because starting the tracer has its own overhead), RECORD, and RECORD_AND_SAVE on the last active step, which ends the cycle and calls on_trace_ready. Each call to prof.step() advances the step counter and applies the transition between the old and new action from a table. With repeat=0 the cycle never ends; the function warns when warmup is 0. With no schedule, every step is recorded.

The schedule from this page's example, skip_first=2, wait=1, warmup=1, active=2, repeat=2, ported line for line:

Traces and tables

on_trace_ready receives the profiler at each save. tensorboard_trace_handler(dir) writes one .pt.trace.json per save, named by host, process and time; export_chrome_trace(path) writes one directly. Either file opens in Perfetto (ui.perfetto.dev) or chrome://tracing as a timeline: CPU threads on top, one row per CUDA stream below, with flow arrows from each launch to its kernel. export_stacks writes stacks for a flame graph, and export_memory_timeline plots memory by category.

key_averages() aggregates events by name (optionally by input shape or stack) and its table(sort_by=..., row_limit=...) prints them. Each row has a total time, the span from entry to exit, and a self time, the total minus the time spent in child operators. aten::linear has a large total but almost no self time, because the work is in its child aten::addmm. Sort by self_cpu_time_total to find host overhead, and by self_cuda_time_total (deprecated in favour of self_device_time_total on the event objects) to find the expensive kernels.

Finding the bottleneck

• Data-loader stalls. In the trace, a long enumerate(DataLoader)#..._next__ span on the CPU with an empty GPU row underneath. Without a profiler, time next(iterator) separately from the step. The fixes are on data loading: more num_workers, pin_memory, prefetching, or cheaper decoding.
• Host overhead. Many short kernels with gaps between them, and CPU time per step close to the step time. Fixes: larger batches, fewer and bigger operations, fused kernels, torch.compile or CUDA graphs, which replay a recorded sequence of launches with one call.
• Hidden synchronizations. .item(), .cpu(), print(loss), indexing with a GPU boolean mask, and nonzero all wait for the GPU. In a trace they show as a cudaStreamSynchronize or cudaMemcpyAsync on the CPU that lasts until the queue drains. One per step for logging is fine; one per layer serializes the step.
• Memory spikes. An out-of-memory error at step 1, or only on some batches. See below.
• Compute-bound. A solid GPU row. Now kernel time matters: check the top self-time kernels, whether matmuls use tensor cores (bf16, dimensions that are multiples of 8), and the MFU.

Memory

The caching allocator (memory management notes) keeps freed blocks for reuse, so nvidia-smi shows the reserved pool, not what tensors use. memory_allocated() is the bytes held by live tensors now; max_memory_allocated() reads "allocated_bytes.all.peak" from memory_stats(); memory_reserved() and max_memory_reserved() are the same for the pool. reset_peak_memory_stats() restarts the peaks, so you can measure one step at a time, and memory_summary() prints the whole table.

For the question "what was allocated at the peak, and by which line", turn on torch.cuda.memory._record_memory_history(), which records a stack trace for every allocation and free (up to max_entries, keeping the latest), run a few steps, and write the history with _dump_snapshot("snapshot.pickle"). Dragging the file into pytorch.org/memory_viz shows every tensor as a band over time; the tallest stack of bands is the peak. The leading underscore means the API may change between releases.

Model FLOPs utilization

Throughput in tokens per second is hard to compare across models and machines. MFU (introduced with PaLM, Chowdhery et al., 2022) divides the FLOPs the model needs per second by the hardware's peak. For a decoder-only Transformer with N parameters, n_layers layers, width d and context T, training costs about

The 6N is 2N for the forward pass's multiply-adds plus 4N for the backward pass (gradients for activations and weights), as on pretraining; the second term is the attention scores and weighted sums, which grow with context. Hardware FLOPs utilization (HFU) also counts work done twice, such as the forward pass repeated by activation checkpointing, so it is never lower than MFU. MFU is the honest number: a run that recomputes everything can raise HFU without training any faster. Well-tuned large dense runs report roughly 40 to 60% MFU.

Worked example

Implementation

Timing a step with events, profiling a few steps on a schedule, measuring the peak memory per step, and computing MFU. train_step and loader stand for your own loop.

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


def time_steps(train_step, batches, warmup=2):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    times = []
    for i, batch in enumerate(batches):
        start.record()
        train_step(batch)
        end.record()
        end.synchronize()
        if i >= warmup:
            times.append(start.elapsed_time(end))
    return sorted(times)[len(times) // 2]


def profile_steps(train_step, loader, steps=14):
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        schedule=schedule(skip_first=2, wait=1, warmup=1, active=2, repeat=2),
        on_trace_ready=tensorboard_trace_handler("./traces"),
        record_shapes=True,
        profile_memory=True,
    ) as prof:
        it = iter(loader)
        for _ in range(steps):
            with record_function("data"):
                batch = next(it)
            train_step(batch)
            prof.step()
    print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=15))


def peak_memory_per_step(train_step, batches, snapshot=None):
    if snapshot:
        torch.cuda.memory._record_memory_history(max_entries=100_000)
    peaks = []
    for batch in batches:
        torch.cuda.reset_peak_memory_stats()
        train_step(batch)
        peaks.append(torch.cuda.max_memory_allocated() / 2**20)
    if snapshot:
        torch.cuda.memory._dump_snapshot(snapshot)
        torch.cuda.memory._record_memory_history(enabled=None)
    return peaks


def mfu(params, layers, d_model, context, tokens_per_step, step_s, gpus, peak_flops):
    per_token = 6 * params + 12 * layers * d_model * context
    return tokens_per_step / step_s * per_token / (gpus * peak_flops)

end.synchronize() waits only for that event, not the whole device. For a single function, torch.utils.benchmark.Timer(stmt="f(x)", globals={"f": f, "x": x}).blocked_autorange() does the warmup, synchronization and median for you. Note that peak_memory_per_step reads the peak after each step, which is itself a sync; that is fine for a diagnostic run and is why it is not left on in production.

Tradeoffs

Check yourself

• Give two correct ways to time a GPU operation in PyTorch, and say why CUDA_LAUNCH_BLOCKING=1 is not one of them.
• For schedule(wait=1, warmup=1, active=2, repeat=1), list the action at steps 0 to 5.
• An operator has a large total time and a small self time. Where is the time?
• What would a data-loader stall look like in a trace, and what would you try first?
• Why is peak memory after step 0 lower than after step 1 with Adam?
• Compute the MFU of a 7B model on 64 GPUs at 989 TFLOP/s peak processing 2M tokens per 9 s step, ignoring attention.

Connections

Why GPU work is asynchronous and what makes kernels fast is on GPU execution. Profiling usually leads to data loading (stalls), mixed precision (kernel time), checkpointing and gradient accumulation (memory), or distributed training, where communication shows up as NCCL kernels in the trace. The FLOP counts come from pretraining, and serving-side measurements of time to first token and inter-token latency are on latency and throughput. The training loop is what gets profiled, and fixing seeds and nondeterminism before comparing runs is on reproducibility.

References

• PyTorch v2.15.0: profile, schedule, ProfilerAction, tensorboard_trace_handler, export_chrome_trace, key_averages, record_function, cuda.synchronize, cuda.Event, max_memory_allocated, reset_peak_memory_stats, _record_memory_history, _dump_snapshot, benchmark.Timer.
• PyTorch v2.15.0 notes: asynchronous execution, memory management.
• Chowdhery et al. PaLM: Scaling Language Modeling with Pathways. 2022, appendix B (model FLOPs utilization). arXiv:2204.02311
• Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022 (hardware versus model FLOPs). arXiv:2205.05198
• Kaplan et al. Scaling Laws for Neural Language Models. 2020, section 2.1 (6N FLOPs per token). arXiv:2001.08361