Speeding Up PyTorch Training with Automatic Mixed Precision

Problem: Training deep models consumes too much GPU memory and time

When you train a large convolutional network or transformer on a single GPU, you often hit two walls: the forward and backward passes take longer than you’d like, and the activations quickly exhaust the device’s memory, forcing you to reduce batch size or gradient accumulation steps. This slows experimentation and increases cost.

Thesis: Using torch.autocast with GradScaler (or bfloat16 without scaling) lets you keep most of the numerical fidelity of float32 while running the heavy math in a lower‑precision format, giving measurable speed‑ups and memory savings on Tensor‑Core GPUs.

How AMP works in PyTorch

Automatic mixed precision (AMP) selects a precision per operation: matrix multiplications and convolutions run in float16 (or bfloat16 on Ampere+), while ops that are sensitive to small‑scale values—such as softmax, loss reductions, and batch‑norm—stay in float32. The torch.autocast context manager handles the casting automatically. When using float16, a GradScaler is required to prevent gradient underflow: it multiplies the loss by a dynamic scale, unscales the gradients before the optimizer step, and backs off the scale if any inf or NaN appears.

If your hardware supports bfloat16 (Ampere GPUs or CPU), you can omit GradScaler because bfloat16 retains the exponent range of float32, making gradient underflow rare.

Worked‑example training loop

The following snippet shows the canonical pattern for float16 AMP on a CUDA device. Replace model, optimizer, loss_fn, and data_loader with your own objects. The code is presented as a template; you should verify it works with your specific PyTorch version.

import torch
from torch.cuda.amp import autocast, GradScaler

model = model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler()

for batch in data_loader:
    inputs, targets = batch
    inputs = inputs.to('cuda', non_blocking=True)
    targets = targets.to('cuda', non_blocking=True)

    optimizer.zero_grad()
    with autocast(device_type='cuda'):          # <-- AMP region
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
    
    scaler.scale(loss).backward()               # scale loss, backward
    scaler.unscale_(optimizer)                  # needed if you clip gradients
    # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # optional
    scaler.step(optimizer)                      # optimizer step with unscaled grads
    scaler.update()                             # adjust scale for next iteration

If you are on an Ampere or newer GPU and prefer bfloat16, the loop simplifies:

for batch in data_loader:
    inputs, targets = batch
    inputs = inputs.to('cuda', non_blocking=True)
    targets = targets.to('cuda', non_blocking=True)

    optimizer.zero_grad()
    with autocast(device_type='cuda', dtype=torch.bfloat16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()

Trade‑offs and limitations

• Numerical stability: float16 AMP can still produce inf/NaN in unstable losses (e.g., very deep unnormalized nets). Monitoring scaler.get_scale() helps detect repeated scale collapses.
• Hardware dependence: Pre‑Volta GPUs lack Tensor Cores, so the throughput gain is minimal; bfloat16 on older CPUs may still help memory but not compute.
• Interaction with other features: When using gradient accumulation, DDP, or torch.compile, the placement of autocast and GradScaler calls must respect the accumulation boundaries and the distributed wrapper.
• Checkpointing: Model parameters stay in float32 master copies by default, but if you manually cast the model to float16 you must keep a separate float32 copy for resuming training.

Practical way to verify the benefit

Run two short experiments on the same dataset:

1. Baseline: train for N epochs with default float32 (no autocast). Record wall‑clock time per epoch and peak memory via torch.cuda.max_memory_allocated().
2. AMP run: repeat with the float16 AMP loop above (or bfloat16 version). Record the same metrics.

Compare the epoch times and memory usage. A typical speed‑up of 1.5‑3× and roughly half the activation memory indicates AMP is effective for your model. Additionally, compare validation loss after the same number of epochs to ensure accuracy has not degraded beyond an acceptable threshold.

Actionable closing

Start by wrapping your forward‑pass and loss computation in torch.autocast(device_type='cuda'). If you are on a Volta‑or‑later GPU and observe stable training, add a GradScaler for float16 or switch to dtype=torch.bfloat16 to drop the scaler altogether. Measure the impact on epoch time and memory, then decide whether to keep AMP for full‑scale training. This incremental approach lets you reap the performance benefits while guarding against numerical surprises.