Title: Memory Efficient Optimizers with 4-bit States
Venue: NeurIPS 2023
Accepted: yes
Keywords: memory efficiency

optimizer

Adam

quantization
Abstract: Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency.

Memory Efficient Optimizers with 4-bit States

Abstract

Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency. *

1 Introduction

Large-scale models with a massive amount of parameters [5, 9, 20, 22, 49, 58] have shown impressive few-shot learning abilities on general tasks [52]. Despite being powerful, training these models is challenging. Memory capacity is one of the main bottlenecks of training large-scale models. Modern neural networks are typically trained with stateful optimizers such as Adam [26], which need to maintain one or two optimizer states (i.e., first and second moments) per each parameter. As the model size grows, the memory consumed by optimizer states can be a dominating factor of memory consumption [40, 41].

There are several attempts to reduce optimizers' memory consumption. Factorization [2, 8, 46] applies low-rank approximation to optimizer states, delta tuning [18, 23, 24, 27, 28] avoids maintaining optimizer states for most parameters by only tuning a small subset, and low-precision optimizers [15, 44] represent their states with low-precision numerical formats, which consume less memory.

Among these methods, low-precision optimizers are attractive due to their simplicity and wide applicability. Dettmers et al. [15] propose an 8-bit optimizer with reparameterized embedding layers ('stable embedding layers') and a block-wise 8-bit dynamic exponential numerical format for optimizer states. Their 8-bit optimizers achieve similar convergence to full-precision optimizers on language modeling, image classification, machine translation, and language understanding tasks.

In this work, we push the required numerical precision of low-precision optimizers from 8 to 4-bit through analyzing the patterns in the first and second moments and designing dedicated quantizers

Algorithm 1 Compression-based Memory Efficient Optimization Framework

Require: black-box stochastic optimization algorithm A , initial parameter w 0 ∈ R p , initial optimizer state ¯ s 0 = 0 , total number of iterations T

1: for t = 1 , 2 , . . . , T do

2: Sample a minibatch ζ t and get stochastic gradient g t = ∇ w f ( w t -1 , ζ t )

3: s t -1 ← decompress (¯ s t -1 )

4: w t , s t ←A ( w t -1 , s t -1 , g t )

5: ¯ s t ← compress ( s t )

6: end for

7: return w T

for optimizer states. Specifically, we find that moments exhibit complicated outlier patterns, that vary across different parameter tensors. The large block size proposed by Dettmers et al. [15] cannot properly handle all different outlier patterns. Based on this observation, we propose to use a smaller block size, which improves the approximation of first moment.

For the second moment, we find its quantization suffers from a zero-point problem . Since the update direction is usually inversely proportional to the square root of the second moment, quantizing non-zero quantities to zero will cause significant deviation. To address this problem, we propose a simple linear quantizer to exclude the zero-point for second moment. We further propose a stronger quantization technique, rank-1 normalization , to improve the approximation of second moment by better handling the outlier patterns. Our proposed quantization techniques are robust enough to achieve lossless convergence under 4-bit, even without the stable embedding layers proposed by Dettmers et al. [15].

Finally, we investigate the combination of factorization methods [2, 46] with low-precision optimizers, and propose a memory efficient optimizer which utilizes quantized 4-bit first moment and factorized second moment. For applicable tasks, the hybrid optimizer enjoys best of both worlds: good convergence and memory efficiency.

We evaluate our 4-bit optimizers on a wide range of tasks, including natural language understanding, machine translation, image classification, and instruction tuning of large language models. On all the benchmarks, our 4-bit optimizer can converge similarly fast with full-precision optimizers, and the converged models do not have noticeable accuracy loss. Our optimizers consumes less memory than existing 8-bit optimizers [15], while improves the throughput of language model finetuning with optimizer offloading [41, 45] due to reduced communication cost.

2 Preliminaries

In this section, we present some preliminaries of compression-based memory efficient optimizers and discuss quantization methods for compression of optimizer states in a general formulation.

2.1 A Framework for Compression-based Memory Efficient Optimizers

Gradient-based stateful optimizers like SGDM [38, 47], Adam [26], AdamW [32] are the principal choices in deep neural network training. However, the memory footprint of stateful optimizers is several times of model itself, which results in a bottleneck for large model pretraining/finetuning. Consider the update rule of the Adam optimizer:

During the training process, the model parameters w t and optimizer states (i.e., first and second moments) m t , v t need to be stored persistently in the GPU memory. As the model grows large, optimizer states are a main source of training memory consumption. Each parameter dimension takes

two 32-bit optimizer states, making the memory consumption of Adam-style optimizers three times larger than stateless optimizers like SGD.

Compressing the optimizer states is a promising method to achieve memory efficient optimization. Formally, given a gradient-based optimizer A , a memory efficient version of the optimization algorithm with compressed optimizer states is given by Alg. 1. The algorithm A can be any gradientbased optimizers like SGDM, Adam, AdamW, etc. See App. F for examples. In Alg. 1, the precise states s t are only temporal, and only compressed states ¯ s t are stored persistently in the GPU memory. Memory footprint reduction can be achieved since in neural networks, the state vectors m t and v t are usually concatenation of state vectors of each parameterized layer. Therefore, we can perform the optimizer steps (Line 3-5) separately for each layer, so only one single layer of the precise states are presented in the memory at a time. The states for all other layers are kept compressed. In the rest of this paper, we focus on the compression and decompression method, to achieve high compression rate of optimizer states while maintaining good convergence of the optimizer.

2.2 Main Compression Method: Quantization

Quantizing the optimizer states to lower precision (e.g. 8-bit integers) is an effective way to compression optimizer states [15]. In this case, the optimizer states are compressed with a quantizer and decompressed with a dequantizer. The low-precision numerical format significantly impacts the accuracy of quantization methods. Here, we present a general framework for numerical formats we considered for compressing optimizer states.

A quantizer converts full-precision tensors to low-precision formats. Based on the formulation proposed by Dettmers and Zettlemoyer [16], we disentangle the quantizer Q ( · ) into two parts: normalization N ( · ) and mapping M ( · ) , which applies sequentially and element-wisely to a tensor to be quantized. For concise presentation, we only discuss quantizers for unsigned inputs, and defer the discussion of signed inputs in App. E.1. Formally, the quantizer for a tensor x ∈ R p is given by

Normalization The normalization operator N scales each elements of x into the unit interval, i.e. [0 , 1] . Normalization can have different granularity, such as per-tensor, per-token (row) [37, 55], per-channel (column) [4], group-wise [37, 55] and block-wise [15]. The per-tensor and block-wise normalization operators are given by

respectively, where the involved scaling factors are called quantization scale , which are persistently stored together with quantized tensor until dequantization. The granularity of normalization presents a trade-off of quantization error and memory overhead. Normalization method with low quantization error and acceptable memory overhead is preferred. In this case, the coarsest per-tensor normalization operator has negligible memory overhead, i.e. only 1 scaling factor regardless of tensor size, but suffers from largest error due to outliers. Block-wise normalization views the tensor as an 1dimensional array, divides the array into blocks of size B called block and assigns a quantization scale within each block, which leads to glyph[ceilingleft] p/B glyph[ceilingright] quantizaton scales totally. The block size could be adapted to trade-off quantization error and memory overhead.

