DL0014 Mixed Precision Training

Mixed precision training runs the compute-heavy parts of a model in FP16 while keeping an FP32 master copy of the weights, so training gets the speed and memory of half precision without sacrificing final accuracy. Modern GPU/TPU tensor cores execute FP16 matrix math several times faster than FP32, and halving activation memory lets you train larger models or use larger batches on the same hardware.

(1) Faster Training: FP16 tensor-core matmuls deliver up to an order of magnitude more throughput than FP32 on supported hardware (e.g., ~312 vs ~19.5 TFLOPS on an A100).
(2) Reduced Memory Usage: FP16 activations and working weight copies occupy half the bytes, freeing room for larger batch sizes or deeper models (master weights and optimizer states stay FP32, so total training memory falls by less than half).
(3) Maintained Accuracy: FP32 master weights plus loss scaling keep small gradient values representable, so final model quality matches full-precision training.

Figure 1: FP16 trades exponent range and mantissa precision for half the storage: gradients below would underflow to zero without loss scaling.

The Training Loop: Weights are stored in FP32 as the master copy. Each step casts them to FP16 for the forward and backward passes, multiplies the loss by a scale factor so that FP16 gradients stay in range, then divides the gradients by and applies the optimizer update to the FP32 master weights.

Figure 2: FP16 does the heavy math while the FP32 master copy absorbs tiny updates; loss scaling shifts gradients into FP16’s representable range.

Measured Benefits: On tensor-core hardware the speedup is substantial, and the halved activation memory (the dominant term at large batch sizes) directly translates into larger feasible models or batches.

Figure 3: Roughly 16x tensor-core throughput and half the activation memory are the headline wins; with Adam states kept in FP32, per-parameter training memory drops only modestly.

Mathematical Formulation:

Where:

• is the loss-scale factor (e.g., , or dynamically adjusted); is the scaled loss used for backprop in FP16.
• are the scaled FP16 gradients; dividing by restores the true gradient .
• is the FP32 master weight set and the learning rate; updates always land on the master copy.

Costs to Manage: FP16’s narrow range causes gradient underflow and occasional activation overflow, requiring loss scaling and careful debugging of NaN/Inf values; efficiency also depends on hardware with fast FP16 paths.