Mixed Precision GPU Training

You’ve built a solid pipeline for fine-tuning LLMs, but when you hit the GPU and watch training crawl — or worse, crash with an out-of-memory error — you’re losing hours and money. The culprit is often full precision (FP32) training, which wastes GPU memory and compute on numbers that don’t need that much detail. In this lesson, you’ll learn how mixed precision training on GPU slashes memory usage and speeds up training by up to 2-3x, without sacrificing model quality.

The problem this lesson solves

Training LLMs in FP32 is simple but expensive. Every weight, gradient, and optimizer state is stored as a 32-bit floating-point number, consuming 4 bytes each. A 7B-parameter model in FP32 alone takes 28 GB—that’s more than most consumer GPUs have, and that’s before you account for activations, gradients, and optimizer states. The result? You either can’t fit the model, or you’re stuck with tiny batch sizes that make training painfully slow.

But here’s the thing: not every number in training needs that much precision. The forward and backward passes—where you compute logits and gradients—can tolerate 16-bit floating point (FP16) or bfloat16 (BF16) without losing significant accuracy. This is the core insight behind mixed precision training on GPU: store and compute most tensors in a lower precision, while keeping critical states (like the optimizer’s master weights) in FP32. The result is up to half the memory usage and often a 2-3x speedup on modern GPUs.

Core concept / mental model

Think of precision as the number of decimal places you keep when writing a number. FP32 gives you about 7 decimal digits of precision; FP16 gives you about 3. For many parts of training, 3 digits are plenty—but not for everything.

Mixed precision training is like using a calculator with a small display for most calculations, but writing down the most important numbers on a full sheet of paper. Here’s how it works:

• Forward and backward passes: Compute in FP16 or BF16 to save memory and speed up matrix multiplications (which GPUs accelerate in half precision).
• Optimizer states: Keep a master copy of weights in FP32 so that updates are accurate and don’t underflow.
• Loss scaling: Multiply the loss by a large factor before backpropagation to prevent gradients from becoming too small and underflowing to zero in FP16.

Pro tip: BFloat16 (BF16) is even better than FP16 for training because it has the same exponent range as FP32, so you don’t need loss scaling. If your GPU supports BF16 (like A100, H100, or RTX 3090+), prefer it.

Here’s a visual in words: imagine a pipeline where data flows as FP16, but every time a gradient is updated, it gets converted to FP32 for the optimizer step, then converted back to FP16 for the next forward pass. That’s the mixed precision dance.

How it works step by step

Let’s break down the mechanics of how a mixed precision training step works:

1. Initialization: Create a model and optimizer. The optimizer holds FP32 master weights (copies of the initial weights).
2. Forward pass: Convert inputs and model weights to FP16 (or BF16) and compute logits.
3. Loss calculation: Compute the loss in FP16.
4. Loss scaling: Multiply the loss by a scale factor (e.g., 1024) to shift gradients into a representable range.
5. Backward pass: Compute gradients in FP16, using the scaled loss.
6. Unscale gradients: Divide gradients by the scale factor, and clip them if needed (to prevent overflows).
7. Optimizer step: Update the FP32 master weights using the unscaled gradients.
8. Synchronization: Copy the updated FP32 weights back to the FP16 model for the next iteration.

Most libraries (like Hugging Face Transformers or PyTorch) automate all of this—you just enable a flag. But understanding the steps helps you debug when things go wrong.

Hands-on walkthrough

Let’s implement mixed precision training from scratch using PyTorch, then see how to use it with Hugging Face Transformers for LLM fine-tuning.

Minimal example in PyTorch

First, here’s a bare-bones training loop with mixed precision, using torch.cuda.amp (automatic mixed precision):

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Create a tiny model and random data
model = nn.Linear(10, 2).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

dataset = TensorDataset(
    torch.randn(64, 10),
    torch.randint(0, 2, (64,))
)
dataloader = DataLoader(dataset, batch_size=16)

# Create a GradScaler for loss scaling
scaler = torch.cuda.amp.GradScaler()

# Training loop
model.train()
for epoch in range(2):
    for inputs, labels in dataloader:
        inputs, labels = inputs.cuda(), labels.cuda()

        optimizer.zero_grad()

        # Enable autocast for forward pass
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)

        # Scale loss and backprop
        scaler.scale(loss).backward()

        # Unscale gradients and step optimizer
        scaler.step(optimizer)
        scaler.update()

        print(f"Epoch {epoch}, loss: {loss.item():.4f}")