Mapping A mapping [15] converts normalized quantities to low-bitwidth integers. Formally, the mapping operator M = M T ,b is equipped with a bitwidth b and a predefined increasing mapping, named quantization mapping T : [0 , 2 b -1] ∩ Z → [0 , 1] . Then M is defined as

The design of T is critical as it could effectively mitigate quantization error by capturing the distribution information of n . There are two kinds mappings that are of specific interest to optimizer states quantization, linear mapping and dynamic exponent (DE) mapping [13]. A linear mapping T ( i ) = ( i + 1) / 2 b defines a linear quantizer, where the quantization intevals distribute evenly in each block. The DE mapping can approximate small values well, similar to floating point numbers. DE splits the binary representation of a low-precision integer i into three parts: a leading

10

3

10

4

0

10

4

10

3

0

5000

10000

15000

20000

10

3

10

4

0

10

4

10

3

0

10000

20000

30000

40000

50000

10

3

10

4

0

10

4

10

3

0

20000

40000

60000

80000

100000

dim0

0

1000

2000

3000

dim1

0

200

400

600

800

0.2

0.4

0.6

0.8

1.0

(a)

(b)

(c)

(d)

sequence of E zeros, followed by an indicator bit one, and remaining F fraction bits. DE defines T ( i ) = 10 -E ( i ) fraction ( i ) , where fraction ( i ) ∈ [0 . 1 , 1] . See App. E.2 for full specifications and visualizations of different quantization mappings.

The normalization N and mapping M roughly play the same role in finding good quantization point candidates and they affect each other. If an oracle normalization scaling the original tensor x to a uniform distribution is accessible, linear mapping could be used generally. On the contrary, if the optimal mapping could be readily identified with respect to certain metrics for a per-tensor normalized tensor, there is no necessity to develop additional normalization methods that incur extra overhead for mitigating the negative effects of outliers. In essence, if one of the two operators approaches optimality, the quantizer can still perform well even when the other operator is set to its most basic configuration. Additionally, as one operator becomes increasingly powerful, the potential for improvement in the other operator gradually diminishes.

Dequantization The dequantizer is just the inverse of the quantizer, which is simply

Based on our formulation, we name quantizers by their normalization and mapping methods as Norm./Map. . For example, 8-bit optimizers [15] use block-wise normalization with a block size of 2048 and dynamic exponent mapping, which we call B2048/DE in the rest of the paper.

3 Compressing First Moment

In the next two sections, we describe our design of compression and decompression methods in Alg. 1 to realize our memory efficient 4-bit optimizers. We first discuss the compression of the first moment. Our compression method for first moment is based on Dettmers et al. [15], which uses block-wise normalization with a block size of 2048 and dynamic exponent mapping [13]. We preliminary reduce the bitwidth from 8-bit to 4-bit and discover that the first moment is rather robust to quantization. This simple optimizer can already converge with 4-bit first moment, though in some cases there is accuracy degradation.

To further improve the performance, we investigate the patterns of the first moment. There are some outliers in the moments and outliers significantly affect quantization scale due to their large magnitude. Outlier patterns have been studied for weights and activations. It is shown that the weights are rather smooth [14, 54, 57], while the activation have column-wise outliers [4, 14, 53, 54], i.e., the outliers in activation always lie in some fixed channels. However, we find that the outlier patterns in optimizer states are quite complicated. Specifically, Fig. 2 shows patterns of first moment in a Swin-T model during training. Outliers of the layers.3.blocks.0.mlp.fc1 layer lie in fixed rows while outliers of the layers.3.blocks.0.mlp.fc2 layer lie in fixed columns. Actually, the outlier patterns vary across different architectures, layers and parameters. See more patterns in App. B.

The complicated outlier pattern makes optimizer states harder to quantize. There exist some moment tensors, that the outliers persist roughly in certain columns (dimension 1), as shown in Fig. 1. Block-wise normalization treats this tensor as an flattened one-dimensional sequence, in the row-first order. Therefore, any block size of 2048 would include an entry in the outlier column, resulting

dim0

0

250

500

750

dim1

0

1000

2000

3000

0.2

0.4

0.6

0.8

1.0

dim0

0

1000

2000

3000

dim1

0

250

500

750

0.2

0.4

0.6

0.8

1.0

(a)

(b)

2

3

4

5

6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

2

3

4

5

6

0.0

2.5

5.0

7.5

10.0

12.5

15.0

2

3

4

5

6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

(a)

(b)

(c)

in a large quantization scale. In this case, block-wise normalization is not better than per-tensor normalization. Therefore, we adopt a smaller block size of 128, as it provides enhanced performance while incurring only a little memory overhead. In Fig. 1, we show that quantizing with a smaller block size approximates the moments better than a large block size.

4 Compressing Second Moment

Compared to the first moment, quantizing the second moment is more difficult and incurs training instability. Besides its sharper outlier pattern and ill-conditioned distribution compared with first moment, we further identify the zero-point problem as the main bottleneck of quantizing the second moment. We propose an improved normalization method with a quantization mapping excluding zero. We also propose a factorization method for compressing the second moment, which leads to even better memory efficiency.

4.1 Zero-point Problem

The problem of quantizing second moment different from quantizing other tensors in neural networks. For weights [25], activations [25], and gradients [1], it is desirable to contain zero in the quantization mapping for lower approximation error and training stability. Empirically, zero is often the most frequent element [57]. But for second moment in Adam, small values around zero significantly impact the update direction, which is proportional to the inverse square root of the second moment, as shown in Eq. 1. In this case, a small quantization error of values around zero will cause catastrophically large deviation of the update direction. Fig. 3 shows the histogram of the inverse square root (i.e., transformed with h ( v ) = 1 / ( √ v +10 -6 ) ) of the second moment quantized with B128/DE . The quantizer pushes most of entries of the tensor to zero, so the inverse square root of most points fall into 10 6 due to zero-point problem, resulting in a complete degradation in approximation. One remedy is to simply remove zero from the DE quantization map, which we call DE-0 . The smallest number representable by DE-0 is 0.0033. With the B2048/DE-0 quantizer applied to second moment, the approximation of second moment is more precise (Fig. 3), and the training with 4-bit optimizer stabilizes (Tab. 1). However, by removing zero, DE-0 wastes one of the 2 4 = 16 quantization points. We propose to use a Linear mapping T ( i ) = ( i +1) / 2 b , whose smallest representable number is 0 . 0625 . The linear mapping performs better than DE-0 for quantizing second moment.

In Tab. 1, we ablate different quantization schemes and show that excluding zero from the mapping is indeed the key factor for second moment quantization, which cannot be replaced by a smaller block size or/and stochastic rounding [11]. We also test Stable Embedding layers proposed by Dettmers et al. [15], which are reparameterized embedding layers which can be more stably optimized. While Stable Embedding could improve training stability, it cannot fully retain accuracy since nonembedding layers still suffer from zero points. On the other hand, when the zero-point problem is properly addressed, Stable Embedding is no longer required to retain accuracy.

4.2 Rank-1 Normalization

We propose an empirically stronger normalization method named rank-1 normalization based on the observation of the heterogeneous outlier pattern in Sec. 3, inspired by the SM3 optimizer [2]. Formally, for a matrix-shaped non-negative (second moment) tensor x ∈ R n × m , its 1-dimensional

statistics r ∈ R n and c ∈ R m are defined as r i = max 1 ≤ j ≤ m x i,j and c j = max 1 ≤ i ≤ n x i,j , which are exactly the quantization scales of per-row and per-column normalization. Rank-1 normalization utilizes two scales jointly and produce a tighter bound for entry, which is defined as

