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.

翻译：我们提出了单元缩放（unit scaling）这一设计深度学习模型的范式，旨在简化低精度数格式的使用。在FP16或近期提出的FP8格式下进行训练可带来显著的效率提升，但往往因数值范围不足而难以实现开箱即用的训练。单元缩放通过引入一种模型数值设计的规范化方法来解决该问题：在初始化时追求所有权重、激活值和梯度的单位方差。与替代方法不同，该方案既不需要通过多次训练运行来寻找合适的缩放比例，也不会产生显著的计算开销。我们在多种模型和优化器上验证了单元缩放的有效性。进一步表明，现有模型可以适配为单元缩放形式，在FP16甚至FP8精度下训练BERT-Large模型且精度无任何损失。