PyTorch Mixed Precision, torch.compile and Performance Tuning

Automatic Mixed Precision

Performance work matters when models become large, datasets grow, or iteration speed slows down. PyTorch gives developers tools such as automatic mixed precision, torch.compile, DataLoader workers, pinned memory, and profilers.

Optimize only after the model is correct. A fast broken training loop is still broken. First verify shapes, loss, gradients, and validation metrics. Then tune data loading, GPU utilization, precision, and compilation.

Mixed precision uses lower precision where safe to speed computation and reduce memory. On CUDA, autocast and GradScaler are common for training. Inference often uses autocast or model-specific precision choices.

Use autocast around the forward pass and loss calculation.
Use GradScaler to reduce underflow risk during backward.
Keep validation numerically checked after enabling AMP.

Throughput Bottlenecks

If the GPU waits for data, tune the DataLoader before changing the model. If memory is full, reduce batch size, use AMP, checkpoint activations, or simplify the architecture.

Increase num_workers carefully and measure.
Use pin_memory when transferring CPU batches to CUDA.
Avoid expensive Python work inside __getitem__ when possible.
Use torch.profiler for evidence instead of guessing.

AMP Training Step

This is the standard CUDA mixed precision shape.

AMP Training Step

scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())

for images, labels in train_loader:
    images = images.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)

    with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
        logits = model(images)
        loss = loss_fn(logits, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

AMP is most useful on GPUs with strong lower-precision support.
Compare metrics before and after enabling AMP.

torch.compile Baseline

Compilation can speed up some models after the first warmup iterations.

torch.compile Baseline

model = model.to(device)

if hasattr(torch, "compile"):
    model = torch.compile(model)

# Run a few warmup batches before measuring throughput.
for step, (images, labels) in enumerate(train_loader):
    if step == 20:
        break
    train_step(model, images, labels)

Compilation benefits vary by model and hardware.
Measure end-to-end time, not only one operation.

Before you move on