Rank-1 normalization could be easily generalized to high-dimensional tensors and signed tensors. See App. G for details and the pseudocode.

Compared with per-tensor, per-token (row), and per-channel (column) normalization, rank-1 normalization utilizes the 1-dimensional information in a more fine-grained manner and gives element-wisely unique quantizaion scales. It deals with the outliers more smartly and effectively when outlier persists in fixed rows or columns but the pattern across tensors are unknown and/or varied (Fig. 2). On the other hand, block-wise normalization is also capable to capture the local information and avoid outliers effectively regardless of the patterns when a small block size is taken, but rank-1 normalization provides a better trade-off between memory overhead and approximation. Rank-1 normalization falls back to per-tensor normalization for 1-dimensional tensors, so we use B128 normalization in those cases. Empirically, rank-1 normalization match or exceed the performance of B128 normalization at a moderate model scale (hidden_size=1024). It is also possible to combine block-wise and rank-1 normalization together, which we leave for future work.

4.3 Factorization

Many memory efficient optimization methods [2, 8, 46] represent the entire second moment with a few number of statistics, which differ from quantization and they take only sublinear memory cost. These works bring more memory saving but they are only applicable to the second moment. In this work, we find the factorization method proposed in Adafactor [46] could also avoid zero-point problem effectively, as shown in Tab. 1. Further, we explore the effects of factorization on second moment based on quantized optimizers to attain maximal memory saving while maintain lossless accuracy. Specifically, when factorization is enabled, we factorize all second moment with dimension greater than 1, and quantize 1d second moment tensors.

5 Experiments

We compare our 4-bit optimizers with their full-precision counterparts, as well as other memory efficient optimizers including 8-bit AdamW [15] † , Adafactor [46] and SM3 [2]. 8-bit AdamW's optimizer states are not quantized for embedding layers. For Adafactor, we compare both the

β 1 > 0 and the β 1 = 0 (no first moment) configuration. For our 4-bit optimizers, we report two versions both based on 32-bit AdamW: (1) '4-bit AdamW' quantizes first moment with B128/DE and second moment with Rank-1/Linear . (2) the more memory efficient '4-bit Factor' quantizes first moment with B128/DE , factorizes second moment when the dimension of tensor is greater than 1, and quantizes lefted 1-dimensional second moment with Rank-1/Linear . See App. D for details about the experiments.

Models, datasets and hyperparameters We report performance metrics on standard benchmarks, including image classification (CLS) with Swin-T [31] ‡ on ImageNet-1k [12], natural language understanding (NLU) by fine-tuning RoBERTa-L [30] § fine-tuning on GLUE [51], question answering (QA) by fine-tuning RoBERTa-L on SQuAD [42, 43], natural language generation (NLG) by finetuning GPT-2 Medium [39] ¶ on E2E-NLG [35], machine translation (MT) by training TransformerBase [50] || on WMT14 en-de [3] and LLaMA [49] fine-tuning. We fine-tune LLaMA-7B, LLaMA13B and LLaMA-33B [49] on the Alpaca dataset [48] ** and evaluate them on MMLU [21] and standard common sense reasoning benchmarks: HellaSwag [56], ARC easy and challenge [10] and OpenBookQA [33].

We mainly follow the hyperparameters in the original paper or/and codebase. In each benchmark, we keep same hyperparameters in one optimizer on different quantization schemes, which gives an out-of-box transfer from full-precision optimizer to low-bit optimizer without extra hyperparameter tuning. See App. D for hyperparameters and training details.

Accuracy of 4-bit Optimizers We first check whether our memory efficient 4-bit optimizers could retain accuracy. According to Tab. 2, our 4-bit optimizers can match or exceed 32-bit AdamW performance on all fine-tuning tasks (NLU, QA, and NLG) and are comparable on all pretraining tasks (CLS and MT). Sublinear memory optimizers Adafactor ( β 1 = 0 ) and SM3 could have better memory efficiency, but they suffer from performance degradation, particularly on the CLS task. According to Tab. 3, our 4-bit AdamW will not destroy the capability of pretrained models while enabling them to obtain instruction-following ability. 4-bit AdamW is comparable with 32-bit AdamW on all tasks and does not get worse when the model size grows. Moreover, their convergence curves closely align (Fig. 4).

Memory and Computing Efficiency We evaluate the memory and computation efficiency of proposed 4-bit optimizers on instruction tuning, NLU, and NLG tasks, in Tab. 4. Our 4-bit optimizer offers more memory saving compared to 8-bit optimizers, reducing the training memory consumption by up to 57.7%. It may look like the memory saving saturates when the bitwidth goes down. This is because we report the total memory consumption (including data, activations, and memory fragments)

rather than the optimizer memory consumption alone. In principle, the optimizer states is 2x smaller for 4-bit AdamW than 8-bit AdamW, and about 4x smaller for 4-bit Factor.

The instruction tuning task uses two Nvidia A100 80GB GPUs, while the model is sharded across GPUs with PyTorch's FSDP. In this case, our 4-bit optimizer speeds up training due to reduced communication cost. For the smaller RoBERTa-L and GPT-2 Medium, our 4-bit optimizers appear to be slower than 8-bit AdamW. This is because we have not yet optimize our implementation with fused operators. The speed of our 4-bit optimizers should match or surpass 8-bit optimizers after operator fusion.

We report the largest OPT and LLaMA models trainable under a given memory budget with fullprecision and our 4-bit optimizers in Tab. 5. Our optimizer allows for the training of 4x large OPT models, and enables the training of LLaMA-7B model with a single 80GB GPU.

0

200

400

600

800

1000

1200

Training Step

0.4

0.6

0.8

1.0

1.2

1.4

Training Loss

32-bit AdamW

4-bit AdamW

Ablation Study Finally, we investigate the effectiveness of our proposed quantization schemes and the sensitivity of each moment to quantization in Tab. 6. We see that quantizing both the first and second moment brings a marginal drop in accuracy. Smaller block size on first moment improves accuracy and takes a step towards lossless performance. Factorizing second moment improves accuracy while leads to better memory efficiency.

6 Related Work

Compression-based memory efficient optimizers There have been some works trying to approximate the gradient statistics with sublinear memory cost relative to the number of parameters. Adafactor [46] achieves memory reduction by approximating the second-moment of matrix-shaped

Largest fine-tunable Model

24 GB

80 GB

80 GB

32-bit AdamW

4-bit AdamW

OPT-350M

OPT-1.3B

-

OPT-1.3B

OPT-6.7B

LLaMA-7B

parameters with the outer product of 'row' and 'column'. SM3 [2] considers the cover of parameters and maintain one statistics for each element in the cover. Experimentally, cover composed of slices of co-dimension 1 for each tensor has been adopted. Extreme tensoring [8], compared with SM3, has a different motivation but similar formulation for the experimental choice. These memory efficient methods only focus on memory reduction on second moment and are applicable to most of the second moment based optimizers, like Adagrad [19], Adam [26], AdamW [32], etc.

Another line of work achieves memory reduction by using compression-based method and maintaining coordinate-wise low-bit precision optimizer states. The stability of 16-bit Adam is firstly studied in DALL-E [44]. Dettmers et al. [15] uses block-wise and dynamic exponent quantization to reduce the coordinate-wise optimizer states from 16-bit to 8-bit and is applicable to SGDM and Adam/AdamW. Compression-based methods only reduce memory by a fixed percentage, irrelevant to the number and shape of parameters. Compression-based methods are less memory efficient compared with sublinear memory methods, but exhibit superior performance across a wider range of benchmarks.

