Profiling PyTorch Training: Find the Bottleneck Before Buying More GPUs

Most "we need more GPUs" conversations start with a training job that is only using a third of the hardware it already has. The pattern is consistent: someone reports step time, someone else proposes a bigger cluster, and nobody has looked at a trace. Profiling is the cheapest capacity you will ever buy — an afternoon with torch.profiler routinely recovers 30–60% of wall-clock on a first-pass training loop.

This tutorial is the workflow we run on client jobs: measure honestly, classify the bottleneck into one of four buckets, open a trace, then fix the specific thing the trace shows. It assumes PyTorch 2.x on NVIDIA hardware, but the method transfers.

Step 0: measure step time honestly

CUDA is asynchronous. time.time() around a forward pass measures how long it took Python to queue the work, not to do it. Two ways to time correctly:

import time, torch

def timed_step(fn, warmup=5, iters=20):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        fn()
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end) / iters  # milliseconds

Warmup matters more than people expect: the first iterations pay for cuDNN/cuBLAS autotuning, allocator growth, and — if you use torch.compile — full graph compilation. Reporting a compiled model's first step as its step time is the single most common benchmarking mistake we see.

Fix your headline metric before optimizing anything. For vision, images/second. For LLMs, tokens/second and model FLOPs utilization (MFU). MFU is worth computing once:

MFU = (6 * params * tokens_per_second) / (num_gpus * peak_flops_per_gpu)

A dense transformer training run in BF16 on H100s should land somewhere around 35–55% MFU. If you are at 10%, the problem is not the hardware.

Step 1: classify the bottleneck in five minutes

Before opening a profiler, work out which of four buckets you are in. Run nvidia-smi dmon -s um (or nvitop) alongside the job for 30 seconds.

Symptom	Bucket	Typical cause
GPU utilization low, spiky; CPU cores pinned	Input pipeline	dataloader workers, decoding, augmentation
GPU utilization low, CPU idle too	Host/launch bound	Python overhead, many tiny kernels, sync points
GPU utilization ~100%, MFU still low	Kernel efficiency	unfused ops, bad shapes, wrong precision, memory-bound layers
Utilization sawtooths across ranks	Communication/imbalance	NCCL waits, uneven sequence lengths, stragglers

A fast discriminator for the input pipeline: train on one batch, repeatedly.

batch = next(iter(loader))
batch = [t.cuda(non_blocking=True) for t in batch]
ms = timed_step(lambda: train_step(batch))

If step time collapses when the dataloader is removed, your GPU is starving and no kernel tuning will help. That case is a data-engineering problem, not a modelling one — see our data processing services for what that work looks like at scale.

Step 2: capture a trace with torch.profiler

Do not profile every step. Use the schedule so you skip warmup, capture a handful of representative steps, and write a Chrome trace.

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

prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=5, warmup=3, active=5, repeat=1),
    on_trace_ready=tensorboard_trace_handler("./trace"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
)

with prof:
    for i, batch in enumerate(loader):
        train_step(batch)
        prof.step()
        if i > 15:
            break

Two habits that pay for themselves:

• Annotate your loop phases with torch.profiler.record_function("dataloading"), ("forward"), ("backward"), ("optimizer"). Traces of real models are unreadable without labels.
• Keep with_stack=True only when you need attribution back to Python lines; it inflates trace size significantly on large models.

For a quick text summary without leaving the terminal:

print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=20))

