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: The authors propose an adaptive loss scaling method during the backpropagation stage for the mix precision training to reduce the underflow. Compared with the previous work, which scales the loss by human design, and needs to be consistent in all layers. The authors state that they can decide the scale rate layer by layer automatically to reduce the underflow in a low precision situation.