Practical Guide to PyTorch Automatic Mixed Precision: When, How, and What to Watch

Problem & Takeaway

Deep learning models routinely hit memory limits or slow down on GPUs when trained in full single precision (float32). PyTorch’s Automatic Mixed Precision (AMP) can cut memory and speed up training by running eligible operations in half precision (float16 or bfloat16) while keeping a master copy of weights in float32. The key takeaway is that AMP works with minimal code change, but only if you follow the correct pattern, manage the GradScaler state, and monitor for numerical failures.

Requirements & Hardware

• PyTorch 1.10+ (AMP API stabilized in 1.10; older versions use torch.cuda.amp).
• CUDA-capable GPU with Tensor Cores (Ampere or newer for float16, Volta or newer for bfloat16).
• Optional: CPU support for bfloat16 autocast (PyTorch 2.0+).
• CUDA toolkit and cuDNN matching the PyTorch build.

Before enabling AMP, confirm that your model contains operations that benefit from half precision (matmul, conv, batchnorm). If the model is small or CPU-bound, the speedup may be negligible.

Minimal Correct Integration

The smallest, safest pattern is a forward pass wrapped in torch.autocast and a GradScaler that scales the loss before backpropagation. Below is a ready‑to‑copy snippet.

# Assume model, optimizer, and dataloader are defined
scaler = torch.cuda.amp.GradScaler()
for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()

        # Forward pass in autocast context
        with torch.autocast(device_type='cuda', dtype=torch.float16):
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)

        # Scale loss, backward, step, and update scaler
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    # Optional: log scaler.scale() and scaler.get_scale() each epoch

Key points:

• Use torch.autocast(device_type='cuda', dtype=…) – the dtype can be torch.float16 or torch.bfloat16 depending on hardware.
• The GradScaler must be instantiated once and reused across epochs.
• Always call scaler.update() after scaler.step() to adjust the scale for the next iteration.

Trust Boundaries Around Loss Scaling

Loss scaling is stateful and must be treated as part of the training state:

• Persist scaler.state_dict() together with the optimizer and model checkpoints.
• When resuming from a checkpoint, load the scaler state before the optimizer step.
• Do not suppress overflow warnings silently; let the scaler record a skipped step and log it.

This ensures that a training interruption does not silently reset the effective learning rate to zero.

Operational Checks & Logging

During training, monitor the following metrics to detect early signs of numerical instability:

• Scaler scale value – log scaler.get_scale() each epoch.
• Skipped steps – scaler._skipped_steps (private API, use scaler.get_scale() history).
• GPU memory usage – torch.cuda.max_memory_allocated() per iteration.
• Step time – wall‑clock time per batch.

Example logging snippet:

import time

start = time.time()
# forward/backward as above
end = time.time()
print(f"Step time: {end - start:.4f}s, Scale: {scaler.get_scale():.0f}, Skipped: {scaler._skipped_steps}")

Use these logs to verify that:

• Skipped steps are rare (e.g., 1–2 per epoch for a typical model).
• Scale values stay within a reasonable range (1e2 – 1e6).
• Memory usage drops by ~30–50% compared to full float32.
• Throughput improves by 10–30% on Tensor Core GPUs.

Failure Modes & When to Pivot

AMP is not a silver bullet. The following conditions should prompt you to disable or adjust AMP:

• Persistent Skips: If scaler.update() keeps halving the scale every step, the effective learning rate is zero. Check the loss magnitude and consider increasing the initial scale or using torch.bfloat16.
• NaN/Inf Loss: Occurs when reductions (e.g., mean over a large batch) produce values outside half precision range. Wrap those ops in torch.autocast(enabled=False) to force float32.
• Softmax over Wide Ranges: Softmax can overflow in float16. Keep the logits in float32 or use the logits_to_softmax=True flag in torch.autocast (PyTorch 2.0+).
• Distributed Data Parallel (DDP): Gradient bucketing may run in half precision. Ensure you apply scaler.unscale_(optimizer) before clipping gradients, otherwise clipping thresholds will be wrong.
• Hardware without Tensor Cores: AMP may not yield speedups and can even slow down due to conversion overhead.

When any of the above occurs, either disable AMP globally or selectively disable it for the problematic ops:

with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=False):
    loss = loss_fn(outputs, targets)  # keep in float32

Concrete Example: Comparing Float32 vs AMP

Below is a minimal script that trains a ResNet‑18 on CIFAR‑10 with and without AMP, logging memory and time. Replace the dataset and model with your own as needed.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# Data
transform = transforms.Compose([transforms.ToTensor()])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)

# Model
model = torchvision.models.resnet18(num_classes=10).cuda()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# AMP setup
scaler = torch.cuda.amp.GradScaler()

for epoch in range(1):
    for inputs, targets in trainloader:
        inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        with torch.autocast(device_type='cuda', dtype=torch.float16):
            outputs = model(inputs)
            loss = criterion(outputs, targets)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    print(f"Epoch {epoch+1} complete")

After training, run torch.cuda.max_memory_allocated() and compare to a float32 run. A typical reduction is ~30% memory usage with a 10–20% speedup on an RTX 3090.

Conclusion

AMP is a low‑overhead way to squeeze more performance from modern GPUs. By wrapping the forward pass in torch.autocast and managing a GradScaler, you can reduce memory and improve throughput with minimal code changes. The critical parts are: persisting the scaler state, monitoring scale and skipped steps, and guarding numerically unsafe ops. If you observe persistent overflows or accuracy regressions, disable or selectively turn off AMP for the offending computations. With these checks in place, you can confidently adopt mixed precision for large‑scale training workloads.