Then open ./trace/*.pt.trace.json in Perfetto (or chrome://tracing). You want the GPU stream rows, not the Python rows, on your first pass.

Step 3: read the trace

Four visual patterns cover the large majority of real findings.

Gaps between kernels on the GPU stream. The GPU is idle waiting for work. Look at the CPU rows underneath the gap: if enumerate(DataLoader) fills it, you are input-bound. If Python operator dispatch fills it, you are launch-bound — the cure is torch.compile (fewer, bigger kernels) or CUDA graphs.

A wall of very short kernels. Hundreds of sub-20-microsecond kernels means launch overhead dominates useful math. This is the classic elementwise-chain signature — a normalization, an activation, a residual add and a dropout each round-tripping to HBM. torch.compile fuses these; our torch.compile in practice walkthrough covers the measurement side.

cudaStreamSynchronize or cudaMemcpyAsync (DtoH) in the middle of the step. Something forced a sync: a .item(), a .cpu(), a print(loss), a Python if loss > x:, or boolean indexing on a GPU tensor. Each one drains the pipeline. Accumulate losses on-device and log every N steps instead:

running_loss += loss.detach()          # stays on GPU
if step % 50 == 0:
    print(running_loss.item() / 50)    # one sync every 50 steps
    running_loss.zero_()

To hunt these down systematically, run once with torch.cuda.set_sync_debug_mode("warn") and read the warnings.

Long nccl:all_gather / all_reduce bars. In FSDP or DDP these are usually waits, not slow networks: one rank arrived late. Suspect uneven batch composition (variable sequence lengths), a rank doing extra logging or checkpointing, or a straggler GPU with thermal throttling. Compare per-rank traces before blaming the interconnect.

Step 4: profile memory separately

Out-of-memory is a different investigation, and the snapshot tool is dramatically better than reading torch.cuda.max_memory_allocated().

torch.cuda.memory._record_memory_history(max_entries=100_000)
try:
    for i, batch in enumerate(loader):
        train_step(batch)
        if i > 3:
            break
finally:
    torch.cuda.memory._dump_snapshot("mem.pickle")
    torch.cuda.memory._record_memory_history(enabled=None)

Upload mem.pickle to pytorch.org/memory_viz — it renders every allocation with the stack that created it. What you are looking for:

• A sawtooth that never returns to baseline — you are retaining tensors across steps. The usual culprit is appending a tensor with a live graph to a list; .detach() or .item() it first.
• A large flat block — optimizer state. AdamW costs 8 bytes per parameter in FP32 moments; consider 8-bit optimizers or sharding.
• Plenty of free memory but still OOM — fragmentation. Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True before reaching for a smaller batch size.
• Activations dominating — enable activation checkpointing on transformer blocks; you trade roughly 30% extra compute for a large memory reduction.

Step 5: fix the input pipeline

If Step 1 said input-bound, work in this order, measuring after each change:

1. num_workers: start at 4–8 per GPU, not os.cpu_count(). Too many workers thrash memory and the page cache.
2. pin_memory=True plus .cuda(non_blocking=True) on the tensors — this is what makes host-to-device copies overlap with compute.
3. persistent_workers=True and prefetch_factor=4 — kills per-epoch worker respawn cost.
4. Stop decoding JPEGs on the CPU one at a time. Pre-resize the dataset, pack it into a sequential format (WebDataset shards, Parquet, FFCV), or decode on the GPU with torchvision.io.decode_jpeg(device="cuda")/DALI.
5. For text, pre-tokenize once and store token IDs. Tokenizing in the training loop is pure waste.
6. Check the filesystem. A network mount with small random reads will defeat every other fix; stage shards to local NVMe first.

Step 6: verify and keep it verified

Re-run Step 0's measurement and record the number. Then stop the regression from coming back:

• Add a smoke benchmark to CI that runs 30 steps on fixed synthetic data and fails if step time regresses more than ~10%.
• Log tokens/second and MFU to your experiment tracker on every run, not just optimization runs.
• Re-profile after any dependency bump. A new PyTorch, driver, or attention-kernel version can change which path your model takes.

The short checklist

Check	Command or setting
TF32 enabled	torch.set_float32_matmul_precision("high")
BF16 autocast on	torch.autocast("cuda", dtype=torch.bfloat16)
No hidden syncs	torch.cuda.set_sync_debug_mode("warn")
Fused optimizer	AdamW(..., fused=True)
Gradients cleared cheaply	zero_grad(set_to_none=True)
Dataloader overlapping	pin_memory, non_blocking=True, persistent_workers
Kernels fused	torch.compile(model) and a trace to prove it
Fragmentation handled	PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Profiling is not a specialist ritual; it is the step that decides whether the next thing you do is worth doing. A trace turns "training is slow" into "38% of the step is idle waiting on JPEG decode", and that sentence is something an engineer can actually fix.

If you would rather have someone else read the traces, that is precisely what our PyTorch performance optimization engagements do: we profile your job, produce a ranked list of fixes with measured impact, and implement the ones you want. Get in touch with your model, hardware and current step time, and we will tell you what we think is recoverable.