

AMP O2: This refers to training in mixed precision but with more weights and ops in FP16 than in O1. Weights do not implicitly remain in FP32 and are cast to FP16. A copy of the master weights is maintained in the FP32 precision that is used by the optimizer. If we want the normalization layer weights in FP32 then we need to explicitly use layer wrapping to ensure that.


Full FP16: This refers to training in full FP16 where weights and operations are in FP16. FP16  is challenging to enable for training due to convergence issues.


We found that AMP O2 with LayerNorm wrapping in FP32 leads to the best performance without sacrificing accuracy.
Kernel Fusion

To reduce GPU kernel launch overhead and increase GPU work granularity, we experimented with kernel fusions, including fused dropout and fused layer-norm, using the xformers library [7].

3.3 Addressing stability challenges by studying ops numerical stability and training recipes