Unit Scaling: Out-of-the-Box Low-Precision Training
This addresses the problem of efficiency in low-precision training for deep learning practitioners, offering a novel paradigm rather than an incremental improvement.
The paper tackles the challenge of training deep learning models in low-precision formats like FP16 and FP8 without degradation, by introducing unit scaling to ensure unit variance of weights, activations, and gradients at initialization, enabling BERT-Large to be trained in FP16 and FP8 with no accuracy loss.
We present unit scaling, a paradigm for designing deep learning models that simplifies the use of low-precision number formats. Training in FP16 or the recently proposed FP8 formats offers substantial efficiency gains, but can lack sufficient range for out-of-the-box training. Unit scaling addresses this by introducing a principled approach to model numerics: seeking unit variance of all weights, activations and gradients at initialisation. Unlike alternative methods, this approach neither requires multiple training runs to find a suitable scale nor has significant computational overhead. We demonstrate the efficacy of unit scaling across a range of models and optimisers. We further show that existing models can be adapted to be unit-scaled, training BERT-Large in FP16 and then FP8 with no degradation in accuracy.