8-bit Optimizers via Block-wise Quantization

Increasing model size is an effective way to achieve better performance for given resources ( Kaplan et al. , 2020 ; Henighan et al. , 2020 ; Raffel et al. , 2019 ; Lewis et al. , 2021 ) . However , training such large models requires storing the model , gradient , and state of the optimizer ( e.g. , exponentially smoothed sum and squared sum of previous gradients for Adam ) , all in a fixed amount of available memory . Although significant research has focused on enabling larger model training by reducing or efficiently distributing the memory required for the model parameters ( Shoeybi et al. , 2019 ; Lepikhin et al. , 2020 ; Fedus et al. , 2021 ; Brown et al. , 2020 ; Rajbhandari et al. , 2020 ) , reducing the memory footprint of optimizer gradient statistics is much less studied . This is a significant missed opportunity since these optimizer states use 33-75 % of the total memory footprint during training . For example , the Adam optimizer states for the largest GPT-2 ( Radford et al. , 2019 ) and T5 ( Raffel et al. , 2019 ) models are 11 GB and 41 GB in size . In this paper , we develop a fast , high-precision non-linear quantization method – block-wise dynamic quantization – that enables stable 8-bit optimizers ( e.g. , Adam , AdamW , and Momentum ) which maintain 32-bit performance at a fraction of the memory footprint and without any changes to the original hyperparameters.1 While most current work uses 32-bit optimizer states , recent high-profile efforts to use 16-bit optimizers report difficultly for large models with more than 1B parameters ( Ramesh et al. , 2021 ) . Going from 16-bit optimizers to 8-bit optimizers reduces the range of possible values from 216 = 65536 values to just 28 = 256 . To our knowledge , this has not been attempted before . Effectively using this very limited range is challenging for three reasons : quantization accuracy , computational efficiency , and large-scale stability . To maintain accuracy , it is critical to introduce some form of non-linear quantization to reduce errors for both common small magnitude values and rare large ones . However , to be practical , 8-bit optimizers need to be fast enough to not slow 1We study 8-bit optimization with current best practice model and gradient representations ( typically 16-bit mixed precision ) , to isolate optimization challenges . Future work could explore further compressing all three . down training , which is especially difficult for non-linear methods that require more complex data structures to maintain the quantization buckets . Finally , to maintain stability with huge models beyond 1B parameters , a quantization method needs to not only have a good mean error but excellent worse case performance since a single large quantization error can cause the entire training run to diverge . We introduce a new block-wise quantization approach that addresses all three of these challenges . Block-wise quantization splits input tensors into blocks and performs quantization on each block independently . This block-wise division reduces the effect of outliers on the quantization process since they are isolated to particular blocks , thereby improving stability and performance , especially for large-scale models . Block-wise processing also allows for high optimizer throughput since each normalization can be computed independently in each core . This contrasts with tensor-wide normalization , which requires slow cross-core synchronization that is highly dependent on task-core scheduling . We combine block-wise quantization with two novel methods for stable , high-performance 8-bit optimizers : dynamic quantization and a stable embedding layer . Dynamic quantization is an extension of dynamic tree quantization for unsigned input data . The stable embedding layer is a variation of a standard word embedding layer that supports more aggressive quantization by normalizing the highly non-uniform distribution of inputs to avoid extreme gradient variation . Our 8-bit optimizers maintain 32-bit performance at a fraction of the original memory footprint . We show this for a broad range of tasks : 1.5B and 355M parameter language modeling , GLUE finetuning , ImageNet classification , WMT ’ 14+WMT ’ 16 machine translation , MoCo v2 contrastive image pretraining+finetuning , and RoBERTa pretraining . We also report additional ablations and sensitivity analysis showing that all components – block-wise quantization , dynamic quantization , and stable embedding layer – are crucial for these results and that 8-bit Adam can be used as a simple drop-in replacement for 32-bit Adam , with no hyperparameter changes . We open-source our custom CUDA kernels and provide a PyTorch implementation that enables 8-bit optimization by changing two lines of code . 1 BACKGROUND . 1.1 STATEFUL OPTIMIZERS . An optimizer updates the parameters w of a neural network by using the gradient of the loss with respect to the weight gt = ∂L∂w at update iteration t. Stateful optimizers compute statistics of the gradient with respect to each parameter over time for accelerated optimization . Two of the most commonly used stateful optimizers are Adam ( Kingma and Ba , 2014 ) , and SGD with momentum ( Qian , 1999 ) – or Momentum for short . Without damping and scaling constants , the update rules of these optimizers are given by : Momentum ( gt , wt−1 , mt−1 ) = m0 = g0 Initialization mt = β1mt−1 + gt State 1 update wt = wt−1 − α ·mt Weight update ( 1 ) Adam ( gt , wt−1 , mt−1 , rt−1 ) = r0 = m0 = 0 Initialization mt = β1mt−1 + ( 1− β1 ) gt State 1 update rt = β2rt−1 + ( 1− β2 ) g2t State 2 update wt = wt−1 − α · mt√rt+ Weight update , ( 2 ) where β1 and β2 are smoothing constants , is a small constant , and α is the learning rate . For 32-bit states , Momentum and Adam consume 4 and 8 bytes per parameter . That is 4 GB and 8 GB for a 1B parameter model . Our 8-bit non-linear quantization reduces these costs to 1 GB and 2 GB . 1.2 NON-LINEAR QUANTIZATION . Quantization compresses numeric representations to save space at the cost of precision . Quantization is the mapping of a k-bit integer to a real element in D , that is , Qmap : [ 0 , 2k − 1 ] 7→ D. For example , the IEEE 32-bit floating point data type maps the indices 0 ... 232 − 1 to the domain [ -3.4e38 , +3.4e38 ] . We use the following notation : Qmap ( i ) = Qmapi = qi , for example Qmap ( 231 + 131072 ) = 2.03125 , for the IEEE 32-bit floating point data type . To perform general quantization from one data type into another we require three steps . ( 1 ) Compute a normalization constant N that transforms the input tensor T into the range of the domain D of the target quantization data type Qmap , ( 2 ) for each element of T/N find the closest corresponding value qi in the domain D , ( 3 ) store the index i corresponding to qi in the quantized output tensor TQ . To receive the dequantized tensor TD we look up the index and denormalize : TDi = Q map ( TQi ) ·N . To perform this procedure for dynamic quantization we first normalize into the range [ -1 , 1 ] through division by the absolute maximum value : N = max ( |T| ) . Then we find the closest values via a binary search : TQi = 2n argmin j=0 |Qmapj − Ti N | ( 3 ) 1.3 DYNAMIC TREE QUANTIZATION Figure 2 : Dynamic tree quantization . Dynamic Tree quantization ( Dettmers , 2016 ) is a method that yields low quantization error for both small and large magnitude values . Unlike data types with fixed exponent and fraction , dynamic tree quantization uses a datatype with a dynamic exponent and fraction that can change with each number . It is made up of four parts , as seen in Figure 2 : ( 1 ) The first bit of the data type is reserved for a sign . ( 2 ) The number of subsequent zero bits indicates the magnitude of the exponent . ( 3 ) The first bit that is set to one indicates that all following values are reserved for ( 4 ) linear quantization . By moving the indicator bit , num- bers can have a large exponent 10−7 or precision as high as 1/63 . Compared to linear quantization , dynamic tree quantization has better absolute and relative quantization errors for non-uniform distributions . Dynamic tree quantization is strictly defined to quantize numbers in the range [ -1.0 , 1.0 ] , which is ensured by performing tensor-level absolute max normalization . 2 8-BIT OPTIMIZERS . Our 8-bit optimizers have three components : ( 1 ) block-wise quantization that isolates outliers and distributes the error more equally over all bits ; ( 2 ) dynamic quantization , which quantizes both small and large values with high precision ; and ( 3 ) a stable embedding layer to improve stability during optimization for models with word embeddings . With these components , performing an optimizer update with 8-bit states is straightforward . We dequantize the 8-bit optimizer states to 32-bit , perform the update , and then quantize the states back to 8-bit for storage . We do this 8-bit to 32-bit conversion element-by-element in registers , which means no slow copies to GPU memory or additional temporary memory are needed to perform quantization and dequantization . For GPUs , this makes 8-bit optimizers faster than regular 32-bit optimizers , as we show in Section 3 . 2.1 BLOCK-WISE QUANTIZATION . Our block-wise quantization reduces the cost of computing normalization and improves quantization precision by isolating outliers . In order to dynamically quantize a tensor , as defined in Section 1.2 , we need to normalize the tensor into the range [ -1 , 1 ] . Such normalization requires a reduction over the entire tensor , which entails multiple synchronizations across GPU cores . Block-wise dynamic quantization reduces this cost by chunking an input tensor into small blocks of size B = 2048 and performing normalization independently in each core across this block . More formally , using the notation introduced in Section 1.2 , in block-wise quantization , we treat T as a one-dimensional sequence of elements that we chunk in blocks of size B . This means for an input tensor T with n elements we have n/B blocks . We proceed to compute a normalization constant for each block : Nb = max ( |Tb| ) , where b is the index of the block 0 .. n/B . With this block-wise normalization constant , each block can be quantized independently : TQbi = 2n argmin j=0 |Qmapj − Tbi Nb | ∣∣∣∣ 0 < i < B ( 4 ) This approach has several advantages , both for stability and efficiency . First , each block normalization can be computed independently . Thus no synchronization between cores is required , and throughput is enhanced . Secondly , it is also much more robust to outliers in the input tensor . For example , to contrast blockwise and regular quantization , if we create an input tensor with one million elements sampled from the standard normal distribution , we expect less than 1 % of elements of the tensor will be in the range [ 3 , +∞ ) . However , since we normalize the input tensor into the range [ -1,1 ] this means the maximum values of the distribution determine the range of quantization buckets . This means if the input tensor contains an outlier with magnitude 5 , the quantization buckets reserved for numbers between 3 and 5 will mostly go unused since less than 1 % of numbers are in this range . With blockwise quantization , the effect of outliers is limited to a single block . As such , most bits are used effectively in other blocks . Furthermore , because outliers represent the absolute maximum value in the input tensor , blockwise quantization approximates outlier values without any error . This guarantees that the largest optimizer states , arguably the most important , will always be quantized with full precision . This property makes block-wise dynamic quantization both robust and precise and is essential for good training performance in practice .

Summary: The paper shows a working implementation of 8bit states for momentum and second momentum. This is achieved by using block-wise dynamic quantization for efficient compression and fast implementation. They show drastic improvement over f32 diagonal optimizers (mainly Adam) and improvement over a subset of previously proposed sub-linear memory optimizer (AdaFactor f32 variant).