Adaptive Loss Scaling for Mixed Precision Training

1 INTRODUCTION . Training deep neural networks ( DNNs ) is well-known to be time and energy consuming , motivating the development of new methods and hardware to make training more efficient . One way to improve training efficiency is to use numerical representations that are more hardware-friendly . This is the reason that the IEEE 754 32-bit single-precision floating point format ( FP32 ) is more widely used for training DNNs than the more precise double precision format ( FP64 ) , which is commonly used in other areas of high-performance computing . In an effort to further improve hardware efficiency , there has been increasing interest in using data types with even lower precision than FP32 for training ( Micikevicius et al. , 2018 ; Kuchaiev et al. , 2018 ; Wang et al. , 2018 ; Kalamkar et al. , 2019 ; Mellempudi et al. , 2019 ; Sakr et al. , 2019 ) . Of these , the IEEE half-precision floating-point ( FP16 ) format is already well supported by modern GPU vendors ( Choquette et al. , 2018 ) . Using FP16 for training can reduce the memory footprint by half compared to FP32 and significantly improve the runtime performance and power efficiency . Nevertheless , numerical issues like overflow , underflow , and rounding errors frequently occur when training in low precision only . Recent works propose various improvements , of which mixed precision training ( MPT ) ( Micikevicius et al. , 2018 ) is the state-of-the-art . Its core idea is to use FP16 for the compute-intensive yet precision-insensitive operations , such as matrix multiplication , for computational efficiency , while using FP32 for the operations that require high precision , such as batch normalization ( Ioffe & Szegedy , 2015 ) and gradient update accumulation . Activations and gradients , which largely contribute to memory consumption , are stored in FP16 , while the weights are stored in FP32 for more accurate accumulation of gradient updates . Even though MPT seems promising and has wide support from both hardware and software frameworks , it still suffers from reliability issues , mainly due to the more limited dynamic range of FP16 being unable to adequately cover possible gradient values during training . The most common issue is for small gradients to fall into the underflow gap and become zero , which makes training less effective . Loss scaling ( Micikevicius et al. , 2018 ; Kuchaiev et al. , 2018 ; Mellempudi et al. , 2019 ) addresses the range limitation in FP16 by introducing a hyperparameter α to scale the loss value before the start of the backward pass so that the computed ( scaled ) gradients can then be properly represented in FP16 without significant underflow . For an appropriate choice of α , loss scaling can achieve state of the art results that are competitive with regular FP32 training . Unfortunately , there is no single value of α that will work in arbitrary models , and so it often needs to be tuned per model . Its value must be chosen large enough to prevent underflow issues from affecting training accuracy . However , if α is chosen too large , it could amplify rounding errors caused by swamping ( Higham , 1993 ) or even result in overflow . This observed sensitivity to the particular choice of loss scale is also reported by Mellempudi et al . ( 2019 ) , who find that different values can lead to very different ResNet-50 MPT convergence behavior . Furthermore , the data distribution of gradients can vary both between layers and between iterations ( Figure 1 ) , which implies that a single scale is insufficient . For instance , gradients closer to the input require a higher loss scale that may cause overflow or severe rounding errors if the same value were used in layers closer to the output . Including the time spent tuning α , the total training time of MPT can even exceed regular FP32 training . We introduce a loss scaling-based training method called adaptive loss scaling that makes MPT easier and more practical to use . We hope that this will help to utilize better existing hardware with support for fast FP16 operations . Our method improves the usability of MPT compared to existing methods by removing the need to tune a model-specific loss scale hyperparameter , while retaining ( and in some cases surpassing ) the accuracy of regular FP32 training . We achieve this by introducing layer-wise loss scale values which are automatically computed and dynamically updated during training to deal with underflow more effectively than existing methods . Experimental results on several examples show that MPT with adaptive loss scaling can achieve the best model accuracy and the shortest overall training time , especially when training deep models on large datasets . 2 BACKGROUND . 2.1 PRELIMINARY . As mentioned above , MPT ( Micikevicius et al. , 2018 ) uses FP16 for storing the activations and gradients and for the most compute-intensive tasks , while FP32 is used only where increased precision is required . FP16 has three fewer exponent bits than FP32 , limiting its dynamic range to magnitudes between umin = 2−24 and umax = 65505 . In practice , the gradients often have a larger range than this , resulting in numerical issues when using MPT . In FP16 , if the absolute actual value of a gradient |g| is smaller than umin , it will become 0 ; and if it is larger than umax , it will be infinite . Also , even if a value is in [ umin , umax ) , the closer it comes to either bound , the less accurate its FP16 form is regarding absolute rounding error , e.g. , 1024.1 is rounded to 1024 . Underflow motivates loss scaling , while the overflow and rounding error are what loss scaling should be careful with . Figure 2 : Comparison between the standard backpropagation algorithm and the loss scaled one . Each layer has a single output in this formulation . Output finds the output layer index , GradW calculates the gradient update with respect to the weights . Backprop calculates the activation gradient given the layer type OP and the input gradient . δN+1 ← initial error gradient ; for i← layer indices in reversed topological order do j ← Output ( i ) ; Wi ←Wi+ GradW ( δj ) ; δi ← Backprop ( OP ( i ) , δj ) ; end Algorithm 1 : Standard backpropagation algorithm . δN+1 ← initial error gradient ; δN+1 ← αδN+1 ; for i← layer indices in reversed topological order do j ← Output ( i ) ; Wi←Wi+ GradW ( δj ) /α ; δi← Backprop ( OP ( i ) , δj ) ; end Algorithm 2 : Standard loss scaling algorithm . α is the loss scale . Figure 2 shows the basic loss scaling algorithm ( Algorithm 2 ) compared to standard backpropagation without loss scaling ( Algorithm 1 ) . Note that they differ only in that in Algorithm 2 , the initial error gradients δN+1 from the output layer are scaled by α before the start of the backward pass , and that the weight gradient update GradW is then unscaled by the same α just before the weight update . Recall that α should be chosen large enough to prevent underflow issues from affecting training , while also being small enough to prevent overflow . Even when kept within this range , using a larger value than necessary could introduce more absolute rounding error as aforementioned . Also , since loss scaling amplifies the ratio between the largest and the smallest elements within each gradient , swamping ( Higham , 1993 ) , the phenomenon that summing small values with larger ones is inaccurate in floating-point , becomes more likely and may hinder training ( Wang et al. , 2018 ) . 2.2 RELATED WORK . Many recent works focus on reducing rounding error to improve the training performance . Wang et al . ( 2018 ) devise a chunk-based accumulation mechanism to mitigate the swamping issue . Sakr et al . ( 2019 ) improve the solution to the same problem by finding lower precision for accumulation by variance analysis . Alternatively , Hoffer et al . ( 2018 ) identify the numerical issues caused by batch normalization and propose to replace it by a more numerically stable and efficient alternative . These methods are orthogonal to loss scaling , and we plan to study the effect of applying them together with adaptive loss scaling as future work . Loss scaling that aims to improve mixed precision training by reducing the underflow rate in the computed gradients can be traced back to Micikevicius et al . ( 2018 ) . They originally suggest to choose a constant loss scale either empirically , or using a factor that can not scale the maximal absolute gradient value to overflow . Kuchaiev et al . ( 2018 ) propose two improved versions : one is called backoff , which simply makes the loss scale smaller if a numerical error is encountered during training ; the other is logmax , which models the maximal absolute gradient value across training iterations by log-normal distribution , in order to estimate the proper loss scale value for the next iteration . We argue that these new solutions are still not ideal , since backoff is simply trial-and-error and can waste training workload ; and logmax is risky to use when we do not have much gradient values to model that log normal distribution , or this assumption can not apply . Mellempudi et al . ( 2019 ) further study the effect of the backoff method for 8-bit floating point . 3 ADAPTIVE LOSS SCALING . As a preview , Figure 3 shows a concrete example of our adaptive loss scaling approach in the backward pass of a 3-layer Multi-Layer Perceptron ( MLP ) . After the forward pass has completed , starting from the rightmost node , the gradients δ are first propagated from loss L and are then scaled by a scalar α4 before being consumed by the last linear layer . The weight gradients for this layer are then scaled by 1/α4 just before the weight update for W3 in order to prevent the particular choice of α4 from affecting the computed gradient magnitudes . It is at this point that our approach begins to differ from the standard loss scaling method . In the standard method , the same α4 would be used for all layers in the network . However , in our method , each layer i calculates its own local loss scale value βi based on the statistics of its output gradients and weights in the current iteration , in order to minimize underflow in its computed input gradients . This βi is then used to scale weightsWi before computing the scaled input activation gradients for layer i . Since the scaling effects from these local loss scales accumulate , when unscaling gradients for updating , we use the product of all previous scale values , indicated by αi = α4 ∏3 j=i+1 βj . Thus , our approach attempts to minimize underflow in every layer simultaneously through the use of layer-local loss scales βi which are also computed automatically based on the current layer statistics . Compared to the standard method , this removes the need to perform model-specific hyperparameter tuning and enables layer-wise loss scaling . 3.1 LOSS SCALED BACKPROPAGATION . We use a 2-tuple notation to denote the propagated entity from layer i : 〈αi , δi〉 , in which αi is the loss scale value for layer i and δi is the gradient that has been scaled by αi . To be more specific about this notation , we can take aN -layer MLP as an example ( see Figure 3 for the notation ) . In this case , layer i takes in 〈αi+1 , δi+1〉 , updates its weight by ( yTi−1δi+1 ) /αi+1 , and produces 〈αi , δi〉 for the previous layer i− 1 . We will elaborate more on how αi is calculated in the following sections . Algorithm 3 : Backpropagation algorithm with adaptive loss scaling , assuming each layer has a single output . Section 3.2.2 shows how multiple outputs work . 〈αN+1 , δN+1〉 ← initial loss scale , and error gradient scaled by α0 for i← layer ID in a reversed topological order do j ← GetLayerOutput ( i ) Wi ←Wi+ GetWeightGradient ( δj ) /αj βi← GetLossScale ( OP ( i ) , 〈αj , δj〉 ) 〈αi , δi〉 ← 〈αjβi , Backprop ( OP ( i ) , βiδj ) 〉 end Algorithm 3 shows the pseudocode for adaptive loss scaled backpropagation for the case where each layer has a single output ( we describe how to handle the multiple-output case in Section 3.2.2 ) : 1 . We start with the error gradients δN+1 computed from the output loss value for the last layer N + 1 . We may optionally scale this gradient by αN+1 . Normally we keep it as 1 . 2 . As visiting each previous layer i in a reversed topological order of the computational graph , we retrieve the tuple 〈αj , δj〉 propagated to it from the next downstream layer that represents the scaled loss for layer i ’ s output . We calculate a local loss scale value βi , which will be used to scale δj before we calculate the activation gradients for the previous layer . 3 . We use δj and other cached inputs ( omitted ) to compute the gradients for Wi . However , since these gradients have been scaled , we must unscale them using αj before performing the weight update . 4 . Since βi contributes to the magnitude of the gradient δi , we calculate the loss scale value δi to be passed to the next previous layer as αjβi .

Summary: In this paper, the authors propose a method to train models in FP16 precision. The authors show that the key reason of training performance drop is the overflow or underflow of back propagation information. Instead of using a fixed value or dynamic value proposed by a previous work, this paper adopts a more elaborate way to minimize underflow