When you run this, you’ll see steady loss values. The key lines are autocast() and scaler.scale(loss).backward(). The autocast context automatically casts tensors to FP16 where beneficial, and the scaler handles gradient scaling to prevent underflow.

Mixed precision in Hugging Face Transformers

For LLM fine-tuning, you rarely write the training loop from scratch. Here’s how to enable mixed precision with the Trainer:

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").cuda()

# Prepare a tiny dataset (replace with real data)
train_data = [
    {"text": "The sky is blue because "},
    {"text": "Cats are tiny because "},
]

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, padding=True, max_length=32)

from datasets import Dataset
train_dataset = Dataset.from_dict({"text": [d["text"] for d in train_data]})
train_dataset = train_dataset.map(tokenize, batched=True)

# Enable mixed precision by setting fp16=True or bf16=True
args = TrainingArguments(
    output_dir="./output",
    fp16=True,          # or bf16=True if your GPU supports it
    per_device_train_batch_size=4,
    num_train_epochs=1,
    save_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)

trainer.train()

This will train GPT-2 with FP16 mixed precision. With fp16=True, the Trainer automatically applies autocast and loss scaling. If you have an Ampere+ GPU, switch to bf16=True for even better stability.

Pro tip: Check your GPU compatibility with torch.cuda.get_device_capability(). If it returns (8, 0) or higher, you can use BF16.

Measuring the speedup

Let’s benchmark the difference to see the impact:

import time
import torch
import torch.nn as nn

model = nn.Linear(1024, 1024).cuda()
device = torch.device("cuda")
input_tensor = torch.randn(256, 1024, device=device)

# FP32
model_fp32 = model.double()  # not actually, but simulate; let's use float32
model_fp32 = model.float()
start = time.time()
for _ in range(100):
    out = model_fp32(input_tensor)
    out.sum().backward()  # approximate
print("FP32 time:", time.time() - start)

# Mixed precision
model_mp = model.half()
input_half = input_tensor.half()
start = time.time()
for _ in range(100):
    with torch.cuda.amp.autocast():
        out = model_mp(input_half)
        out.sum().backward()
print("Mixed precision time:", time.time() - start)

In practice, you’ll see a 1.5–2x speedup on matrix-heavy layers.

Compare options / when to choose what

• FP16: Best for older architectures (Pascal, Volta) where BF16 isn’t supported. Requires careful handling of loss scaling.
• BF16: Preferred on Ampere and newer because it avoids underflow, making training more robust.
• FP32: Reserve for debugging or models where precision is critical (e.g., some inference scenarios).

Troubleshooting & edge cases

Loss spikes or NaN

• Cause: Gradient overflow due to loss scaling that’s too high, or underflow when scaling is too low.
• Fix: Try bf16=True instead of fp16=True if your GPU supports it. If you must use FP16, adjust the GradScaler’s growth interval or set the initial scale factor manually.

Model doesn’t converge

• Cause: The learning rate is too aggressive when using mixed precision, or certain layers (like BatchNorm) are sensitive to low precision.
• Fix: Lower the learning rate slightly (common in fine-tuning, anyway). Also, ensure your model is not using layers that require high precision (e.g., some attention implementations).

Out-of-memory errors persist

• Cause: You’re not actually using mixed precision (e.g., forgot to set fp16=True), or you’re still keeping too many FP32 tensors.
• Fix: Double-check the training arguments. Also, consider gradient checkpointing to trade compute for memory.

BF16 not supported

• Error: AssertionError: Mixed precision training with BF16 is not supported
• Cause: GPU architecture is older than Ampere (Compute Capability 8.0+).
• Fix: Switch to fp16=True or upgrade your GPU.

Pro tip: Always monitor nvidia-smi during training. If you see memory usage drop by ~40% after enabling mixed precision, you’re on the right track.

What you learned & what's next

In this lesson, you learned how mixed precision training on GPU reduces memory usage and speeds up training by storing and computing in FP16 or BF16 while keeping FP32 optimizer states. You now know how to implement it from scratch with GradScaler and autocast, and how to enable it in Hugging Face Trainer. You also know when to choose FP16 vs BF16 based on your GPU.

Next, in the track, you’ll learn how to combine mixed precision with parameter-efficient fine-tuning methods like LoRA to train even larger models on a single GPU. That’s a powerful combination that lets you fine-tune 70B models on consumer hardware.