PyTorch GPU Utilization Low? Profile the Training Loop

The PyTorch profiler records operations and timing so you can see where a training step spends its time. Inspect a short trace after warm-up, then test one explanation for the idle gaps before changing batch size or distributing the job.

Read one profiler trace before changing the loop

Start with a short capture after warm-up. Look for CPU work, host-to-device transfers, and device kernels on the same timeline. A summary table shows expensive operations but can hide idle gaps between them. Keep input shapes and the number of measured steps fixed when you compare traces. The example below assumes you supply the model and training functions.

Define the baseline before making it faster

Save the code revision, dataset version, hardware, software stack, random seeds, batch size, precision, and evaluation metric. Measure examples or tokens per second after warm-up, peak device memory, data-loader wait, checkpoint time, and the time required to reach a target validation score. Samples per second alone can reward a change that trains faster but converges worse.

Run a short correctness test first. Confirm that loss decreases, gradients are finite, checkpoints restore, and evaluation matches the reference within an agreed tolerance. Optimization on top of a silent data or numerical bug only produces wrong results faster.

Profile the step as a timeline

Use a framework profiler to capture CPU work, accelerator kernels, memory, shapes, and call stacks for a small number of representative steps. Include warm-up steps so one-time compilation and allocation do not dominate the trace.

import torch

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    record_shapes=True,
    profile_memory=True,
) as profiler:
    for step, batch in enumerate(loader):
        train_step(batch)
        profiler.step()
        if step == 4:
            break

print(profiler.key_averages().table(sort_by="self_cuda_time_total"))

Look for gaps where the accelerator waits, repeated small kernels, expensive copies, synchronization points, and operators that allocate large temporary tensors. Profile representative sequence lengths and input shapes; a convenient small batch may hide the real bottleneck.

Fix the input pipeline when the accelerator waits

If compute repeatedly pauses for the next batch, optimize the path from storage to device. Cache deterministic preprocessing where appropriate, reduce tiny file reads, move expensive transforms out of the hot loop, and test data-loader worker counts. More workers are not always faster: they consume memory and can compete for CPU, storage, or network bandwidth.

Pinned host memory and asynchronous transfers can overlap copies with compute when the hardware and code path support it. Verify overlap in the trace. Pre-fetching too far ahead can raise memory use without improving throughput.

Track corrupted samples and preprocessing failures explicitly. A fast pipeline that silently skips hard examples changes the training data.

Use mixed precision with numerical checks

Automatic mixed precision can reduce memory use and accelerate supported operations. Lower precision also narrows numerical range, so use the framework's autocast and gradient-scaling mechanisms instead of converting the entire model blindly.

# Assumes model, optimizer, loader, and loss_fn are defined.
model = model.to("cuda")
scaler = torch.amp.GradScaler("cuda")

for inputs, targets in loader:
    inputs, targets = inputs.to("cuda"), targets.to("cuda")
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Monitor non-finite loss, skipped steps, gradient norms, and validation quality. BF16 often offers a wider exponent range than FP16 on supporting hardware, but the best choice is hardware- and model-dependent.

Increase useful work per step carefully

Larger batches can improve accelerator utilization, but they change optimization dynamics and require more memory. If the target batch does not fit, gradient accumulation reproduces the number of examples per optimizer update but not every behavior of a true large batch. Batch-normalization statistics, dropout sequences, and learning-rate schedules can differ.

Measure tokens or valid elements, not padded examples, for variable-length data. Bucket similar lengths to reduce padding. For sequence models, a smaller batch of long examples can cost more than a larger batch of short ones.

Reduce memory before distributing the job

Activation checkpointing trades extra forward computation for lower activation memory. Optimizer-state sharding, parameter sharding, and offloading trade communication or host transfers for capacity. Apply these when memory limits the batch or model, then re-profile because a capacity fix can become a throughput bottleneck.

Remove accidental tensor retention from logging and metrics. Storing a tensor with its computation graph in a Python list can retain the entire graph. Detach values needed only for reporting:

loss_history.append(float(loss.detach()))

Scale out only when communication pays

Data parallelism copies the model and synchronizes gradients. It works well when each worker performs enough compute between synchronizations. Small models, slow networks, imbalanced batches, or frequent barriers can produce poor scaling.

Report scaling efficiency, not just total throughput. If one device processes 1,000 samples per second and four devices process 3,000, the speedup is 3x and efficiency is 75%. Include evaluation, checkpointing, and data access in end-to-end timing. A training loop that scales while shared storage collapses is not a production win.

Compile stable regions, then check graph breaks

Compiler features can fuse operations and reduce framework overhead, especially when shapes and control flow are stable. Compilation has an upfront cost and may recompile for changing shapes. Measure a full representative run, inspect graph breaks, and keep an eager-mode fallback for debugging.

Do not combine compilation, mixed precision, new kernels, and a batch-size change in one experiment. If quality or performance moves, you will not know why.

Make checkpoint and recovery time part of performance

A job that trains quickly but loses hours after an interruption is not efficient. Measure checkpoint write time, restore time, storage bandwidth, and how much work can be lost between checkpoints. Save the model, optimizer, scheduler, scaler, sampler position, and random state required for a faithful resume. Test recovery before a long run.

Checkpoint frequency is a cost decision. Frequent checkpoints reduce lost work but consume I/O and storage; infrequent checkpoints make failures expensive. Distributed jobs also need a coordinated strategy so one worker does not serialize a huge state while every other accelerator waits. Where supported, asynchronous or sharded checkpoints can reduce the pause, but restored quality and step numbering still need verification.

A safe optimization loop

Lock a reproducible quality and performance baseline.
Profile after warm-up with production-like shapes.
Name the limiting resource and form one hypothesis.
Change one major variable.
Compare throughput, peak memory, convergence, final metric, and cost.
Repeat the result across several runs.
Keep the change only if the end-to-end tradeoff improves.

Frequently asked questions

What is the first optimization to try?

Profile first. Mixed precision is often valuable on modern accelerators, but it will not solve a data-loader stall or a synchronization bottleneck.

Does a larger batch always train faster?

It can improve hardware utilization, but it may change convergence, require learning-rate adjustments, and increase memory. Compare time to a target quality, not only step time.

When should I use more GPUs?

When one device is well utilized, the workload is large enough to amortize synchronization, and the expected wall-clock or capacity benefit justifies added cost and complexity.

Primary sources

Compare a synthetic batch with the real loader

As a diagnostic experiment, repeat a representative batch that is already on the device, then compare that trace with the real input pipeline. If the gaps disappear, investigate loading and transfer work before changing the model. The repeated batch is useful for timing, not for measuring learning quality. Keep it out of any claim about convergence.

The profiler example assumes a CUDA device and your own loader and train_step. For a CPU-only profile, omit the CUDA activity and use CPU timing fields. These fragments are integration examples, not standalone scripts.