Other memory efficient techniques Some works focus on the memory efficiency of activations. Activation compressed training [6, 29] keeps low-precision activations in GPU memory via quantization in forward phase and dequantize the low-precision activations to full-precision layer-wisely in backward phase. Gradient checkpointing [7] only keeps activation in a small number of layers in forward phase and recompute the activations in all layers when gradient computation is needed, which leads a trade-off between storage overhead of activations and additional computation cost. These methods can be combined with our optimizer for better memory efficiency. LoRA [24] freezes the pretrained weights and only tunes new initialized low-rank parameters, which greatly reduce the size of computation graph but is only applicable to language model finetuning.

Sharding [40] divides model parameters and optimizer states among multiple devices, allowing for more efficient model training through the use of additional GPUs. Offloading [41, 45] reduces GPU memory consumption by transferring data to CPU memory. Our optimizer can be utilized by these methods to reduce the communication overhead.

7 Conclusions, Limitations, and Broader Impact

Wepropose compression-based memory efficient 4-bit optimizers using quantization and factorization. We observe the heterogeneous outlier patterns in optimizer states and identify the zero-point problem in quantizing second moment as the main bottleneck. 4-bit optimizers achieve lossless performance in finetuning and comparable accuracy in pretraining on a wide range of tasks.

Limitations The optimal quantization setting probably depends on task, datasets, and training details. While rank-1 normalization and linear mapping for second moment quantization performs consistently well in our experiments, task-specific quantization settings not in the scope of the study might perform better and be helpful to achieve lossless performance. Our evaluation is currently limited to language and vision tasks, while the applicability of our method to reinforcement learning, audios, and graph learning tasks still needs further study.

Broader Impact Our work can facilitate the access to large models for pretraining and finetuning, which were previously constrained by GPU memory limitations. This could help democratizing large models and opens up new avenues for research that were previously unattainable due to restricted GPU memory, especially benefiting researchers with limited resources. However, our work can also exaggerate the abuse large models.

References

Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-efficient sgd via gradient quantization and encoding. Advances in Neural Information Processing Systems , 30, 2017.

Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems , 32, 2019.

Ondˇ rej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, et al. Findings of the 2014 workshop on statistical machine translation. In Proceedings of the ninth workshop on statistical machine translation , pages 12-58, 2014.

Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Understanding and overcoming the challenges of efficient transformer quantization. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 7947-7969, 2021.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems , 2020.

Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Actnn: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning , pages 1803-1813. PMLR, 2021.

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016.

Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. In International Conference on Learning Representations , 2020.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 , 2018.

Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. Binaryconnect: Training deep neural networks with binary weights during propagations. In Advances in Neural Information Processing Systems , pages 3123-3131, 2015.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , 2009.

Tim Dettmers. 8-bit approximations for parallelism in deep learning. arXiv preprint arXiv:1511.04561 , 2015.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022a.

Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations , 2022b.

Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. arXiv preprint arXiv:2212.09720 , 2022c.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT , pages 4171-4186, 2019.

Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al. Delta tuning: A comprehensive study of parameter efficient methods for pre-trained language models. arXiv preprint arXiv:2203.06904 , 2022.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research , 23(120):1-39, 2022.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 , 2020.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack William Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning , pages 2790-2799. PMLR, 2019.

Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2704-2713, 2018.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, 2021.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 , 2021.

Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for general architectures. arXiv preprint arXiv:2206.11357 , 2022.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789 , 2018.

Yurii Nesterov. Introductory lectures on convex optimization: A basic course , volume 87. Springer Science & Business Media, 2003.

Jekaterina Novikova, Ondˇ rej Dušek, and Verena Rieser. The e2e dataset: New challenges for end-to-end generation. arXiv preprint arXiv:1706.09254 , 2017.

Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint arXiv:1904.01038 , 2019.

Gunho Park, Baeseong Park, Se Jung Kwon, Byeongwook Kim, Youngjoo Lee, and Dongsoo Lee. nuqmm: Quantized matmul for efficient inference of large-scale generative language models. arXiv preprint arXiv:2206.09557 , 2022.

Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks , 12(1):145-151, 1999.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-16. IEEE, 2020.

Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zeroinfinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-14, 2021.

Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. arXiv preprint arXiv:1806.03822 , 2018.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250 , 2016.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , pages 8821-8831. PMLR, 2021.

Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. { ZeRO-Offload } : Democratizing { Billion-Scale } model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21) , pages 551-564, 2021.

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018.

Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning , pages 1139-1147. PMLR, 2013.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca , 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018.

Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. Transactions on Machine Learning Research , 2022.

Xiuying Wei, Yunchen Zhang, Xiangguo Zhang, Ruihao Gong, Shanghang Zhang, Qi Zhang, Fengwei Yu, and Xianglong Liu. Outlier suppression: Pushing the limit of low-bit transformer language models. arXiv preprint arXiv:2209.13325 , 2022.

Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438 , 2022.

Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems , 35:27168-27183, 2022.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830 , 2019.

Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130b: An open bilingual pre-trained model. In International Conference on Learning Representations , 2023.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

A Additional Experiment Results

In this section, we show additional experiment results beyond Tab. 2. Tab. 7 shows the results of RoBERTa-L finetuning on each task in GLUE datasets. Tab. 8 shows the results of GPT-2 Medium finetuning on E2E-NLG via different metrics. Tab. 9 shows the EM and F1 of RoBERTa-L finetuning on SQuAD and SQuAD 2.0 datasets.

β

1

= 0

.

B Outlier Patterns of Moments

In this section, we give a comprehensive visualization about the outlier pattern of optimizer states. [2] did similar analysis for Adagrad's second moment but here we give a better demonstration about various patterns in optimizer states. The same technique has been applied to parameters and activations in [54].

The outlier pattern of first moment depends on many factors such as data and training hyperparameters. Here we mainly focus on different transformer models and different layers inside. In one transformer block, there is one Attention module and one MLP module, including 6 main parameter matrices. We do not focus on additional parameters including bias and parameters in layer normalization (LN) since they only account a small portion. We denote the 6 matrices by W Q , W K , W V , W O , W 1 , W 2 , respectively. When the matrices across different layers are involved at the same time, we add a subscript indicating layer index. Note that W has shape R C o × C i in a linear layer, we call the output and input dimension by dim0/row and dim1/column, respectively. The per-channel quantizaition used in other works actually correspond to per-row(dim0) normalization here.

Swin Transformer ImageNet pretraining In Fig. 5,6,7, the magnitude of first moment in transformer blocks at different depths are shown. It can be seen that the 1-dimensional structure in all parameter matrices are vague at the initial layer. At layer 2, the pattern in W O and W 1 , that outliers occur at fixed columns, becomes obvious while other parameter matrices remain noisy. At layer 3, the patterns in W O , W 1 , W 2 are quite obvious. In W O and W 2 , the outliers occur at fixed rows while the pattern in W 1 remain unchanged. The 1-dimensional structure in W Q , W K , W V also seem to appear even though not remarkable. It is notable that the pattern of same parameter matrix at different depths are not necessarily same. See W O in Fig. 6,7.

WQ

WK

W

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

80

80

80

60

60

60

0

20

dimo

40

