Lensa ML
Lensa ML

Mixed Precision Training

Step 1 of 6: Memory per Parameter

16 bytes for every weight with Adam

Training in
Tokens per GPU per step2,048
Training memory for the 7B chat model with Adamfp16 weights2 B · 13.5 GBfp16 gradients2 B · 13.5 GBfp32 master weights4 B · 27.0 GBAdam m (fp32)4 B · 27.0 GBAdam v (fp32)4 B · 27.0 GBactivations, 2048 tokens9.1 GBAll of it on one GPU (dashed line: 80 GB)0501001502006.74B params × 16 bytes = 107.8 GB of model statesactivations: about 34 × d_model bytes per token per layer
Bytes per parameter
16
Model states
107.8 GB
Activations
9.1 GB
Total
117.0 GB

Mixed precision still needs 16 bytes per parameter, because the fp32 master copy and both Adam moments stay in fp32. The saving is in activations: 9.1 GB here against 18.3 GB in fp32. The model states alone, 107.8 GB, are more than one 80 GB GPU holds.

Switch between fp32 and mixed precision, then drag Tokens per GPU per step to 8,192. The ZeRO paper calls the 16 bytes "model states" and writes them as 2 + 2 + K bytes per parameter, where K = 12 for mixed-precision Adam. The activation estimate comes from Korthikanti et al.; it assumes attention scores are not stored, as with FlashAttention. Real runs also need workspace buffers and lose some memory to fragmentation. Activation checkpointing cuts the activation term by recomputing parts of the forward pass during backward.

Step 2 of 6: fp16's Narrow Range

Small gradients round to zero

Typical gradient size (center of the spread)1.0e-6
400 illustrative gradient magnitudes on a log axisrounds to 0subnormalnormal fp16inf1e-121e-101e-81e-61e-41e-211e21e41e61e8fp16: smallest subnormal 6.0e-8, smallest normal 6.1e-5, max 65504
Rounded to 0
4.3%
Subnormal
91%
Normal fp16
4.8%
Largest gradient
1.6e-3

Few values are lost outright (4.3%), but 91% land in the subnormal range below 6.1e-5, where fp16 keeps fewer than 10 significant bits and the smallest ones keep only 1 or 2.

Drag Typical gradient size from 1e-9 to 0.1 and watch the bars change color as they cross each fp16 limit. fp16 has 5 exponent bits with a bias of 15, so normal values run from 2^−14 to 65,504, and subnormals extend down to 2^−24 by giving up mantissa bits. Gradients of activations shrink as backpropagation multiplies them through many layers, so the deep end of a network sees the smallest values. The 400 gradients here are seeded samples, not measurements from a real model.

Step 3 of 6: Loss Scaling

Shift gradients into range, then shift back

Loss scale S2^0 = 1
Illustrative gradients × S, stored in fp16rounds to 0subnormalnormal fp16inf1e-121e-101e-81e-61e-41e-211e21e41e61e8backward on S × loss, then divide the gradients by S in fp32Dynamic loss scaling, illustrative run: log2 S by step1620241214160× = inf found, step skipped, S halved
Rounded to 0
7.2%
Overflow to inf
0%
Run: S at step 1
2^16
Run: skipped steps
0

At S = 2^0, 7.2% of the gradients still round to 0. Each doubling of S moves the whole histogram right by a factor of 2 without changing its shape. In the run, step 1 at 2^16 is clean; S doubles after 5 clean steps in a row.

Drag Loss scale S from 2^0 to 2^24, then press Play to run the dynamic scaler. Backpropagation is linear in the loss, so multiplying the loss by S multiplies every gradient by S. The gradients are divided by S before gradient clipping and the optimizer step, so the learning rate keeps its meaning. PyTorch's GradScaler starts at 2^16, halves on an inf or NaN, and doubles after 2,000 clean steps. The run here doubles after 5 to keep the chart short, and its spikes are made up.

Step 4 of 6: fp32 Master Weights

Small updates vanish in fp16

Update size (learning rate × gradient)1.00e-6
Weight w
One update to w in fp16 (ticks: every fp16 value, 9.77e-4 apart)ww + update (exact)stored in fp16Total change after 1,000 identical updatesexact1.00e-3fp32 master9.54e-4fp16 weights only0
fp16 spacing at w
9.77e-4
Update ÷ spacing
1.0e-3
fp16 total change
0
fp32 master change
9.54e-4

The update is 1.0e-3 of the fp16 spacing at w = 1, under half a step, so w + update rounds back to w. All 1,000 updates vanish in fp16; the fp32 master copy keeps 9.54e-4 of the 1.00e-3.

Drag Update size down to 1e-7 with w = 1, then try w = 10. In fp16 the spacing between neighbors is 2^−10 times the weight's power of two, so any update below half of that rounds away. A learning rate of 1e-4 times a gradient of 1e-3 gives 1e-7, far below the 4.9e-4 cutoff at 1.0. The fp32 master copy has 13 more mantissa bits. Adam's m and v stay in fp32 for the same reason: v holds squared gradients, which underflow 16-bit formats first.

Step 5 of 6: bf16

fp32's range with fewer mantissa bits

Typical gradient size (center of the spread)1.0e-8
Range: smallest subnormal to largest finite valuegradientsfp32bf16fp161e-401e-301e-201e-1011e101e201e301e40Precision: every stored value between 1 and 1.0625fp16bf16spacing at 1: fp16 2^-10 = 9.77e-4, bf16 2^-7 = 7.81e-3
fp16 rounded to 0
68%
bf16 rounded to 0
0%
fp16 mean rel. error
74%
bf16 mean rel. error
0.14%

