Mixed Precision Training

You’ve just spent 12 hours training a model, only to hit out-of-memory errors that kill your session—or worse, your GPU is fast but your training loop crawls because it’s stuck in single precision. Every Applied AI engineer hits this wall: your model needs to be bigger, but your GPU runs out of memory; your training time balloons, and your experiments feel like an eternity. The fix isn’t a fancier GPU or shrinking your model—it’s mixed precision training, a technique that can cut memory usage by up to half and accelerate training by 2–3x on modern hardware, all with minimal code changes and little to no accuracy loss. In this lesson, you’ll learn exactly how to use mixed precision training in PyTorch, from the core concept to a hands-on exercise that you can run yourself.

The problem this lesson solves

Training deep learning models is a resource-hungry endeavor. By default, PyTorch—and most deep learning frameworks—store all parameters, activations, and gradients in FP32 (32-bit floating point). That’s 4 bytes per value. A modest model with 100 million parameters consumes ~400 MB just for the parameters, and activations and gradients easily multiply that several times over. On a typical GPU with 8–16 GB of VRAM, you’ll quickly hit memory limits, preventing you from increasing batch size, model size, or input resolution—all of which often boost accuracy.

Beyond memory, FP32 arithmetic is slower than lower-precision operations on modern GPUs. NVIDIA’s Tensor Cores, available on V100, T4, A100, and RTX 20+ series GPUs, are designed to perform FP16 (16-bit floating point) matrix multiplications dramatically faster than FP32. But naively casting everything to FP16 causes your model to diverge—gradients underflow to zero, and training becomes unstable. That’s the core problem: you want speed and memory savings, but you can’t sacrifice numerical stability.

The solution is mixed precision training: use FP16 where it’s safe (matrix multiplications, convolutions) and keep FP32 where it matters (master weights, loss scaling, gradient accumulation). This hybrid approach—first popularized by Baidu and NVIDIA in 2018—lets you get the best of both worlds.

Core concept / mental model

Think of mixed precision training like a chef who uses a chef’s knife (FP32) for precision tasks like slicing delicate herbs, but switches to a food processor (FP16) for heavy chopping that doesn’t need precision. FP16 is faster and uses half the memory, but it has a limited dynamic range: it can represent numbers between about 1e-5 and 65,504, and it suffers from rounding errors when values get too small. FP32 covers a vastly larger range (down to 1e-38) with more precision, making it safe for storing weights and accumulating gradients.

Here’s the mental model: FP16 is the sports car—fast but fragile; FP32 is the work truck—slow but reliable. Mixed precision uses the sports car for the heavy lifting (matrix math) while the work truck handles the critical payload (weights and gradients). To make the sports car safe, you add a seatbelt: loss scaling. Since gradients in FP16 can underflow (turn to zero), you multiply the loss by a scalar factor (e.g., 512) before backpropagation, so gradients stay within FP16’s range. After computing gradients, you divide them back down before updating the master FP32 weights.

Key terms you’ll see throughout: - FP32: Single-precision floating point (32-bit). Default in most frameworks. - FP16: Half-precision floating point (16-bit). Half the memory, faster on Tensor Cores, but narrower range. - Master weights: An FP32 copy of your model’s weights used for updates; the model itself runs in FP16. - Loss scaling: Multiplying the loss by a factor to prevent gradient underflow. - Automatic Mixed Precision (AMP): A PyTorch API that automates the casting and scaling for you.

How it works step by step

Let’s break down what happens inside PyTorch’s torch.cuda.amp (AMP) when you enable mixed precision training. The magic is in two components: GradScaler and autocast.

Step 1: Run forward pass in autocast

torch.cuda.amp.autocast is a context manager that tells PyTorch to use FP16 for certain operations (like torch.matmul, torch.nn.Conv2d, torch.nn.Linear) and FP32 for others (like torch.nn.BatchNorm, which requires FP32 for stability). This happens automatically—you don’t manually cast tensors. The context manager ensures that ops within its scope are cast appropriately.

Step 2: Compute loss, then scale it

After the forward pass, you still compute the loss in FP32 (PyTorch does this automatically). But before calling loss.backward(), you scale the loss: scaler.scale(loss).backward(). The scaler multiplies the loss by a large factor (default 2^16 or 65536) so that gradients computed in the backward pass are larger and won’t underflow to zero in FP16.

Step 3: Scale optimizer’s gradients and unscale

Before the optimizer updates the weights, you must unscale the gradients: scaler.unscale_(optimizer). This divides the gradients by the scaling factor, returning them to their true magnitude. Then the optimizer applies the update to the FP32 master weights.

Step 4: Update the scaler

After each step, you call scaler.step(optimizer) (which, if gradients are finite, performs the update) and then scaler.update(). The scaler may adjust the scaling factor dynamically: if gradients overflow (becoming inf or NaN), it reduces the scale; if they stay under a threshold for many steps, it increases it. This adaptive behavior keeps gradients in a safe range.

Why it works