60

80

0

20

40 dim1

0

20

dimo

40

60

80

0

20

40 dim1

0

20

dimo

40

60

80

0

20

40 dim1

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

400

80

80

300

20

dimo

40

60

80

0

20

40 dim1

60

0

100

dimo

200

300

20

40 dim1

60

0

20

owp

40

60

80

0

100dim1

200

400

WQ

WK

W

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

400

400

400

300

300

300

0

100

200

100

200

dim1

0

100

200

100

200

dim1

0

100

200

100

200

dim1

dimo

300

400

0

dimo

300

400

owp

300

400

0

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

400

400

1500

300

300

1250

0

100

200

100dim1

200

0 25050075009259500

100dim1

200

100

200

250

500 dim1

750

1000

dimo

300

0

dimo

0

dimo

300

0

400

400

WQ

WK

W

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

800

800

800

600

600

600

0

200

400

200

400

dim1

0

200

400

200

400

dim1

0

200

400

200

400

dim1

dimo

600

800

0

dimo

600

800

dimo

600

800

0

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

600

800

600

800

2500

3000

0

200

400

200dim1

400

0 00090095090000

200dim1

400

0

200

400

500

1000dim1

1500

2000

dimo

600

0

dimo

dimo

600

800

800

RoBERTa-Large GLUE finetuning In Fig. 8,9,10,11,12,13, the magnitude of first moment in transformer blocks of RoBERTa-Large at different depths are shown. At layer 0 and layer 1 (initial layers), patterns in W O , W 2 are obvious. At layer 11 and layer 12 (intermediate layers), patterns are all noisy. At layer 22 and layer 23 (last layers), patterns in W Q , W K are obvious. Patterns in other matrices are weak.

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0

200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

0.0

0.0

1000

1000

4000

800

800

3000

0200

400

dimo

600

800

1000

0

200

600

0

1000

dimo

2000

3000

4000

200

600

0

200

400

owp

600

800

1000

0

1000

2000

dim1

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

200

400dim1

600

0

200

400dim1

600

。

200

400dim1

600

400

dimo

600

800

1000

0

200

400

dimo

600

800

1000

0

200

400

dimo

600

800

1000

0

200

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

0.0

0.0

0200

400

dimo

600

800

1000

200

400dim1

600

800

1000

0

1000

dimo

2000

3000

4000

0

200

400dim1

600

800

1000

200

400

dimo

600

800

1000

1000

2000

3000

dim1

4000

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

0.0

1000

1000

4000

0200

400

600

200

400dim1

600

800

0

1000

2000

200

400dim1

600

800

200

400

600

1000

2000

dim1

3000

dimo

800

1000

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

0200

400

600

200

400dim1

600

800

0

1000

2000

200

400dim1

600

800

200

400

600

1000

2000

dim1

3000

dimo

800

1000

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

dimo

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

200

400dim1

600

0

1000

2000

200

400dim1

600

600

1000

2000

dim1

dimo

800

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

W

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

8000

1000

0

200

400dim1

600

0

200

400

dimo

600

800

1000

0

200

400dim1

600

0200 400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

dimo

800

1000

0

200

400dim1

600

0

1000

dimo

2000

3000

4000

0

200

400dim1

600

200

400

dimo

600

800

1000

0

1000

2000

dim1

GPT-2 Medium E2E-NLG finetuning In Fig. 14,15,16,17,18,19, the magnitude of first moment in transformer blocks of GPT-2 Medium at different depths are shown. At layer 1 and layer 2 (initial layers), patterns in W O are obvious. At layer 13 and layer 14 (intermediate layers), patterns in W K , W O are obvious. At layer 21 and layer 22 (last layers), patterns in W Q , W K , W V , W O are obvious. First moment of W 1 , W 2 are consistently noisy throughout layers. It is notable that the rows(or columns) that gather outliers are different across different layers.

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

400

dimo

600

800

1000

0

200

400dim1

600

200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

dimo

800

0

200

400dim1

600

0

1000

dimo

2000

3000

4000

0

200

400dim1

600

0

200

400

dimo

600

800

1000

0

1000

2000

dim1

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0

200

400

dimo

600

800

1000

0

200

400dim1

600

0

200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

dimo

600

800

1000

0

200

400dim1

600

0

1000

dimo

2000

3000

4000

200

400dim1

600

0

200

dimo

600

800

1000

0

1000

2000

dim1

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800，

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

600

200

400dim1

600

0

1000

2000

200

400dim1

600

200

400

600

1000

2000

dim1

dimo

800

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

600

200

400dim1

600

0

1000

2000

200

400dim1

600

0

200

400

600

1000

2000

dim1

dimo

800

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

WV

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

800

800

3000

0200

400

600

200

400dim1

600

0

1000

2000

200

400dim1

600

0

200

400

600

1000

2000

dim1

dimo

800

1000

0

dimo

3000

4000

0

dimo

800

1000

0

WQ

WK

WV

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

1000

800

800

800

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

0200

400

dimo

600

800

1000

0

200

400dim1

600

wo

W1

1.0

1.0

1.0

0.8

0.8

0.8

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.2

1000

1000

4000

0200

400

600

200

400dim1

600

800

0

1000

2000

200

400dim1

600

800

200

600

1000

2000

3000

dim1

dimo

800

0

dimo

3000

4000

0

dimo

800

1000

0

C Quantization Quality via Histogram

C.1 Zero-point Problem

In Fig. 20,21,22, we show the effect of zero-point on quantization error for second moment via histogram. All those figures show the negative impact of zero-point on quantizing second moment. After removing zero-point, the quantization quality improves at a great scale.

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

0

2

4

6

8

10

12

14

16

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

10

1

10 0

10 1

10 0

10 1

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

1

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

10 2

1

2

3

4

5

6

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

10

2

10

1

10 0

10 1

10 2

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

1.2

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

3

4

5

6

7

8

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

10 1

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

2

3

4

5

6

7

8

0.0

0.5

1.0

1.5

2.0

2.5

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

10

2

10

1

10 0

10 1

10 0

10 1

10 2

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

2

3

4

5

6

7

8

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2

3

4

5

6

7

8

C.2 Comparison between Block-wise and Rank-1 Normalization

To show the differences in quantization error for second moment between block-wise normalization and rank-1 normalization, some cases where rank-1 normalization approximates better than blockwise normalization are shown in Fig. 23, 25, 27. Also, some cases where rank-1 normalization approximates worse than block-wise normalization is shown in Fig. 24, 26, 28. Empirically, it has been observed that rank-1 normalization yields superior results when the distribution exhibits long-distance multimodal characteristics. On the other hand, block-wise normalization tends to outperform when the distribution displays short-distance multimodal patterns and/or intricate local structures.

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

10 1

10 2

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

10

2

10

1

10 0

10 1

10 2

10 1

10 2

10 3

10 4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

C.3 Effectiveness of Block Size in Block-wise Normalization

In Fig. 29,30,31, we show the effect of block size on quantization error for both first and second moments. Fig. 29,30 shows that B2048 normalization quantizes a significant portion of the points to zero, resulting in poor approximation based on the histogram. However, when we utilize a smaller block size of 128, the quantization performance improves. Fig. 31 shows smaller block size improves quantization quality on second moment.

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

10

2

10

1

10 0

10 1

10 2

10 1

10 2

10 3

10 4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

10

3

10

2

10

1

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

10 1

10 2

10 3

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

10

2

