Mixed precision and GPU performance

autocast and GradScaler, channels_last, pinned memory, gradient accumulation, and profiling to find where the frame time actually goes.

Automatic mixed precision

import torch
from torch.amp import autocast, GradScaler

model = model.cuda()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler("cuda")

for x, y in loader:
    x = x.cuda(non_blocking=True)
    y = y.cuda(non_blocking=True)

    opt.zero_grad(set_to_none=True)
    with autocast("cuda", dtype=torch.float16):
        logits = model(x)
        loss = criterion(logits, y)

    scaler.scale(loss).backward()
    scaler.unscale_(opt)                                  # before clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(opt)                                      # skips on inf/nan
    scaler.update()

print(scaler.get_scale())     # the current loss scale, useful when debugging

autocast chooses the dtype per operation: matmuls and convolutions run in float16, reductions and the loss stay in float32. You do not have to cast the model.
GradScaler multiplies the loss to keep small gradients representable and skips a step when an infinity appears, halving the scale and retrying. Without it, training frequently produces NaN in float16.
scaler.unscale_(opt) must run before gradient clipping, or you clip the scaled gradients and the threshold becomes meaningless.
On CPU, mixed precision support is limited; torch.amp.autocast("cpu", dtype=torch.bfloat16) works on recent hardware but is a different story from CUDA.

Memory and throughput tricks

# channels_last: better tensor-core utilisation for convolutions
model = model.to(memory_format=torch.channels_last)
x = x.to(memory_format=torch.channels_last)

# pinned memory and non-blocking copies overlap transfer with compute
loader = torch.utils.data.DataLoader(
    dataset, batch_size=64, num_workers=8,
    pin_memory=True, persistent_workers=True, drop_last=True)
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)

# gradient accumulation: a larger effective batch without more memory
ACCUM = 4
opt.zero_grad(set_to_none=True)
for i, (x, y) in enumerate(loader):
    with autocast("cuda"):
        loss = criterion(model(x.cuda()), y.cuda()) / ACCUM
    scaler.scale(loss).backward()
    if (i + 1) % ACCUM == 0:
        scaler.step(opt)
        scaler.update()
        opt.zero_grad(set_to_none=True)

# a quick memory report
print(torch.cuda.memory_allocated() / 1e9, "GB allocated")
print(torch.cuda.max_memory_allocated() / 1e9, "GB peak")

Technique	Saves	Cost
Mixed precision	Memory and time (1.5-2.5x)	Occasional instability without a scaler
`channels_last`	Convolution time on tensor cores	Nothing if the model is not convolutional
Gradient accumulation	Memory at a larger effective batch	Longer wall clock, worse batch-norm statistics
`set_to_none=True`	Memory and a little time	None; it is strictly better than zeroing
Freezing the backbone	Memory and backward time	No gradients for the frozen part
`torch.compile`	Kernel fusion and launch overhead	A long first compile, some op coverage gaps

⚠️

Gradient accumulation does not reproduce a large batch for batch normalisation. BN statistics are still computed per micro-batch, so a model trained with accumulation behaves differently from one trained with the true batch size. Switch to group or layer norm when this matters.

Profiling the step

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("logs/prof"),
    record_shapes=True,
    profile_memory=True,
) as prof:
    for step, (x, y) in enumerate(loader):
        if step >= 6:
            break
        with record_function("train_step"):
            train_step(x, y)
        prof.step()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

Always warm up before profiling. The first steps include cuDNN autotuning and kernel compilation, which dominate the trace and mislead you.
Read the table by cuda_time_total first, then by cpu_time_total. A large gap between them means the GPU is idle waiting for the CPU.
Use record_shapes=True to catch an unexpectedly large tensor — a batch dimension in the wrong place is visible immediately.
Set torch.backends.cudnn.benchmark = True when your input shapes are fixed; it autotunes kernels once and can give a double-digit speedup, but it thrashes if shapes vary.

FAQ

Why is my GPU utilisation so low?

In order: input pipeline or CPU preprocessing, host-to-device transfer without pinned memory, small kernels with launch overhead, or synchronisation from a .item() inside the loop. Profile before changing the model.

Does mixed precision train well without a scaler?

On some models, but you are relying on the gradients never underflowing. Use GradScaler: it costs almost nothing and turns an occasional silent NaN into an automatic skipped step.

Data pipelines and augmentation with torchvision Debugging PyTorch models

Last refreshed 2026-09-18.

Mixed precision and GPU performance

Automatic mixed precision

Memory and throughput tricks

Profiling the step

FAQ

Related