fp32 (float32)
fp16 (float16)
bf16 (bfloat16)
tf32 (CUDA internal data type)
Here is a diagram that shows how these data types correlate to each other.
(source: NVIDIA Blog)
While fp16 and fp32 have been around for quite some time, bf16 and tf32 are only available on the Ampere architecture GPUS. TPUs support bf16 as well.
fp16
AMP = Automatic Mixed Precision
If we look at what's happening with FP16 training (mixed precision) we have:
- the model has two copies in memory: one in half-precision for the forward/backward computations and one in full precision - no memory saved here
- the forward activations saved for gradient computation are in half-precision - memory is saved here
- the gradients are computed in half-precision but converted to full-precision for the update, no saving there
- the optimizer states are in full precision as all the updates are done in full-precision
