Mixed precision and checkpoints

Two habits make real training runs cheaper and safer. Mixed precision does most of the math in a smaller number format, which saves memory and time. Checkpoints save enough state to restart a run where it stopped, rather than from scratch.

Autocast with bfloat16

Here a linear layer runs its matrix multiply in bfloat16, but the weights stay float32:

import torch
import torch.nn as nn

model = nn.Linear(4, 1)
x = torch.randn(8, 4)
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    out = model(x)
print(out.dtype)           # the matmul ran in bfloat16
print(model.weight.dtype)  # the weights stay float32

Output:

torch.bfloat16
torch.float32

On an NVIDIA GPU, the same pattern uses device_type="cuda". The float32 weights are the reason this works: the optimizer updates full-precision values.

bfloat16 has the same exponent range as float32, so tiny gradients rarely underflow. You usually do not need a scaler with it. float16 has a smaller range, so gradients can flush to zero. For float16 add a GradScaler. It multiplies the loss before backward() and divides the gradients before step().

Choosing a dtype

Do not call .half() on the whole model for training. Autocast keeps the sensitive reductions in float32 for you.

What a checkpoint holds

Saving only the weights lets you test a model but not resume training. A resume-friendly checkpoint stores more:

import torch
import torch.nn as nn

model = nn.Linear(4, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
ckpt = {"model": model.state_dict(), "optimizer": opt.state_dict(),
        "step": 1200, "config": {"lr": 1e-3}}
torch.save(ckpt, "ckpt.pt")

restored = torch.load("ckpt.pt", weights_only=True)
model.load_state_dict(restored["model"])
print(restored["step"], restored["config"])

Output:

1200 {'lr': 0.001}

This writes ckpt.pt to the current folder. The full version also stores the scheduler state, the scaler state, the epoch and the best validation score.

Restore the optimizer state when you continue the same run. Skip it when you fine-tune with a new learning rate schedule. Always store the config. Six months later, you will not remember whether weight decay was 0.01 or 0.1.

To load a checkpoint saved on a GPU onto a laptop without one, pass map_location="cpu" to torch.load.

Reproducibility

Seeding makes runs repeatable:

import random
import numpy as np
import torch

def seed_all(s=0):
    random.seed(s)
    np.random.seed(s)
    torch.manual_seed(s)
    torch.cuda.manual_seed_all(s)

Full determinism costs speed, and it is still fragile across CUDA versions. For research, run a few seeds and report the spread. For shipping, prefer robust metrics over bit-identical replays.

Common mistakes

• Converting the whole model to half precision. Use autocast, which keeps the reductions safe.
• Saving weights without the optimizer or config. You cannot resume exactly, and you cannot explain the run later.
• Loading a GPU checkpoint on CPU without map_location. It fails on machines without a GPU.

Check yourself

1. Why does bfloat16 usually skip the GradScaler while float16 needs it? Answer: bfloat16 has the exponent range of float32, so gradients rarely underflow.
2. What must a checkpoint hold to resume a run exactly? Answer: the model weights, the optimizer state, the step count and the config, plus the scheduler and scaler when you use them.

Key takeaways

• Autocast runs matrix math in a lower precision while weights stay in float32.
• Prefer bfloat16 when the hardware supports it. Use float16 with a GradScaler on older GPUs.
• Checkpoints store weights, optimizer state, step and config, so runs can resume and be explained later.
• Seed runs and report variance, rather than chasing identical replays.

Further learning

When one GPU is not enough, the next step is distributed training.