BFloat16 in general but with LayerNorm in FP32
The bfloat16 (BF16) [8] floating-point format provides the same dynamic range as FP32 with a memory footprint identical to FP16. We found that we could train models in the BF16 format using the same set of hyperparameters as in FP32, without special parameter tuning. Nevertheless, we found that we need to keep LayerNorm in FP32 mode in order for the training to converge.
3.4 Final training recipe
A summary of the final training recipe.

Wrap the outer model in an FSDP instance. Enable parameter sharding after the forward pass.
Wrap individual ViT blocks with activation checkpointing, nested FSDP wrapping, and parameter flattening.
Enable mixed precision mode (AMP O2) with bfloat16 representation. Maintain the optimizer state in FP32 precision to enhance numerical stability.
Wrap normalization layers like LayerNorm in FP32 for better numerical stability.