Tensor Cores are specifically designed for mixed precision matrix math, giving you 2x–3x throughput on supported GPUs. By keeping weights in FP32, you avoid the accumulation of rounding errors that would occur if you updated FP16 weights directly. The result: faster training, lower memory footprint, and comparable model accuracy.

Hands-on walkthrough

Now let’s put this into practice. We’ll set up a simple CNN on CIFAR-10 and train it with mixed precision. This example is complete and runnable—make sure you have PyTorch with CUDA support installed.

Setup

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.cuda.amp import autocast, GradScaler

# Check GPU availability
assert torch.cuda.is_available(), "CUDA GPU required for mixed precision"
print(f"Using GPU: {torch.cuda.get_device_name(0)}")

Define a simple model

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc = nn.Linear(32 * 8 * 8, 10)

    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        return self.fc(x)

Training loop with mixed precision

def train_with_mixed_precision(model, trainloader, epochs=5):
    model = model.cuda()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    criterion = nn.CrossEntropyLoss()
    scaler = GradScaler()  # Automatically handles loss scaling

    for epoch in range(epochs):
        running_loss = 0.0
        for images, labels in trainloader:
            images, labels = images.cuda(), labels.cuda()

            optimizer.zero_grad()

            # Forward pass with autocast
            with autocast():
                outputs = model(images)
                loss = criterion(outputs, labels)

            # Backward pass with scaling
            scaler.scale(loss).backward()

            # Unscale gradients, clip if needed, then step
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()

        print(f"Epoch {epoch+1}, loss: {running_loss/len(trainloader):.4f}")

Expected output (loss will vary):

Epoch 1, loss: 1.8421
Epoch 2, loss: 1.2335
Epoch 3, loss: 0.9234
...

This code is the entire essence of mixed precision training. Notice we only added scaler = GradScaler(), with autocast():, scaler.scale(loss).backward(), and scaler.step(optimizer) / scaler.update(). The model itself doesn't change.

Measuring speed and memory benefits

To really see the impact, run the same training loop without mixed precision (just remove the scaler and autocast) and compare GPU memory and time. Use torch.cuda.max_memory_allocated() to see peak memory:

# After training
print(f"Peak GPU memory allocated: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")

You’ll typically see a 30–50% reduction in memory and 1.5–2x speedup on Tensor Core GPUs.

Compare options / when to choose what

When you decide to use mixed precision training, you have a few API choices. Here’s a comparison:

Recommendation: If you’re writing raw PyTorch, use torch.cuda.amp (or torch.amp for newer PyTorch). If you’re using Lightning or Hugging Face, just set the precision parameter. The manual model.half() approach is almost always a mistake because it often breaks batch normalization and causes gradient underflow—avoid it.

Troubleshooting & edge cases

Despite the simplicity, you’ll likely hit a few common issues. Here’s how to fix them.

Loss or gradients become NaN or inf

This is the classic sign that the scale factor is too high or too low. PyTorch’s GradScaler handles this automatically by skipping optimizer steps when gradients are non-finite and reducing the scale. If you see NaNs constantly, check: - Your learning rate might be too high (lower it). - Your loss function might be numerically unstable (use log-softmax or label smoothing). - If you’re not using scaler.update() after each step, the scaler can’t adapt.

BatchNorm layers misbehave

BatchNorm expects FP32, and autocast automatically keeps it in FP32. But if you manually cast the model to .half(), BatchNorm will fail—don’t do that.

No speedup observed

Mixed precision only accelerates on GPUs with Tensor Cores (NVIDIA V100, T4, A100, RTX 20+). On older GPUs, FP16 may be slower. Also, if your model has many small layers, the overhead of casting can negate benefits. Profile your GPU utilization to ensure you’re actually hitting Tensor Core paths.

CPU training

torch.cuda.amp is CUDA-only. For CPU, you can use torch.amp.autocast(device_type='cpu') (available in PyTorch 2.x), but the speedup is minimal. Don’t expect miracles.

GradScaler with gradient clipping

If you use gradient clipping, you must unscale first:

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)

If you skip unscale_, your clipping will use scaled gradients, giving wrong results.

What you learned & what's next

You now know how to use mixed precision training to cut memory usage and accelerate training without sacrificing accuracy. You’ve learned the core concepts (FP16 vs FP32, master weights, loss scaling), the step-by-step implementation in PyTorch, and how to troubleshoot common issues. You can apply this to your own models—whether you’re training a vision classifier or fine-tuning a transformer.

Next lesson in the Applied AI engineering track will likely cover gradient accumulation or distributed training—both of which build on this speed and memory optimization mindset. You’ve already mastered the foundation: don’t waste GPU resources when mixed precision is one line of code away.

Pro tip: Always benchmark your system. Implement mixed precision, but measure actual wall-clock time and memory before committing to it in production. The theoretical gains don’t always translate, but with the right hardware, you’ll see dramatic improvements.

Now, go ahead and try it: modify your existing training script with these three lines, and see the difference on your next run.