Mixed Precision in PyTorch: What autocast Changes and When It Pays Off

Your GPU has Tensor Cores — circuitry built for half-precision matrix math. If your training loop is plain float32, they sit mostly idle while every weight, activation, and gradient moves through the full-precision pipeline. PyTorch's answer is automatic mixed precision (AMP): a context manager around the forward pass plus, for float16, a gradient scaler. It's one of the highest payoff-per-line changes you can make to a training loop on a modern NVIDIA GPU — but the payoff is conditional, and the float16 path has a failure mode worth understanding before you rely on it.

What autocast actually changes

AMP is per-operation, not per-model. Inside torch.amp.autocast, ops on its cast-to-fp16 list — matrix multiplies, convolutions, linear layers — run in half precision. Ops on its keep-fp32 list — including cross-entropy and softmax — stay full precision for numerical stability. Everything else runs in whatever dtype its inputs arrive in.

Your parameters are not converted. Weights stay float32 in memory and in checkpoints; autocast casts tensors on the fly, op by op. That makes AMP a training-time decision with no deployment lock-in — the model you save is an ordinary float32 model. The speedup comes from hardware: on GPUs with Tensor Cores (Volta and newer), fp16 and bfloat16 matrix math runs at higher throughput than fp32, and half-sized tensors also halve the bytes moved through memory.

The training-loop change, end to end

Run this in your normal training script on a CUDA machine; replace the model, loader, and loss placeholders. Examples assume PyTorch 2.x — newer 2.x releases deprecate the torch.cuda.amp namespace in favor of torch.amp (used here), so check your installed version's docs if the imports differ.

import torch

model = MyModel().cuda()                    # your model, unchanged
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler("cuda")       # torch.cuda.amp.GradScaler on older 2.x

for inputs, targets in train_loader:        # your DataLoader
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)
    optimizer.zero_grad(set_to_none=True)

    # Wrap forward + loss only — not backward, not optimizer.step()
    with torch.amp.autocast("cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)    # e.g. cross_entropy: runs in fp32

    scaler.scale(loss).backward()  # backward runs on the scaled loss
    scaler.step(optimizer)         # unscales, then steps only if grads were finite
    scaler.update()                # adapts the scale factor for the next iteration

The scaler exists because float16 is narrow: its smallest representable values are around 6e-8, and many gradients are smaller, so they'd flush to zero and silently stall learning. GradScaler multiplies the loss so gradients land in a representable range, unscales before the optimizer step, and — the part people miss — skips the step entirely if any gradient came out inf or NaN, then lowers the scale and tries again. Watch scaler.get_scale() across iterations: repeated drops mean persistent overflow or underflow. If you clip gradients, call scaler.unscale_(optimizer) first.

Checking that it actually worked

Confirm the casting lands where you expect:

with torch.amp.autocast("cuda"):
    layer = torch.nn.Linear(64, 64).cuda()
    out = layer(torch.randn(8, 64, device="cuda"))

print(layer.weight.dtype)  # torch.float32 — parameters are untouched
print(out.dtype)           # torch.float16 — the matmul ran in half precision

Then measure. Run a fixed number of steps with and without AMP — the enabled= flag makes toggling trivial — with the same seed and data order, and compare wall-clock time plus peak memory:

torch.cuda.reset_peak_memory_stats()
# ... run N fixed training steps ...
print(torch.cuda.max_memory_allocated() / 1e6, "MB peak")

On Tensor Core GPUs with matmul-heavy models, speedups on the order of 1.5–3x are commonly reported; your numbers depend on the model, batch size, and GPU. Treat any figure — including this one — as a hypothesis until your own benchmark confirms it. Compare final validation metrics between the fp32 and AMP runs too: results differ slightly by design, and parity is verified, not assumed.

float16 or bfloat16?

bfloat16 keeps float16's speed but uses float32's exponent range, so gradients don't underflow and no scaler is needed:

	float16	bfloat16
Dynamic range	Narrow (max ~65,504; tiny gradients underflow)	Same exponent range as float32
Precision per value	10 mantissa bits	8 mantissa bits
GradScaler	Required	Not needed
Fast hardware	Volta (7.0) and newer	Ampere (8.0) and newer; many recent CPUs

print(torch.cuda.get_device_capability())  # e.g. (8, 0) = Ampere

A practical rule: on compute capability 8.0 or higher, start with dtype=torch.bfloat16 and a plain loss.backward(); optimizer.step() — the simplest correct loop. On 7.x (Volta/Turing), use float16 with GradScaler as above. On older GPUs, AMP gains little; skip it. The trade: bfloat16 spends its bits on range rather than precision — 8 mantissa bits versus float16's 10 — so individual values are coarser. It's the more robust default in practice, but validate your final metric either way.

Where the speedup doesn't show up

• Memory-bound elementwise ops, and small models dominated by kernel-launch overhead, see little gain.
• Pre-Volta GPUs lack Tensor Cores, so expect little to no speedup there.
• Distributed training: DDP works with autocast, but gradients are still all-reduced in float32 because your parameters are float32 — on bandwidth-limited clusters, communication, not compute, can become the bottleneck.
• Loss functions and other stability-sensitive ops intentionally stay in fp32, so their cost doesn't shrink.

A checklist for your next run

1. Check torch.cuda.get_device_capability(): bf16 on 8.0+, fp16 plus GradScaler on 7.x, skip AMP on older hardware.
2. Wrap forward and loss in autocast; add the scaler only for float16.
3. Benchmark a fixed step count with enabled=True/False; record wall-clock and max_memory_allocated().
4. Compare final validation metrics against the fp32 run on the same seed and data order.
5. If scaler.get_scale() keeps dropping or you hit NaNs, switch to bfloat16 if your hardware allows, or find the op producing inf/NaN before blaming AMP.

The change is a handful of lines and trivially reversible in code — but its value is empirical. Measure on your own workload before you keep it.