10

1

10 0

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2

3

4

5

6

1

2

3

4

5

0.0

0.2

0.4

0.6

0.8

0.0

0.2

0.4

0.6

0.8

1.0

10

3

10

2

10

1

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

10 2

1

2

3

4

5

0.0

0.2

0.4

0.6

0.8

0.0

0.2

0.4

0.6

0.8

10

3

10

2

10

1

10 0

10 1

10 2

10 0

10 1

10 2

0.0

0.5

1.0

1.5

2.0

1

2

3

4

5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

1

2

3

4

5

2.0

2.5

3.0

3.5

4.0

4.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

2.0

2.5

3.0

3.5

4.0

4.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

2.0

2.5

3.0

3.5

4.0

4.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

10 1

10 2

2.0

2.5

3.0

3.5

4.0

4.5

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

2.4

2.6

2.8

3.0

3.2

3.4

3.6

0.0

0.5

1.0

1.5

2.0

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

2.4

2.6

2.8

3.0

3.2

3.4

3.6

0.0

0.5

1.0

1.5

2.0

2.4

2.6

2.8

3.0

3.2

3.4

3.6

0.0

0.5

1.0

1.5

2.0

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

0.0

0.2

0.4

0.6

0.8

1.0

10

1

10 0

10 1

10 2

2.4

2.6

2.8

3.0

3.2

3.4

3.6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

10

4

10

5

0

10

5

10

4

0

10000

20000

30000

40000

50000

60000

70000

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0

10

20

30

40

50

60

70

80

10

4

10

5

0

10

5

10

4

0

50000

100000

150000

200000

10

4

10

5

0

10

5

10

4

0

10000

20000

30000

40000

50000

60000

70000

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0

2

4

6

8

10

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0

50

100

150

200

10

4

10

5

0

10

5

10

4

0

100000

200000

300000

400000

500000

10

3

10

4

10

5

0

10

5

10

4

10

3

0

10000

20000

30000

40000

50000

60000

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0

5

10

15

20

25

10

3

10

4

10

5

0

10

5

10

4

10

3

0

20000

40000

60000

80000

100000

120000

140000

10

3

10

4

10

5

0

10

5

10

4

10

3

0

10000

20000

30000

40000

50000

60000

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

0

50

100

150

200

10

3

10

4

10

5

0

10

5

10

4

10

3

0

200000

400000

600000

800000

2

3

4

5

6

7

0.0

0.2

0.4

0.6

0.8

1.0

1.2

0.0

0.2

0.4

0.6

0.8

1.0

10

2

10

1

10 0

10 1

10 2

0.0

0.2

0.4

0.6

0.8

1.0

10 0

10 1

10 2

2

3

4

5

6

7

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.2

0.4

0.6

0.8

1.0

1.2

10

1

10 0

10 1

10 2

10 0

10 1

10 2

10 3

0

1

2

3

4

5

6

2

3

4

5

6

7

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

2

3

4

5

6

7

D Experimental Details

D.1 Quantization

There are several parameters in deep neural networks that play a delicate role without occupying too much memory, such as normalization layers and bias. In this context, we establish a rule to determine which parameters should not be quantized. For all the experiments we conducted, the rule is straightforward: tensors with a size smaller than or equal to 4096 will not be quantized. However, for larger models with a hidden size exceeding 4096, it is advisable to exclude the bias and normalization layers from quantization. Regarding the quantization settings, as stated in Sec. 5, we employ block-wise normalization with a block size of 128, dynamic exponent mapping for first moment and rank-1 normalization, linear mapping for second moment. When we apply factorization on second moment, only tensors with a dimension greater than or equal to 2 will be factorized while 1-dimensional tensors that meet the specified rule will still be quantized.

8-bit Adam [15] also uses the threshold of 4096 about size to determine whether or not to quantize parameters. Additionally, the implementation in huggingface does not quantize the parameters in Embedding layers regardless of the model used. Consequently, we compare our method with the 8-bit Adam that does not quantize Embedding .

D.2 Hyperparameters and Training Details

In each benchmark, unless otherwise specified, we maintain the same hyperparameters for a given optimize across different quantization schemes. Additionally, we use same optimizer hyperparameters across various optimizers, including our 4-bit optimizers, 8-bit Adam [15], SM3 [2], Adafactor [46] and the full precision counterpart AdamW [32]. For Adafactor, we use β 1 > 0 as default setting which is same as the β 1 value used in AdamW. Also, the case where β 1 = 0 is compared. The other newly introduced hyperparameters in Adafactor are set to their default values and remain fixed throughout the experiments. For SM3, we compare with the β 1 > 0 configuration, same as the β 1 value used in AdamW.

RoBERTa We train all of our RoBERTa-L models with PyTorch Huggingface †† . On GLUE benchmark, we mainly follow the hyperparameters in fairseq [36]. We use β 1 = 0 . 9 , β 2 = 0 . 98 , glyph[epsilon1] = 1e-6, a weight decay factor of 0.1 and linear learning rate schedule. Other hyperparameters are listed in Tab. 10. On SQuAD benchmark, we mainly follow the reported hyperparameters in RoBERTa paper [30]. We use β 1 = 0 . 9 , β 2 = 0 . 98 , glyph[epsilon1] = 1e-6, a weight decay factor of 0.01 and linear learning rate schedule. The other hyperparameters are listed in Tab. 11. On both datasets, we report the median and standard deviation results over 5 runs, the result in each run is taken from the best epoch. We utilize single RTX 3090 or 4090 GPU for runs of each task in GLUE datasets and four RTX 3090 or 4090 GPUs for SQuAD and SQuAD 2.0.

On SQuAD 2.0, there may be a performance gap observed between the reproduced results using 32-bit AdamW and the original results reported in the original paper. This is because there are some questions without answers in SQuAD 2.0. It is worth noting that the approach employed by Liu et al. [30] to handle unanswered questions may differ from the solutions utilized in the BERT paper [17], which is the reference implementation we are using

GPT-2 We train all of our GPT-2 Medium models with the LoRA codebase ‡‡ . We mainly follow the hyperparameters in [28] and [24]. We use β 1 = 0 . 9 , β 2 = 0 . 999 , glyph[epsilon1] = 1e-6, a weight decay factor of 0.01 and linear learning rate schedule. The other hyperparameters used in GPT-2 are listed in Tab. 12. We report the mean and standard deviation results over 3 runs, the result in each run is taken from the best epoch. We utilize fours RTX 3090 or 4090 GPUs for runs of this task.

Transformer We train all of our Transformer-Base models for machine translation with codebase §§ . We completely follow the hyperparameters in the codebase. We report the mean and standard deviation results over 3 runs, the result in each run is taken from the best epoch. We utilize eight RTX 3090 or 4090 GPUs for runs of this task.

Swin We train all of our Swin-T models with its official codebase ¶¶ . We completely follow the hyperparameters in the codebase. We report the mean and standard deviation results over 3 runs, the result in each run is taken from the best epoch. We utilize eight RTX 3090 or 4090 GPUs for runs of this task.

LLaMA We fine-tune LLaMA-7B, LLaMA-13B and LLaMA-33B with Alpaca codebase *** . We follow the hyperparameters in the codebase for LLaMA-7B and LLaMA-13B, and the hyperparameters of LLaMA-33B are consistent with LLaMA-13B. We fine-tune LLaMA-7B with two A100 80GB GPUs. The training loss curve is the mean results over 3 runs. For LLaMA-7B, we enable Fully Sharded Data Parallelism (FSDP), which packs parameters into 1-dimensional array. This packing process makes it difficult to apply factorization directly without additional engineering efforts. Consequently, we only compare the performance of 4-bit AdamW with its full precision counterpart.