Centered at 1.0e-8, fp16 loses 68% of the gradients to 0 while bf16 loses 0%, since it shares fp32's 8 exponent bits. Without loss scaling, fp16's mean error here is 74% against bf16's 0.14%.

Drag Typical gradient size from 1e-12 to 0.01 and compare the four cards. bf16 is the top 16 bits of an fp32 number: 1 sign bit, 8 exponent bits and 7 mantissa bits. Google Brain designed it for TPUs, and NVIDIA added it with the A100. Converting from fp32 only drops low mantissa bits, so overflow and underflow behave as in fp32 and no loss scale is needed. Training in bf16 still keeps fp32 master weights, because its spacing at 1.0 is 2^−7, eight times coarser than fp16's.

Step 6 of 6: Why It's Faster

Tensor cores, memory and fp8

Tokens per GPU per step4,096
One training step of the 7B chat model on an illustrative GPUTime: 6 × params × tokens ÷ (peak × 40%)fp32 (no tensor cores)21 sTF32 tensor cores2.7 sbf16 mixed precision1.3 sMemory: 16 B/param of states + activations (0 to 200 GB)fp32144 GBtf32144 GBbf16 mixed126 GB80 GB
fp32 step
21 s
bf16 step
1.3 s
Speedup over fp32
16×
bf16 memory
126.1 GB

At 4,096 tokens the bf16 step drops from 21 s to 1.3 s, a 16× gain at equal utilization, and activations shrink from 36.5 GB to 18.3 GB. Newer GPUs add fp8 tensor cores at about twice the bf16 rate.

Drag Tokens per GPU per step and compare the three time bars. The FLOP count uses the rule of thumb of 6 FLOPs per parameter per token: 2 for the forward pass and 4 for the backward pass. The peaks are A100 datasheet figures, 19.5 TFLOPS for fp32, 156 for TF32 and 312 for bf16, and real runs reach a fraction of them. Autocast keeps softmax, normalization and loss reductions in fp32. fp8 training uses E4M3 for weights and activations and E5M2 for gradients, with a scale factor per tensor.

What is mixed precision training?

Mixed precision training runs most of a neural network's math in a 16-bit float format, fp16 or bf16, while keeping a few sensitive values in 32-bit fp32. The forward and backward passes use 16-bit weights and activations, and the optimizer updates an fp32 master copy of the weights. With fp16, a loss scale keeps small gradients from rounding to zero.

Nearly every large model today is trained this way. Matrix multiplications in 16 bits run on GPU tensor cores at many times the fp32 rate, and 16-bit activations halve the memory that the backward pass has to keep. Frameworks expose it as a switch: PyTorch's autocast and GradScaler, or a bf16 setting in DeepSpeed and Megatron-LM. The same ideas carry over to fp8 training on the newest GPUs.

Key ideas

  • Mixed precision runs matrix math in fp16 or bf16 and keeps master weights and optimizer state in fp32.
  • Training with Adam needs about 16 bytes per parameter in mixed precision, the same as plain fp32; the memory saving comes from activations.
  • fp16's normal range ends at 2^−14 below and 65,504 above, so small gradients round to zero without help.
  • Loss scaling multiplies the loss by a power of 2 before backward and divides the gradients by it afterwards.
  • Updates smaller than half the fp16 spacing vanish, which is why the optimizer updates an fp32 copy of the weights.
  • bf16 keeps fp32's exponent range with fewer mantissa bits, so it trains without loss scaling.

Common questions

Should I train with fp16 or bf16?

Use bf16 if the hardware supports it, which includes NVIDIA GPUs from the A100 onward and Google TPUs. It needs no loss scaling and rarely overflows, so runs are simpler and more stable. fp16 remains the choice on older GPUs such as the V100 and T4, where bf16 tensor cores are missing.

Does mixed precision hurt model accuracy?

With loss scaling and fp32 master weights, Micikevicius et al. matched fp32 accuracy across image, speech, translation and language models. Most later large models were trained in mixed precision with no reported loss. Problems usually show up as loss spikes or NaNs, which point to overflow or to an operation that should stay in fp32.

What causes NaN losses in fp16 training?

An activation or gradient above 65,504 becomes inf, and inf turns into NaN through operations such as inf − inf or 0 × inf. Dynamic loss scaling catches overflow in the gradients, but not in the forward pass. Common fixes are keeping softmax and normalization in fp32, lowering the learning rate, or switching to bf16.

What is TF32?

TF32 is a tensor core mode on A100 and later GPUs that reads fp32 inputs, rounds them to 10 mantissa bits with fp32's 8 exponent bits, and accumulates in fp32. Code sees ordinary fp32 tensors, so it speeds up fp32 training without other changes. It is slower than bf16, because each value still takes 4 bytes in memory and the peak rate is half.

Related explainers

References

  1. Micikevicius et al. (2017), "Mixed Precision Training"
  2. Rajbhandari et al. (2019), "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"
  3. Kalamkar et al. (2019), "A Study of BFLOAT16 for Deep Learning Training"
  4. Korthikanti et al. (2022), "Reducing Activation Recomputation in Large Transformer Models"
  5. Micikevicius et al. (2022), "FP8 Formats for Deep Learning"

Last updated