Mixed-precision training stores activations and most compute in lower precision (FP16 or bfloat16) while keeping a master copy of model weights in FP32. During each step: forward/backward use FP16 for speed and reduced memory; gradients are accumulated/scaled, converted to FP32 to update the master weights; updated FP32 weights are cast to FP16 for the next forward.Benefits- Memory: ~2x reduction for activations/weights stored in FP16 → larger batch sizes or bigger models.- Throughput: tensor cores on modern GPUs/TPUs accelerate FP16/bfloat16 math → 2–8x faster depending on layer types and hardware.Why loss scaling- FP16 has smaller dynamic range; small gradients can underflow to zero. Loss scaling multiplies the loss by a factor S before backprop, preventing underflow. Gradients are divided by S before updating. Dynamic loss scaling adjusts S to avoid overflow.Numerical pitfalls- Overflow: large intermediate values can become Inf; detect and reduce scale or skip the step.- Underflow: tiny gradients become zero without scaling.- Accumulation error: repeated FP16 ops reduce precision; master FP32 weights mitigate this.- Ops not safe in FP16 (softmax reduction, layernorm) should run in FP32 if necessary.Enable mixed precisionPyTorch (native AMP):python
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for x,y in loader:
optimizer.zero_grad()
with autocast(): # mixed precision context
loss = model(x).loss_fn(y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
TensorFlow (Keras mixed precision):python
from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')
# optimizer automatically uses a loss scale; or use LossScaleOptimizer wrapper
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['acc'])
When mixed precision can degrade performance- Models dominated by small-batch or memory-bound CPU/GPU kernels with little FP16 acceleration.- Numerical sensitivity: very deep recurrent models, tiny gradients, or algorithms relying on high-precision accumulation (some GANs, certain optimizers) may show instability or convergence regression.- Hardware without efficient FP16/bfloat16 support — may be slower.Best practices- Use bfloat16 on TPUs/modern accelerators as it has wider dynamic range.- Enable AMP + dynamic loss scaling; monitor for NaNs/Infs and validate final accuracy vs FP32 baseline.