D.3 Memory and Computing Efficiency

In Tab. 4, we present measurements of memory usage in practical settings, i.e. training configuration described in Sec. D.2. Specifically, we measure the memory usage for LLaMA-7B using 2 A100 80G GPUs, RoBERTa-L using 1 RTX 4090 GPU, and GPT-2 Medium using 4 RTX 4090 GPUs. Additionally, the time measurement for RoBERTa-L is conducted on the RTE task.

E Quantization Formulation Details

E.1 Signed Case

In this section, we discuss quantization function for signed tensors and the differences compared to unsigned case. Regarding the normalization operator, the only difference lies in the fact that the sign of the tensor remains unchanged before and after normalization. Formally, let N be the normalization operator for the unsigned cases. For the signed case, the normalization can be defined as

Therefore, the unit interval for signed case is [-1, 1]. Regarding the mapping operator, the difference lies in the values of quantization mappings. See App. E.2 for more details.

E.2 Quantization Mappings

In this work, we mainly consider linear mapping and dynamic exponent mapping [13]. See Fig. 32 for illustration of quantization mappings.

0.0

0.2

0.4

0.6

0.8

1.0

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

linear

dynamic exponent

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

linear

dynamic exponent

Linear mapping It is notable that the linear mapping considered in our work does not include zero in both signed case and unsigned case. Actually, we only use linear mapping in unsigned case, which is defined as torch.linspace(0, 1, (2 ** b) + 1)[1:] .

Dynamic exponent mapping Let b be the total bits. In the main text, we mentioned that dynamic exponent takes the form T ( i ) = 10 -E ( i ) fraction ( i ) . In following paragraphs, we will define the dynamic exponent mapping formally based on the binary representation.

In unsigned case, dynamic exponent mapping [13] is composed of exponent bits E , one indicator bit and fraction bits F , where b = 1 + E + F . It uses the number of leading zero bits E represents the exponent with base 10. The first bit, which is one, serves as an indicator bit that separates the exponent and the unsigned linear fraction. The remaining bits F represent an unsigned linear fraction distributed evenly in (0.1, 1), which is formally defined as

Therefore, a number with E exponent bits and F fraction bits valued k has a value of

For signed case, the only difference is that dynamic exponent mapping additionally uses the first bit as the sign thus we have b = 1 + E +1+ F . Specially, at 8-bit case, we learn from the codebase ††† that dynamic exponent mapping assign 00000000 2 = 0 10 , 00000001 2 = 1 10 in unsigned case and assign 10000000 2 and 00000000 2 with 1 10 and 0 10 , respectively. This means -1 10 is not defined and the mapping is not symmetric in signed case. Finally, after collecting all the represented numbers and arranging them in a sorted, increasing list, which has a length of 2 b , the quantization mapping T ( i ) returns the i -th element of this list.

The construction of dynamic exponent mapping is unrelated to the number of bits. Therefore, when we say we barely turn the 8-bit optimizer into 4-bit optimizer, it just use 4 total bits. The corner cases mentioned in last paragraph remain unchanged.

E.3 Stochastic Rounding

Stochastic rounding is only used in Tab. 1. In this section, we talk about how to integrate stochastic rounding into our formulation of quantization. When stochastic rounding is used, the definition of mapping M has some minor changes. Specifically, M is still an element-wise function and defined as

In other words, M maps each entry n j to the maximal index set M ( n j ) such that for any i ∈ M ( n j ) there is no other 0 ≤ k ≤ 2 b -1 with T ( k ) lying between T ( i ) and n j . Actually, M acts as a filter of T and give a more fine-grained range of quantized output candidates. In this definition, M ( n j ) has only one or two points since only stochastic rounding is considered in the final step.

Finally, we define stochastic rounding R s . When M ( n j ) only has one point, R s just output this point. When M ( n j ) has two points q 1 and q 2 with T ( q 1 ) < n j < T ( q 2 ) , stochastic rounding ( R s ) is defined as

F Compression-based Memory Efficient Optimizer Instances

In this section, we present some examples about compression-based memory efficient optimizers. See Compression-based Memory Efficient SGDM in Alg. 2 and Adam in Alg. 3.

Algorithm 2 Compression-based Memory Efficient SGDM

Require: initial parameter θ 0 ∈ R p , learning rate α , initial first moment ¯ m 0 = 0 , total number of iterations T and momentum parameter β .

1: for t = 1 , 2 , . . . , T do

2: Sample a minibatch ζ t and get stochastic gradient g t = ∇ θ f ( θ t -1 , ζ t )

3: m t -1 ← decompress ( ¯ m t -1 )

4: m t ← β · m t -1 + g t

5: θ t ← θ t -1 -α · m t

6: ¯ m t ← compress ( m t )

7: end for

8: return θ T

Algorithm 3 Compression-based Memory Efficient Adam

Require: initial parameter θ 0 ∈ R p , learning rate α , initial moments ¯ m 0 = 0 , ¯ v 0 = 0 , total number of iterations T and hyperparameters β 1 , β 2 , glyph[epsilon1] .

1: for t = 1 , 2 , . . . , T do

2: Sample a minibatch ζ t and get stochastic gradient g t = ∇ θ f ( θ t -1 , ζ t )

3: m t -1 , v t -1 ← decompress ( ¯ m t -1 ) , decompress (¯ v t -1 )

4: m t ← β 1 · m t -1 +(1 -β 1 ) · g t

6: ˆ m t ← m t / (1 -β t 1 )

5: v t ← β 2 · v t -1 +(1 -β 1 ) · g 2 t

7: ˆ v t ← v t / (1 -β t 2 )

8: θ t ← θ t -1 -α · ˆ m t / ( √ ˆ v t + glyph[epsilon1] )

9: ¯ m t , ¯ v t ← compress ( m t ) , compress ( v t )

10: end for

11: return θ T

G Rank-1 Normalization

In this section, we present the detailed formulation of rank-1 normalization in Alg. 4.

H Theoretical Analysis

The convergence of low-bit optimizers can be guaranteed if their fp32 counterparts converge. Here we provide a theorem about the convergence of quantized SGDM (Alg. 2) under some assumptions. We believe the convergence of low-bit AdamW could be inferred from the convergence of AdamW.

Algorithm 4 Rank-1 Normalization

Require: tensor x ∈ R d 1 ×···× d p ; statistics µ r ∈ R d r for 1 ≤ r ≤ p ; permutation function Φ mapping { 1 , . . . , d } to indices of tensor x , where d = d 1 ×··· × d p .

1: for r = 1 , 2 , . . . , p do

2: for j = 1 , 2 , . . . , d r do

4: end for

5: end for

6: for i = 1 , 2 , . . . , d do

7: M i = min 1 ≤ r ≤ p µ r, Φ( i ) r

8: end for

9: reshape 1-dimensional array M to the same shape as x

10: return x/M

First, we make some assumptions. The first three are rather standard in stochastic optimization literature, while last two depict properties of stochastic quantizers.

(Convexity) The objective function is convex and has an unique global minimum f ( θ ∗ ) .

(Smoothness) The objective f ( θ ) is continuous differentiable and L -smooth;

