Photo by Norbert Braun on Unsplash

E31 : Mixed Precision Training

Training the neural networks using half-precision floating point numbers(FP16) alongside adopting three techniques, nearly halves the memory requirement and improves the compute speed without any impact to the accuracy of the same networks when trained on single-precision floating point numbers(FP32)

Paper Name : Mixed Precision Training

Paper URL : https://arxiv.org/abs/1710.03740

Authors : Baidu Research - Sharan Narang, Gregory Diamos, Erich Elsen NVIDIA - Paulius Micikevicius , Jonah Alben, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, Hao Wu

Conference : ICLR 2018

Please find the annotated paper **here**.

Problem Statement :

• Training a neural network using FP32 precision, requires abundant memory requirements to store the model weights, activations and gradients.
• Arithmetic operations using FP32 precision floating point numbers increases the arithmetic bandwidth as well during training and inference times

Solution :

• Instead of using single precision (FP32) to store the model parameters and use for arithmetic operations, opting for half-precision (FP16) helps to reduce the memory and arithmetic operations overhead at training and inference time.
• Storing a number in FP32 precision requires 8 bytes of memory, whereas storing the same number in FP16 precision requires 4 bytes of memory.
• Assume a neural network with 1B parameters in FP32 precision would approx. require 4GB of memory, whereas the same network in FP16 precision would require only 2GB approx.

• No changes to model architecture or hyperparameter changes were done as part of this training process

MIxed Precision Training (end to end flow)

Approach :

• In addition to the original training process, three techniques are followed to avoid model accuracy loss. (i) maintain master copy of model weights in FP32 precision (ii) loss scaling to minimize the number of weights gradients from becoming zero during backward pass (iii) FP16 for arithmetic operations but accumulated in FP32 precision

(i) Maintain master copy of model weights in FP32 precision

• A master copy of the model weights is stored in FP32 precision and the original model weights are then converted from FP32 to FP16 and the activations are calculated using FP16 precision weights during the forward pass. The activations thus calculated would also be in FP16 precision.
• The loss is then calculated and backpropagated and thus the weight gradients are calculated. The gradients thus calculated would be very small and experiments show that about 5% of the gradients have exponents less than -24 which when represented in FP16 precision would become zero.
• Hence the gradients would become zero either before or after multiplication with the learning rate. In case if it becomes zero, then the new updated weights would be the same the old weights and hence no learning would happen in the neural network. Hence the gradients are converted into FP32 precision.

Multiplying the weight gradient by the learning rate would further make it smaller

• These gradients are then multiplied with the learning rate which might result in much more smaller values but it would not get converted zero as they are already in FP32 precision now. The FP32 master copy weights are updated with the multiplied value.

(ii) Loss Scaling

• The loss calculated after the forward pass is scaled by a factor ‘S’. This in turn increases the weight gradients by a factor ‘S’. This prevents the weight gradients from rounding upto zero in FP16 precision even before getting multiplied by a learning rate.
• The scaled weight gradients are then descaled by diving by the factor ‘S’ and then converted to FP32 and then multiplied with the learning rate.

(iii) FP16 for arithmetic operations but accumulated in FP32 precision

• Arithmetic operations involved dot product (vector-vector multiplication, matrix vector multiplication, matrix-matrix multiplication) are not limited by compute speed but rather bound by memory bandwidth.
• Hence when doing a dot product calculation during training, the parameters were multiplied in FP16 precision and before accumulating them was converted to FP32 precision so that the precision or accuracy of the parameter is not lost when adding multiple elements of the matrix and the accumulated vector is then cast back to FP16 precision before writing it to memory as it constrained by memory bandwidth.
• In case of reductions like batch-normalization, softmax which are memory bandwidth constrained, the arithmetic operation like adding the parameters were done in FP32 precision to maintain the accuracy and then cast back to FP16 precision before writing it to memory.
• Point-wise operations like activation functions would not be impacted by the compute speed or by memory bandwidth so the choice can be made to maintain FP16 or FP32 precision

Experiments : Tasks tested - Classification, regression, Generative Application - Image classification, Image generation, Object detection, Language modelling, Machine Translation and Speech recognition

Observations :

• Experiments show that about 5% of the gradients have exponents less than -24 which when represented in FP16 precision would become zero

5% weight gradients would become zero in FP16

• It was observed that training a model by maintaining a master copy of FP32 weights and updating it by casting the FP16 precision gradients to F32 precision and then multiplying it with learning rate matched the performance of a model trainined only with FP32 precision, whereas a model completely trained in FP16 precision observed 80% accuracy loss

Training and validation loss for baseline (compelete FP32 training) vs mixedprecision (with mastercopy FP32 weights) vs mixedprecision (without mastercopy of FP32 weights)

• 5 out of 6 CNNs trained with MPT were able to outperform the results of baseline training and this was achieved only master copy of FP32 weights and without adopting loss scaling technique
• During training with Multibox SSD detector network for object detection, it was observed that scaling the loss only by a factor 8 was enough to match the performance of mixed precision training with baseline FP32 training.Loss scaling helped to shift some of the non-representable range parameters in FP16 precision to representable values in FP32 precision.
• MP training outperformed baseline training in both English and Mandarin dataset for speech recognition task. Results also show that the half-precision storage may act as a regularizer during training.
• MP training was able to match the performance of baseline training with loss scaling technique included on Machine Translation and Language Modeling tasks as well
• The qualitative comparison of images generated by GAN model trained using MP training showed better results compared to baseline testing

Conclusion :

• Given the size of the language models being built these days, a technique like Mixed Precision Training where a small change in the precision of the parameters from FP32 to FP16 can have a huge positive impact on the memory and arithmetic requirements of the model being trained without reducing the model accuracy is a huge plus.