Revisiting BFfloat16 Training

State-of-the-art generic low-precision training algorithms use a mix of 16-bit and 32-bit precision , creating the folklore that 16-bit precision alone is not enough to maximize model accuracy . As a result , deep learning accelerators are forced to support both 16-bit and 32-bit compute units which is more costly than only using 16-bit units for hardware design . We ask can we do pure 16-bit training which requires only 16-bit compute units , while still matching the model accuracy attained by 32-bit training . Towards this end , we study pure 16-bit training algorithms on the widely adopted BFloat16 compute unit . While these units conventionally use nearest rounding to cast output to 16-bit precision , we show that nearest rounding for model weight updates can often cancel small updates , which degrades the convergence and model accuracy . Motivated by this , we identify two simple existing techniques , stochastic rounding and Kahan summation , to remedy the model accuracy degradation in pure 16-bit training . We empirically show that these two techniques can enable up to 7 % absolute validation accuracy gain in pure 16-bit training . This leads to 0.1 % lower to 0.2 % higher matching validation accuracy compared to 32-bit precision training across seven deep learning applications . 1 INTRODUCTION . Recently there has been an explosion in the compute resources required for training deep learning models ( Shoeybi et al. , 2019 ; Rajbhandari et al. , 2019 ; Real et al. , 2019 ) . As a result , there has been broad interest in leveraging low-precision ( < 32-bit ) training algorithms to reduce the required compute resources ( De Sa et al. , 2017 ; Hubara et al. , 2017 ; Gupta et al. , 2015 ) . Among these algorithms , mixed-precision training—in which model activations and gradients are stored using a 16-bit floating point format while model weights and optimizer states use 32-bit precision—is commonly used when training generic deep learning models ( Micikevicius et al. , 2017 ; Kalamkar et al. , 2019 ) . While there is a wide body of literature showing that low-precision training can minimally impact accuracy on specific models ( Wang et al. , 2018b ; De Sa et al. , 2015 ; Zhang et al. , 2017 ) , conventional wisdom suggests that at least some 32-bit computation is required as a fail-safe in generic deep learning training . As such , new accelerator architectures for deep learning are forced to support both 32-bit and 16-bit compute units . This is much more costly in terms of area , power , and speed when compared to hardware with only 16-bit compute units ( Horowitz , 2014 ; Galal et al. , 2013 ) . In this paper we question if 32-bit compute units are truly needed for new deep learning hardware accelerators . Namely , can we match the model accuracy of 32-bit-precision algorithms while leveraging only 16-bit compute units ? To answer this question , we study pure 16-bit training algorithms , ones which use only 16-bit compute units and which store activations , gradients , model weights , and optimizer states all in a 16-bit precision . Specifically , we focus on training with the BFloat16 compute unit which is widely adopted in modern deep learning accelerators ( Jouppi et al. , 2017 ; Burgess et al. , 2019 ) . Such units take 16-bit inputs , perform computation , and then round the results to a 16- bit output . BFloat16 compute units can provide 3⇥ higher power efficiency , 1.5⇥ lower latency , and 1.5⇥ less chip area than 32-bit units ( Horowitz , 2014 ; Galal et al. , 2013 ) . In addition , pure 16-bit training algorithms can reduce the memory footprint and bandwidth consumption of model weights and optimizers by 2⇥ compared to mixed precision or 32-bit precision training , especially for large models with billions of weights ( Shoeybi et al. , 2019 ; Rajbhandari et al. , 2019 ) . Developing reliable pure 16-bit training algorithms will enable hardware designers to realize these advantages . The simplest approach to pure 16-bit training is to take a 32-bit baseline and “ make it low-precision ” by replacing all the 32-bit numbers with 16-bit numbers and replacing each 32-bit floating-point op- eration with its 16-bit analog , using nearest rounding1 to quantize as necessary : we call this approach the standard algorithm . Unfortunately , we show empirically that standard pure 16-bit training does not match 32-bit training on model accuracy across deep learning models . For example , the standard pure 16-bit training algorithm one would run on conventional hardware attains 16 % and 7 % lower training and validation accuracies than a 32-bit baseline . Motivated by this observation , we start by analyzing what factors limit the model accuracy of this standard pure 16-bit algorithm . The goal of our analysis is to inspire a clean , minimal set of simple techniques that allow pure 16-bit training to attain strong model accuracy for state-of-the-art deep learning models across application domains . Towards this end , we derive insights from a simple least-squares regression model in Section 3 . Using this least-squares regression model , we reveal that nearest rounding of compute unit outputs causes significant convergence degradation and consequent model accuracy loss . More concretely , we show a key theoretical insight hidden in existing work : when running stochastic gradient descent on a least-squares regression model , nearest rounding while updating model weights ignores small updates . This phenomenon significantly degrades the convergence of stochastic gradient descent when model updates become small relative to model weights , which is also what we observe when training deep learning models . In comparison , nearest rounding in the forward and backward pass of backpropagation has a negligible impact on convergence . These insights lead us to consider two simple existing techniques to achieve high-accuracy pure 16-bit training . First , we can use stochastic rounding instead of nearest rounding for the model weight updates . Here , the rounded weights become an unbiased estimate of the precise weights without rounding : thus , regardless of the magnitude of updates , the expectation of rounded weights converges at the same speed as the precise weights . Second , we can use the well-known Kahan summation algorithm ( Kahan , 1965 ) to accumulate model updates while still keeping nearest rounding for all operations . This method tracks and compensates weight rounding errors across iterations with auxiliary 16-bit values , which avoids catastrophic cancellation of many small model weight updates . Empirically , in Section 4 we first validate that , as suggested by our theory , nearest rounding for model weight updates is the sole bottleneck for convergence and model accuracy on several deep learning models . We then demonstrate that pure 16-bit training using stochastic rounding or Kahan summation on model weight updates can match 32-bit training in model accuracy across a wide range of applications ( He et al. , 2016 ; Amodei et al. , 2016 ; Devlin et al. , 2018 ; Naumov et al. , 2019 ) . To validate that nearest rounding for model weight updates is the cause of the accuracy degradation , we show that if we store model weights in 32-bit precision without rounding during weight updates , and we keep using 16-bits and nearest rounding for all other operations , then the attained model accuracy matches full 32-bit precision training . Next , we demonstrate that 16-bit training with stochastic rounding for weight updates attains model accuracy matching 32-bit training for five out of seven applications in our study . Note that while it works most of the time , this is not a silver bullet , as using stochastic rounding alone could not fully match 32-bit training on all models . To address this , we show that Kahan summation for model weight updates closes remaining gaps on all the models we consider ; this Kahan summation comes with a trade off , as it requires 2⇥ weight memory , but achieves up to 0.2 % higher validation accuracy than stochastic rounding . Our results suggest that deep learning accelerators using only 16-bit compute units are feasible if stochastic rounding and Kahan summation are supported respectively by the hardware and the software stack . 2 PRELIMINARY . In this section we establish the background and notation for our study and present the preliminary observations that motivate our work . We focus on the case of stochastic gradient descent ( SGD ) , which is the primary workhorse used to train deep learning models . SGD computes gradients from a subset of training samples , and uses them to update the model weights so as to decrease the loss in expectation . In the classic supervised learning setting , let ( X , y ) be a dataset where X = [ x1 , x2 , ... , xn ] 2 Rn⇥d and y = ( y1 , y2 , ... , yn ) 2 Rn . On this dataset , we use stochastic gradient descent to optimize a loss function f ( w ) = 1/n Pn i=1 fi ( w , xi , yi ) defined by the model . At the t-th iteration , we sample an index subset ( t ) ⇢ { 1 , 2 , .. , n } and compute a sample gradient rf ( t ) ( wt ) as an unbiased estimate of the full gradient rf ( w ) . In deep learning , model training can be described as a compute graph where the compute graph operators such as addition and ma- 1This nearest rounding is the standard rounding mode for compute unit output commonly supported across hardware platforms ( Intel , 2018 ; Nvidia , 2020 ) . trix multiplication are the nodes . For example , the model weight update operator is defined as the subtraction in the operation wt+1 = wt ↵rf ( t ) ( wt ) which updates the model weight w. 16-bit Compute Units On modern hardware , numerical operators in the compute graph are supported by fused multiply-and-accumulation ( FMAC ) compute units . These units work by computing a a + ( x ⇥ y ) , where x and y are input floating point numbers , and a is an accumulator that is part of the FMAC unit . Importantly , for a 16-bit FMAC unit , the accumulator a has higher-than-16bit precision . This higher precision accumulator is inexpensive compared to the multiply in terms of chip area and energy consumption in FMAC units , but is critical to the numerical accuracy of operations such as matrix multiplication and convolution . Thus 16-bit FMAC units with higher precision accumulator is standard for modern hardware including TPUs and GPUs ( Chao & Saeta , 2019 ; Markidis et al. , 2018 ; Stephens , 2019 ) ; this will likely continue to be standard in emerging accelerators . Because of the higher precision accumulator , the result in the accumulator then needs to be rounded to 16-bits before it is output from the FMAC unit ( e.g . before writing to memory ) . FMAC units use the same hardware implementation to support all operators from simple additions to computationally intensive convolutions , so this output-rounding step happens for all the operators in a compute graph . Nearest Rounding FMAC output rounding is widely implemented with nearest rounding as the standard mode , due to its efficient support across hardware platforms . “ Nearest rounding ” means rounding a higher precision floating point number to the closest low-precision representable value . Given that the add step already uses accurate higher preicision accumulators , this nearest rounding is the primary source of numerical errors for training using 16-bit FMAC units . In this context of 16 bit FMAC units and nearest rounding , we discuss the following training-precision approaches . • In 32-bit precision training , all the compute graph operators read and write memory using a 32- bit precision . These operators require 32-bit compute units , which constrains the compute and memory efficiency . • In mixed precision training , model weights are stored in 32-bit precision while activations and gradients use 16-bit precision . Thus , new accelerators customized to maximize efficiency for mixed precision training still require 32-bit compute units for operators involving 32-bit weights as the input ; this has lower efficiency in power , speed and chip area than only using 16-bit units . • In pure 16-bit training , activation , gradients and model weights are all stored in 16-bit precision . All operators use pure 16-bit input and write out 16-bit output after rounding . Thus , aside from just saving memory , pure 16-bit training can eliminate the requirement for 32-bit compute units . This opens the possibility for highly-efficient accelerators using only 16-bit compute units . In spite of this favorable efficiency , we now show that it can be surprisingly challenging for standard pure 16-bit training to match 32-bit training on model accuracy . Motivating Observations Although recent works have shown that certain models are robust to numerical error during training ( Wang et al. , 2018b ; De Sa et al. , 2015 ; Zhang et al. , 2017 ) , surprisingly , we observe that it is challenging for pure 16-bit training to attain the same accuracy as 32-bit precision training on several state-of-the-art deep learning models . To demonstrate this , we compare 32-bit precision training and standard pure 16-bit training ( using nearest rounding for all its operator outputs ) . For example , Figure 1 illustrates that for a BERT-Base model for natural language inference , the standard pure 16-bit training algorithm demonstrates 16 % and 7 % lower training and validation accuracies than 32-bit precision training2 . This gap suggests that nearest rounding is the primary source of numerical error in pure 16-bit algorithms , significantly degrading the convergence and model accuracy . To alleviate this problem , in Section 3 , we study how nearest rounding impacts convergence , and we expose insights which lead to simple techniques to improve the model accuracy in pure 16-bit training . 2We note that 32-bit and standard 16-bit training in Figure 1 start from the same initialization with the same model accuracy . The gap at the beginning of the two curves is due to curve smoothing for visualization clarity . We refer to Appendix D.1 for the unsmoothed curves which start from the same model accuracy .

Summary: This paper explores the possibilities of reducing the precision of the weight update operation (i.e. AXPY ops) from 32 bit to 16 bit in today’s BFloat16 training framework. To enable 16 bit update, the authors proposed two techniques, i.e. stochastic rounding and Kahan summation. The authors use a simple least-squares regression model to theoretically explain the model accuracy degradation due to nearest rounding, then experimentally demonstrate the effectiveness of two techniques on a range of deep learning models and dataset.