(Moments of stochastic gradient) The stochastic gradient g is unbiased, i.e., E [ g ( θ )] = ∇ f ( θ ) , and has bounded variance, i.e., E [ ‖ g ( θ ) -∇ f ( θ ) ‖ 2 ] < σ 2 , ∀ θ ∈ R d .

(Unbiased quantizer) ∀ x ∈ R d , E [ Q ( x )] = x .

(Bounded quantization variance) ∀ x ∈ R d , E [ ‖ Q m ( x ) -x ‖ 2 ] ≤ σ 2 m .

Then, we have following theorem:

Theorem 1. Consider the Algorithm 2 with Assumptions 1-5. Let α ∈ (0 , 1 -β L ] , then for all T > 0 we have

where ¯ θ T = 1 T ∑ T -1 i =0 θ i .

H.1 Proof of Theorem 1

To prove Theorem 1, we need some useful lemmas.

Lemma 1. In Algorithm 2, The conditional first and second moments of g t satisfies

Proof. By assumption, we easily have

With Assumption 3, it holds true that

which implies the second part.

Lemma 2. If Assumptions 3-5 hold, then sequence { z t } satisfies

Proof. By definition of z t , we have the first equation immediately. Take expectation on the first equation and we get

Note that and

which gives the second equation.

For the last equation, since

Take expectation and we have

Proof of Theorem 1. From Lemma 2, we have

Substituting Lemma 1 gives

Suppose θ ∗ is the optimal parameter and f ∗ = f ( θ ∗ ) is the minimal objective value. First, we have

Take expectation over the randomness in the ( t +1) -th step, we have

Since f is continuously differentiable and L-smooth, we have the following inequalities. [34]

Substitute them and get

where ρ ∈ (0 , 1] is a parameter used to balance the first two inequalities. Denote M = 2 ( α 1 -β ) 2 ( σ 2 + σ 2 m ) . Substitute Eq. 8 into this inequality and collect the terms, we get

When α satisfies the condition 2 α 2 (1 -β ) 2 -α (2 -ρ ) L (1 -β ) ≤ 0 , i.e. 0 ≤ α ≤ (1 -β )(2 -ρ ) 2 L , the term about ‖∇ f ( θ t ) ‖ 2 is non-positive, thus we have

Summing this inequality from 0 to T -1 and taking full expectation gives

which implies that

Since f is convex , we have Tf ( ¯ θ T ) ≤ 1 T ∑ T -1 i =0 f ( θ i )) . Subsequently we have

Finally, when α ∈ (0 , 1 -β L ] , we can take ρ = 1 , use L-smooth condition again and substitute M , which gives

Peer reviews:
comments: SUMMARY:
This paper concentrates on quantizing the optimizer states into low-bit to reduce training memory consumption. By thoroughly analyzing the first and second momentums, the authors quantize the optimizer states to 4-bit via the use of a smaller block size and both row-wise and column-wise information. They perform extensive experiments across various tasks.

STRENGTHS:
(1) The authors thoroughly analyze the outliers of momentums, which makes the existing block-wise method not work well.

(2) They perform extensive experiments across various tasks.

WEAKNESSES:
Compressing the optimizer states is essential especially for training large language models, but the paper only deals with how much memory is saved when fine-tuning LLaMA-7B on Alpaca. It would be necessary to evaluate LLaMA-7B fine-tuned via 4-bit AdamW on common sense reasoning tasks (to check whether the common season reasoning could be preserved even after fine-tuning via 4-bit AdamW) and/or MMLU (to check whether fine-tuned models via 4-bit AdamW can possess the instruction-following ability as a few-shot learner). In addition, it would be much more convincing if the experimental results of LLaMA-13B, 33B, and/or 65B are provided.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

comments: SUMMARY:
Today’s machine learning pipeline places a heavy load on GPU memory during training, due to the large number of parameters the optimizers have to maintain during training. This paper proposes a method to reduce the internal states of optimizers from their full-width counterparts to 4-bit numbers, while maintaining nearly the same performance across a range of tasks, such as vision and language models.

STRENGTHS:
This paper presented some intriguing observations, such as the 'zero-point problem' discussed in relation to second-order momentum, and the first-order term appears to be robust under quantisation. Additionally, this paper examined how quantisation might be combined with factorisation. The results presented covered a variety of language tasks as well as a vision task with relatively advanced models.

WEAKNESSES:
The figures are too small, making it hard to read. It is also unclear to me how the block size was chosen and whether there are theories to support this decision. Additionally, it would be beneficial to demonstrate a wider range of vision tasks to demonstrate the breadth of the proposed approach.

QUESTIONS:
- How do one generally pick the block size when we have an unseen model? Do we have theories for backing up the block size picking?
- Are there any other ways to demonstrate the convergence (from proof) rather than showing empirical results?
- Regarding the empirical results, more results might be need on 1) generative models (eg. diffusion) and 2) more vision tasks (eg. segmentation or object detection). The core idea is to demonstrate an efficeint optimizer, then the capability of this optimizer should be then tested on a range of different tasks.

LIMITATIONS:
I do not think this is applicable to this paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper develops 4-bit optimizers by using a smaller blocksize for the first moment and by analyzing and finding solutions to outlier patterns in the second moment. In particular it is found that quantization to the zeropoint are problematic for the second moment. The analysis is extensive and robust leading to novel and valuable insights into how optimizer states work in the low-bit regime.\n\n\nRecommendation:\n\nThis is a high-quality paper with extensive and valuable analysis and new methods with very robust evaluations. I am happy to fight for acceptance of this paper. I recommend that this paper be highlighted at the conference.\n\nSTRENGTHS:\n- The analysis in this paper is very robust, both theoretically and empirically. In particular, the identification of the zero-point problem is important.\n- Rank-1 normalization for quantization is ingenious and a valuable solution that goes beyond blocking/grouping that is done in many other quantization papers.\n- This work will have a strong impact on the general quantization literature way beyond quantized optimizers.\n\nWEAKNESSES:\n- This is outstanding work. I cannot find any weakness.\n\nQUESTIONS:\n- How do you store Rank-1 normalization constants to perform the dequantization? As I understand it, you use min(absmax row, absmax col) as the normalization constant for each x_{i, j}, but in this case, you need to store which constant is valid for each index (i, j), is that correct?\n\nSuggestion: You often write 2nd momentum. I believe the correct terminology would be first and second moment (look for central moment on Wikipedia). It is correct that the first moment is equivalent to the momentum term in Adam, but I think the second moment is usually not referred to as second momentum (rather a RMSProp term).\n\nLIMITATIONS:\nLimitations are fully and honestly discussed.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper proposes a 4-bit optimizer to save the memory of model training. They use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. They identify a zero-point problem of quantizing the second-order momentum and solve this problem with a linear quantizer that excludes the zero point.

STRENGTHS:
1. aims at solving a key problem to reduce the memory occupation for model training
2. conduct experiments on vision and language tasks

WEAKNESSES:
1. lack experiments on large models like 65B or 175B
2. The accuracy achieved by the 4-bit optimizer is still lower than the fp16 or int8 version. It is unsure whether it can behave stably under various settings.

QUESTIONS:
1. What about the effect on large models
2. Comparison with other optimization methods for large models (e.g., https://arxiv.org/abs/2306.09782)

LIMITATIONS:
This method can not achieve lossless accuracy.

The stability is not clear and sufficiently validated.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07