GPUs and mixed precision

Why the GPU wins

A CPU has a handful of cores optimised for latency, which is the delay before one task finishes. A GPU has thousands of lanes optimised for throughput, which is the number of tasks finished per second. A matrix multiply is the same small operation repeated over independent elements. That shape is why the GPU is the default for training. It is also why a model built mostly from small elementwise operations does not benefit as much as its FLOP count suggests. A FLOP is one floating-point operation.

Arithmetic intensity is the number of floating-point operations per byte moved. A large matrix multiply loads each value once and reuses it many times, so it is compute-bound. An elementwise operation reads and writes memory for every operation, so it is bandwidth-bound. A kernel is the compiled function that runs on the GPU. The ratio decides which limit a kernel hits, and therefore what to optimise.

Kernel fusion and batch size raise reuse. Fusing two elementwise operations avoids a round trip to memory. A larger batch reuses the same weights across more rows. Launch overhead matters for the same reason. A kernel that runs for a few microseconds can spend more time being scheduled than doing arithmetic.

A float is range and resolution

A floating-point number has a sign, an exponent and a mantissa. The exponent sets the range, and the mantissa, the significant digits, sets the resolution. fp32 spends 8 bits on the exponent and 23 on the mantissa. fp16 spends 5 and 10, bf16 spends 8 and 7, and tf32 spends 8 and 10.

fp16 has fine resolution but a small range. Its largest finite value is 65,504. A gradient smaller than about 6 × 10⁻⁸ becomes zero. bf16 has the range of fp32 and less resolution. That is usually the better trade for deep learning, because the values that matter span many orders of magnitude.

Rounding is not free. Summing many small values in a low-precision accumulator loses the small ones once the running total is large. An accumulator is the running total. The standard pattern is therefore to multiply in low precision and accumulate in fp32. A dot product that accumulates in fp16 can be visibly wrong after a few thousand terms. The error is largest exactly where the sum is dominated by a few big values.

Mixed precision in a training loop

import torch
from torch.amp import GradScaler, autocast

scaler = GradScaler("cuda")
for inputs, targets in loader:
    optimizer.zero_grad()
    with autocast("cuda", dtype=torch.float16):
        loss = criterion(model(inputs), targets)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

The mixed-precision recipe

Keep a master copy of the weights in fp32. Cast them to fp16 or bf16 for the forward and backward matrix multiplies, and accumulate the results in fp32. The optimiser updates the master weights. A change smaller than the low-precision resolution is then not lost between steps.

Autocast chooses the precision per operation. Matrix multiplies and convolutions run in low precision. Reductions, losses and anything numerically delicate stay in fp32. It is a policy rather than a global switch. That policy is why mixed precision usually matches a full-precision run without hand-editing every layer.

With fp16 the gradients can underflow, which means they become zero because the value is too small. Loss scaling multiplies the loss by a large constant before backward, then divides the gradients afterwards. The scaler raises the constant when it sees an overflow and lowers it when the gradients are healthy. bf16 has enough range that scaling is usually unnecessary. That is one reason it is the default on hardware that supports it.

Memory is the real budget

Training memory has four parts: weights, gradients, optimiser state and activations. Adam keeps two fp32 moments per parameter, and mixed precision keeps an fp32 master copy. A 7B model therefore needs about 84 GB for the optimiser state alone, before the low-precision weights and the activations. That number, not the parameter count, is what decides how many devices a run needs.

Activations scale with batch, sequence length, width and depth. At long context they are often the largest term. Gradient checkpointing stores a few activations and recomputes the rest during backward. It trades compute for memory, and it is the standard way to fit a longer sequence without shrinking the model.

Measure before tuning. Four numbers tell you whether a change helped: peak allocated memory, tokens per second, time per step, and a loss curve against a full-precision baseline. A run that is twice as fast and quietly diverging is not faster. It is a shorter run to the same failure.