Memory Efficient Optimizers with 4-bit States
Authors: Bingrui Li
Jianfei Chen
Jun Zhu
Venue: NeurIPS 2023 spotlight (2023)
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
Bingrui Li', Jianfei Chen 't Jun Zhu' Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua- Bosch Joint ML Center, Tsinghua University
1br220mails. wttiinaisssssssssiiiee..eeee tsinghua. edu. cn; {jianfeic, dcszj) @tsinghua edu. cn
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 8-bit. this work, we push optimizer states bitwidth down 4-bit through 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 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- optimizers are evaluated on 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
Introduction
Large-scale models with a massive amount of parameters 20 22 49 58 have shown impressive
few-shot learning abilities 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 typically trained with stateful optimizers such 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 dominating factor of memory consumption 40 41]
There are several attempts to reduce optimizers' memory consumption. Factorization low-rank approximation to optimizer states, delta tuming [18 23 24
8 46 applies
27 28 avoids maintaining
optimizer states for most parameters by only tuning small subset, and low precision optimizers [15]
44 represent their states with low- precisiionn numerical formats, which consume less memory.
Among these methods, low-precision optimizers are attractive due to their simplicity and wide applicability. Dettmer et al. [15] propose an 8-bit optimizer with reparameterized embedding layers ("stable embedding layers") and 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- low-rrecision optimizers from 8 to 4-bit through analyzing the patterns in the first and second moments and designing dedicated quantizers
*Code is available at /htts::/// htgp::////thut www.toumtooeeeeeee..e m1/1om/ httm://uuunl/oooootttooootooooi bit.ooooiittt--ttttieiee.e *Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Algorithm Compression- based Memory Efficient Optimization Framework
Require: black-box stochastic optimization algorithm A, initial parameter wo E RP, initial opti- mizer state So O, total number of iterations T
1: for = 1, ... do
2: Sample minibatch ç and get stochastic gradient S: (wf((e-1..)
3: 5--1 f d)compress($t-1)
4: Wt; < AW
5: S f- compress(s:)
6: end for
7: return
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 this observation, we propose use smaller block size, which improves the approximation of first moment.
For the second moment, find its quantization suffers from zero-point problem. Since the update direction is usually inversely proportional to the square root of the second moment, quantizing non-zero quantities
will cause significant deviation. To address this problem, we propose simple linear quantizer to exclude the zero-point for second moment. We further propose 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, investigate the combination of factorization methods 2 2,46 with low- www-precision optimizers, and propose 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 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 [4 due to reduced communication cost.
Preliminaries
In this section, present preliminaries of compression- based memory efficient optimizers and discuss quantization methods for compression of optimizer states a general formulation.
2.1 A Framework for Compression- -bssd Memory Efficient Optimizers
Gradient-bast stateful optimizers like SGDM 38 47
Adam (26]. Adam W 32 are the principal
choices deep neural network training. However, the memory footprint of stateful optimizers several times of model itself, which results in a bottleneck for large model pretrainingfiietuning Consider the update rule of the Adam optimizer:
m: = -,mm--1 + (1 31)g.
Ve. 12 + +(1 32)g?
Adam(w; (w,-1, m;-1, Vt-1.8,)
m = -m,/(1 -8{
(1)
V. =v.((1 -8%)
We W 7) mi/ (Vv. + +c)
During the training process, the model parameters w, and optimizer states (i.e., first and second moments) m,, V, need to be stored persistently in the GPU memory. As the model grows large, optimizer states are 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 gradient-based optimizer A, memory efficient version of the optimization algorithm with compressed optimizer states is given by Alg. The algorithm A can be any gradient- based optimizers like SGDM, Adam, AdamW, etc. See App. F for examples. In Alg. the precise states S, are only temporal, and only compressed states S, are stored persistently in the GPU memory. Memory footprint reduction can be achieved since in neural networks, the state vectors m, and Vi 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 lower precision (e.g. 8-bit integers) is an effective way to com- pression optimizer states [5]. In this case, the optimizer states are compressed with a quantizer and decompressed with a dequantizer. The low 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 [I6), we disentangle the quantizer Q(-) into two parts: normalization (-) 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. Formally, the quantizer for a tensor x E R2 is given by
Ij Qxr,)) M N(rz).
Normalization The normalization operator N scales each elements of into the unit interval, i.c. [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 [1). The per-tensor and block-wise
normalization operators are given by
n2 == N per-tensor (Iz) = = max
SiSP
n2 Nniock wise (1z) = / max (la]::1+BB/B] SiSBU/B)))BBB/B++
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 scaling factor regardless of tensor size, but suffers from largest error due to outliers. Block-wise normalization views the tensor dimensional array, divides the array into blocks of size B called block and assigns a quantization scale within each block, which leads to [p/ B quantizaton scales totally. The block size could be adapted to trade-off quantization error and memory overhead.
1-
Mapping A mapping [15] converts normalized quantities to low-bitwidth integers. Formally, the mapping operator M MT, is equipped with a bitwidth b and a predefined increasing mapping, named quantization mupping T [0, 2 1] n Z + [0, 1]. Then M is defined as
4 M(n,) arg min |ng T())) 05i<2'
The design of T 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 [33. linear mapping T()) (i + 1)/2b 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 low-precision integer i into three parts: leading
3

---

(a)
(b)
(c)
(d)
20009
50000
100000
0.5
15000
0.6
0.4
30000
00000
10000
800
600
5000
10000
20000
1000
200 diml
2000
-10
10:1
10
-10-1
10"
10"
-10
-10-
10
10
Figure 1: Visualization of the first moment the layers. 3. blocks. 1. m.m.p.ffc layer in Swin-T model. (a): Magnitude of the first moment. (b): Histogram of the first moment. (c): Moment approximated by B128/DE.
(d): Moment approximated by B2048/DE.
sequence of E zeros, followed by indicator bit one, and remaining F fraction bits. DE defines T(i) 10-E()) fraction()),) 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 oracle normalization scaling the original tensor to uniform distribution is accessible, linear mapping could be used generally. On the contrary, the optimal mapping could readily identified with respect to certain metrics for per-tensor normalized tensor, there is necessity to develop additional normalization methods that incur extra overhead for mitigating the negative effects of outliers. In essence, 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
3 = N T(q)).
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 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 compression decompression methods in Alg. to realize memory efficient 4-bit optimizers. We first discuss the compression of the first moment. Our compression method for first moment based on Dettmers et al. (15), which uses block-wis. normalization with a block size of 2048 and dynamic exponent mapping [3]. 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 there is accuracy degradation.
To further improve the performance, we investigate the patterns of the first moment. There are some outliers the moments and outliers significantly affect quantization scale due their large magnitude. Outlier patterns have been studied for weights and activations. It is shown that the weights are rather smooth [4 [54 577, while the activation have column outliers 4 14 53 54). i.c., the outliess in activation always lie in some fixed channels. However, find that the outlier patterns in optimizer states are quite complicated. Specifically, Fig. shows patterns of first moment a Swin-T model during training. Outliers of the layers. 3. blocks 0 .lp. fc1 layer lie in fixed rows while outliers of the layers... ..3rs.3.k blocks...... O. m.mpp.p. fc2 layer 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. [ Block-wise normalization treats this tensor as flattened one-dimensional sequence, in the row-first order. Therefore, any block size of 2048 would include an entry in the outlier column, resulting

---

(a)
(b)
10
2218
Figure 2: Outlier patterns vary across two first moment tensors. (a): outliers lie in fixed rows (dimension O). (b): outliers fixed columns (dimen- sion l).
Figure 3: Histogram of the inverse square rool of second
moment. (a) full-precision; (b) quantized with B128/DE: (c) quantized with B128/DE-O, O. All figures at log scale and y-axis represents density.
in large quantization scale. In this case, block wise normalization is not better than per-tensor normalization. Therefore, we adopt a smaller block size 128, as it provides enhanced performance show that quantizing with smaller
while incurring only little memory overhead. In Fig. [ 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 improved normalization method with 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 (0, is desirable to contain zero in the quantization mapping for lower approximation етгor 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. I 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 (v) = 1/ (Vo + 10-6) of the second moment quantized with B128/DE The quantizer pushes most of entries of the tensor to zero, the inverse square root of most points fall into 10 due to zero-point problem, resulting in complete degradation in approximation. One remedy is to simply remove zero from the DE quantization map, which call DE-O. The smallest number representable by DE-O is 0.0033. With the B2048/ DDE-O quantizer applied to second moment, the approximation of second moment is more precise (Fig. 3 and the training with 4-bit optimizer stabilizes (Tab. However, by removing zero, DE-O 0 wastes one of the 2 16 quantization points. We propose to use Linear mapping T(i) = (i 1)/2b, whose smallest representable number is 0.0625. The linear mapping performs better than DE-O O for quantizing second moment.
In Tab. I ablate different quantization schemes and show that excluding from the mapping indeed the key factor for second moment quantization, which cannot replaced by a smaller block size or/and stochastic rounding [IT]. 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 non- embedding layers still suffer from zero points. On the other hand, when the zero-point problem is properly addressed, Stable Embedding is longer required retain accuracy.
4.2 Rank-1 Normalization
We propose empirically stronger normalization method named rank-1 normalization based on the observation of the heterogeneous outlier pattern in Sec. inspired by the SM3 optimizer [2] Formally, for matrix-shaped non-negative (second moment) tensor x E gnxm its I-dimensional
5

---

Table I:
Ablation analysis 4-bit optimizers on the second moment on the GPT-2 Medium E2E-NLG finetuning task. The first line barely turns 8-bit Adam [3] into 4-bit, ie. B2048/DE for both first and second moments. We only vary the quantization scheme for second moment. SR=stochastic rounding (see App. B.3 for details). Stable Embedding layers are quantized. 32-bit AdamW achieves BLEU of 67.7.
Normalization
Mapping
Stable Embed."
Factorized
(
BLEU
B2048
DE
X
X
33
66.6 + 0.61
B2048
DE
V
X
66.9 + 0.52
B128
DE
X
66
65.7 N/A
B128
DE+SR*
X
X
33
65,4 + 0.02 67.2 + 1.13
B128
DE
/
B2048
DE-0
X
X
67.5 + 0.97
B2048
DE-0
/
X
67.1 = 1.02
B128
DE-0
X
X
67.4 + 0.59
Rank-1
DE-0
X
X
67.5 + 0.58
Rank-1
Linear
X
X
67.8 + 0.51
Rank-1
Linear
X
67.6 + 0.33
statistics r E R" and E Rn are defined as r2 = maXi<jśm Tij and cj = max <ien Zi.. which 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
N (Ir.) min Z4g
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 1 normal- ization utilizes the I-dimensional information 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 out- liers effectively regardless of the patterns when a small block size is taken, but rank-1 normalization provides better trade-off between memory overhead and approximation. Rank-1 normalization falls back to per-tensor normalization for I-dimensional tensors, so we use B128 normalization those cases. Empirically, rank-1 normalization match exceed the performance of B128 normalization a moderate model scale (hidden thidden_size=1099 102e=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
few number of statisties, 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. Further, explore the effects of factorization 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.
second
Experiments
We compare our 4-bit optimizers with their full-precision counterparts, as well 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

---

Table 2: Performance on language and vision tasks. Metric: NLU=Mean AccuracyCCorelation. CLS-Aceuracy. NLG=BLEU. QA=FI. MT=SacreBleu. do not quantize optimizer states for embedding layers; B1 = 0. See more results in App. A
NLU RoBERTa-L
CLS
NLG GPT-2 M
QA RoBERTa-L
MT Transformer
Optimizer
Swin-T
32-bit AdamW
88.9 + 0.01
81.2 + 0.05
67.7 + 0.67
94.6 + 0.13
26.61
0.08
32-bit Adafactor
89.1 + 0.00
80.0 + 0.03
67.2 + 0.81
94.6 + 0.14
26.52 + 0.02
32-bit Adafactor!
89.3 0.00
79.5 + 0.05
67.2 : 0.63
94.7 0.10
26.45 0.16
32-bit SM3
87.5 + 0.00
79.0 = 0.03 0.03
0.29
22.72 = 0.09
8-bit AdamW!
89.1 + 0.00
$1.0
6
94.5
26.66 + 0.10
4-bit AdamW (ours)
89.1 + 0.01
80.8 + 0.02
67.8 + 0.51
94.5 0.10
26.28 0.05
4-bit Factor (ours)
88.9 + 0.00
80.9 + 0.06
67.6 + 0.33
94.6 + 0.20
26.45 = 0.05
31 > 0 and the 31 (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 I-dimensional second moment with Rank monk-1/Lihn 1/Linerr. See App. Diro details about the experiments.
Models, datasets and hyperparameters We report performance metrics standard benchmarks, including image classification (CLS) with Swin-T 31
on ImageNet- 1k [12, natural language [30 fine- time-ttuing on GLUE [51], question answering
understanding (NLU) by fine-tuning RoBERTa-L
(QA) by fine-tuning RoBERTA-L on SQuAD
42
43
natural language generation (NLG) by fine-
tuning GPT-2 Medium (39 E2E-NIG 335
machine translation (MT) by training Transformer-
Base 550 on WMT14 en-de
3 and LLLMMA [49
fine-tuning. We fine-tune LLaMA-7B, LLaMA-
13B and LLMMA-33B [49] on the Alpaca dataset
48 and evaluate them MMLU 21 and
standard common sense reasoning benchmarks: HellaSwar wag [56). ARC easy and challenge 10 and
OpenBookQA (33
We mainly follow the hyperparameters the original paper or/and codebase. In each benchmark. we keep same hyperparameters in one optimizer different quantization schemes, which gives an out-of-box transfer from full full-precision optimizer to low-bit optimizer without extra hyperparameter tuning. See App. D for hyperparameters and training details.
Accuracy 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 exceed 32-bit AdamW performance on fine-tuning tasks (NIJU, QA, and NL.G) and are comparable on all pretraining tasks (CLS and MTT). Sublinear memory optimizers Adafactor (3, = 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 W will not destroy the capability of pretrained models while enabling them to obtain instruction-following ability. 4-bit Adam is comparable with 32-bit AdamW W on all tasks and does get worse when the model size grows. Moreover, their convergence curves closely align (Fig.)
Memory and Computing Efficiency We evaluate the memory and computation efficiency of proposed 4-bit optimizers on instruction tuning, NLU, and NL. 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 because we report the total memory consumption (including data, activations, and memory fragments)
Attpa://g thw..//httttttt/// www//iirooseeeoo /Buint//////iiiee Syin- TransT-rrnn
sttp::/ wttp:///ctnnttceeeeeaeee httpr//cccctaatttrrrrrrr..
tttt://stthbb. cow/no Loo
attpa:/
then: tho.. wwwD/ plasrn Esamp
attps://g2ta http://tttbboottttteeeeeee con/tatsu lab/stanford alpaca
7

---

Table Performance on LLaMA fine-tuning on MMLU and commonsense reasoning tasks across
different sizes.
Model
Optimizer
MMLU (5-shot)
HellaSwag
ARC-e
ARC-c
OBQA
Original
33.1
73.0
52,4
40.9
42,4
LLaMA-7B
32-bit AdamW
38.7
74.6
61.5
45.1
43.4
4-bit A
38.9
74.7
61.2
44.4
43.0
Original
47.4
76.2
59.8
44.5
42.0
LLaMA-13B
32-bit AdamW
46.5
78.8
63.6
48.3
45.2
4-bit AdamW
47.4
79.0
64.1
48.0
45.2
Original
54.9
79.3
58.9
45.1
42.2
LLaMA-33B
32-bit Adam
56.4
79.2
62.6
47.1
43.8
4-bit AdamW
54.9
79.2
61.6
46.6
45.4
rather than the optimizer memory consumption alone. In principle, the optimizer states is 2x smaller for 4-bit Adam W than 8-bit AdamW, and about 4x smaller for 4-bit Factor.
The instruction tuning task uses two Nvidia A 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 reduced communication cost. For the smaller RoBERTa-L and GPT-2 Medium, our 4-bit optimizers appear 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 full- precision and our b-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 single 80GB GPU.
1.4
bil AcammW bit AdomWW
9.8
0.6
200
400
800
1005
Wairing Stop
Figure 4: Training loss curve of LLaMA-7B fine-tuning Alpaca dataset (averaged over runs). Only result 4-bit AdamW is reported since parameters are I-dimension FSDP packing. See more details App. D
Ablation Study Finally, we investigate the effectiveness of our proposed quantization schemes and the sensitivity of each moment to quantization in Tab. 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 step towards lossless performance. Factorizing second moment improves accuracy while leads to better memory efficiency.
6 Related Work
Compression-basee memory efficient optimizers There have been some works trying to ap- proximate the gradient statistics with sublinear memory cost relative the number parameters. Adafactor [46] achieves memory reduction by approximating the second-momen of matrix-shape

---

Table 4: Memory and Time of 4-bit optimizers compared with 32-bit Adam and 8-bit Adam [15].
Task
Optimizer
Time
Total Mem.
Saved Mem.
32-bit AdamW
3.35 h
75.40 GB
0.00 GB (0%)
LLaMA-7B
4-bit Adam
3.07 h
31.87 GB
43.54 GB (57.7%)
4-bit AdamW (fused)
3.11 h
31.88 GB
43.53 GB (57.7%)
32-bit A
3,93 min
5.31 GB
0,00 GB (0%)
8-bit Adam W
3,38 min
3.34 GB
1.97 GB (37.1%)
RoBERTa-L
4-bit A
5,59 min
3.02 GB
2.29 GB (43.1%)
4-bit A (fused)
3,17 min
3,00 GB
2.31 GB (43.5%)
4-bit Factor
4.97 min
2.83 GB
2.48 GB (46.7%)
32-bit AdamW
2.13
6.89 GB
0.00 GB (0%)
8-bit AdamW
2.04 h
4.92 GB
1.97 GB (28.6%)
GPT-2 Medium
4-bit A
2.43 h
4.62 GB
2.37 GB (34.4%)
b-bit AdamW (fused)
2.11 h
4.62 GB
2.37 GB (34.4%)
4-bit Factor
2.30 h
4.44 GB
2.45 GB (35.6%)
Table Largest trainable model under given memory budget. We use batch size of and max length of 512 for this comparison. FSDP is enabled at GPUs of 80 GB.
Table 6: Ablation study on the impact of compress- ing different moments to Swin-T pretraining on Ima- geNet1k.
Quant.
Quant. Factor, 2nd
Acc.
Largest fine-tunable Model
X
81.2 + 0.05
GPU Mem.
32-bit AdamW
4-bit AdamW
B2048/DE
X
80.9 1 0.04
24 GB
OPT-350M OPT-1.3B
OPT-1.3B OPT-6.7B
B128/DE
X
81.0 + 0,06
1128/DE
Rank-I/LLneas i/Liear
X
80.8 + 0.02
80 GB
B128/DE
Lonk-I/Linear
V
80.9 + 0.06
80 GB
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 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 memory reduction second moment and are applicable to most of the second moment based optimizers, like Adagrad 19 Adam 26),, Adam W
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 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 applicable to SGDM and Adam/ Adam W. 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 quanti- zation in forward phase and dequantize the low-precision activations to full- full-precision -nieryy in backward phase. Gradient checkpointing 77] only keeps activation 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 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
We propose 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 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.
Acknowledgements
This work was supported by the National Key Research and Development Program of China (No. 20217D01 10502), NSFC Projects (Nos. 62061 136001, 62106123, 62076147, U19A2081, 61972224, 62106120) Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
References
[1] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-efficiene sgd via gradient quantization and encoding. Advances in Neural Information Processing Systems, 30, 2017.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
[3] Ondřej 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 statistical machine translation. In Proceedings of the ninth workshop statistical machine translation, pages 12-58, 2014.
[4] 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.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Merbert-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, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[6] Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Actmn: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, pages 1803-1813, PMLR, 2021.
[7] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint 1rXwv.1604.0617, 2016.
10

---

[8] Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. In International Conference on Learning Representations, 2020.
[9] 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 preprini arXiv:2204.023) 2022.
[10] 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. 05545, 2018.
[11] Matthieu Courbariaux, Yoshua Bengio, and Sean-Pierre David. Binaryconnect: Training deep neural networks with binary weights during propagations. Advances in Neural Information Processing Systems, pages 31 3123-313,, 2015.
[12] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, 2009.
[13] Tim Dettmers. 8-bit approximations for parallelism in deep learning. arXiv preprint 1..05:1511.0446/ 2015.
[14] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O: 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07333 2022a.
[15] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations, 2022b.
[16] Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. arXiv preprint arXiv:2212.0973)) 2022c.
[17] 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.
Ming-Wei
[18] 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.066 06904, 2022.
[19] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[20] 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.
[21] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03000 2020.
[22] 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 Damoe, 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.
[23] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficien transfer learning for nlp. In International Conference Machine Learning, pages 2790-2799. PMLR, 2019.
11

---

[24] Edward 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.
[25] 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
inference. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2704-2713, 2018
[26] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[27] 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.
[28] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2 201.00/900 2021.
[29] 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. 1357, 2022.
[30] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint
2019.
[31] 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.
[32] Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
[33] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. arXiv preprint arXiv: wrXiv:1809.02788 2018.
[34] Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
[35] Jekaterina Novikova, Ondřej Dušek, and Verena Rieser. The e2e dataset: New challenges for end-to-end generation. arXiv preprint arXiv: / 171.70066000924 2017.
[36] 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
1rXiv:190..01038 2019.
[37] 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 22xx52206.09957, 2022.
[38] Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks, 1555)1145-151, 1999.
[39] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18)99 2019.
[40] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimiza- tions toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEEE, 2020.
12

---

[41] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: 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.
[42] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. arXiv preprint arXiv: /806.03322 2018.
[43] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: 100.02560000525000, 2016.
[44] 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 Machine Learning, pages 8821-8831 PMLR, 2021.
[45] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. {ZeRO- Democratizing (Billion-Scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564, 2021.
[46] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[47] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International Conference on Machine Learning. pages 1139-1147. PMLR, 2013.
[48] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-follooinn llama model. https://gt thub. com/ http:/tttullaabsaanne ord alpaca 2023.
[49] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 33771, 2023.
[50] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[51] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 11X4v:180.004400 2018.
[52] 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.
[53] 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 A2Ii:2220911325.. 2022.
[54] 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.
[55] 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-22 177777788711,,,,, 2022.
[56] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint
2019.
13

---

[57] 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.
[58] 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
2022.
14

---

A Additional Experiment Results
In this section, we show additional experiment results beyond Tab. 2 RoBERTa-L. finetuning on each task in GLUE datasets. Tab. shows the results of GPT-2 Medium finetuning on E2E-NLG via different metrics. Tab. shows the EM and FI of RoBERTa-L finetuning on SQuAD and SQuAD 2.0 datasets.
Tab. shows the results of
Table 7: Performance of MoBERTa-Large finetuning GLUE with diverse optimizers. Medians and std over 5 runs are reported on all tasks. I: do not quantize optimizer states for embedding layers; i 31 =
Optimizer
MNLI
QNLI
QQP
RTE
SST-2
MRPC
CoLA
STS-B
32-bit bit Adumi
90.2 1 0.00
94.9
0.00
92.2 0.00
85.2 1: 0.14
96.3 1. 0.00
93.2 0.01
66.9 1 0.01
92.3.0000 0.00
32-bit Adafactor
90.4 + 0.00
94.7 0.00
92.2
85.9
0.00
92.8 0.00
67.3 0.01
0.00
32-bit Adafactor!
90.5 + 0.00
94.8 0.00
92.2
87.0 0.03
92.9 0.00
68.2
9.00
90.6
9..0+00.00
90.5
9..4+000
bit AdamW"
90.4 1 0.00
94.8 0.00
92.2 0.00
84,8 1 0.02
96.2 16.2.10000
93.2 0.00
68.0 1 0.00
92.2.100 0.00
4-bit Adami
90.2 1: 0.00
94.5 1: 0.00
92.0 1 0.00
85.2 1: 0.12
96.3 0:00
92.8 0.00
67.3 1: 0.01
92.3
4-bit
90.1
9..4
92.7
68.1
92.3
Table 8: Performance of GPT-2 Medium finetuning E2E-NLG Challenge with diverse optimizers. Means and std over 3 runs are reported.
Optimizer
BLEU
NIST
METEOR
ROUGE-L
CIDEr
32-bit AdamW
67.7 1: 0.67
8.60 1 0.08
45.7 1: 0.28
68.7 1: 0.61
2.35 1 0.04
32-bit Adafactor
67.2 0.81
8.61 = 0.60
45.3 = 0.08
68.3 = 0.22
2.35 + 0.01
32-bit Adafactor
67.2 1.22:0.6 0.63
454
45.6 0.32
68.5 0.30
2.32 0.02
32-bit SM3
66.9 1 0.58
8.59
45.4 0.32
68.2 + 0.49
2.33 0.03
8-bit Ad
67.5 1 0.87
$.59
45.7 0.52
68.7 1 0.97
2.34 1 0.06
4-bit AdamW
67.8 + 0.51
8.61 + 0.08
45.8 + 0.23
68.9 = 0.33
2.35 + 0.07
4-bit Factor
67.6 + 0.33
8.59 + 0.03
45.6 + 0.43
68.6 0.60
2.34 0.06
Table 9: Performance of RoBERTa-Large on SQuAD and SQuAD 2.0 with diverse optimizers. Medians and std over runs are reported.
SQuAD
SQuAD 2.0
Optimizer
EM
F1
EM
F1
32-bit AdamW
89.0 + 0.10
94.6 + 0.13
85.8 0.18
88.8 0.15
32-bit Adafactor
88.8 = 0.12
94.6 0.14
85.8 - 0.44
88.7 0.21
32-bit Adafactorf
89.0 + 0.18
94.7 0.10
85.9 + 0.15
88.8 0.15
32-bit SM3
84.2 + 0.49
91.7 +
77.2 + 0.71
81.1 0.66
8-bit AdamW!
88.8 + 0.15
9..5
86.1 0.26
89.0 0.26
4-bit AdamW
88.8 = 0.08
94.5 0.10
85.4 + 0.28
88.4 0.26
4-bit Factor
88.8 = 0.38
94.6
85.9 + 0.36
88.9 = 0.18
15

---

B Outlier Patterns of Moments
In this section, give a comprehensive visualization about the outlier pattern of optimizer states. [ 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 parameters and activations in [54].
The outlier pattern of first moment depends many factors such data and training hyperparameters. Here we mainly focus on different transformer models and different layers inside. In one transformer block, there is 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 small portion. We denote the matrices by Wº, W w W W W respectively. When the matrices across different layers are involved the same time, add subscript indicating layer index. Note that W has shape C, in a linear layer, we call the output and input dimension by dim0/row and dim l/columnn, respectively. The per-channel quantizaition used in other works actually correspond to per-row(dim() normalization here.
Swin Transformer ImageNet pretraining Fig. 7 the magnitude of first moment in trans- former 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 and W that outliers occur fixed columns, becomes obvious while other parameter matrices remain noisy. At layer the patterns in Wo W1 w quite obvious. In W and W the outliers occur fixed rows while the pattern in W remain unchanged. The 1-dimensional structure in Wq w w 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 in Fig. 6
VQ
10
10
cm
53
100
1513
Figure 5: Outlier patterns of first moment in transformer block layers. O. blocks. Swin-T at epoch 210.
16

---

6.
00
98
102
223
200
262
200
1230
aer
1200
305
300
233
Figure 6: Outlier patterns of first moment in transformer block layers 2. blocks 0 of Swin-T at epoch 210.
03
08
200
06
66
500
100
500
2006
200
700
300
833
Figure 7: Outlier patterns of first moment in transformer block layers 3. blocks........ 0 of Swin-T epoch 210.
17

---

12 the magnitude of first moment in transformer blocks of RoBgRTa-Large at different depths are shown. At layer 0 and layer (initial layers), patterns W w 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 w are obvious. Patterns in other matrices are weak.
RoBERTa-Large GLUE finetuning In Fig. 11
203
253 499
268 cst
202 XX
590
2000
der
co0
1000
arr
1063
WO
06
06
3000
257
457mmm
1062
2002
200
2000
0m
Figure 8: Outlier patterns of first moment transformer block layer-0 of RoBERTa-Large at epoch 8.
500
203
400
203 400
200
40
WO
WL
10
10
14
66
86
200
LEC0
1005
3000
1082
2007
2000
Figure 9: Outlier patterns of first moment transformer block layer-1 of RoBERTa Large at epoch 8.
18

---

30
10
1000
1301
630
1003
2000
203
200
2000
4000
2003
Figure 10: Outlier patterns of first moment in transformer block layer-11 1 of RoBERTa-Large at epoch
2000
1000
1903
106
020
000
3003
81
00
30
303
300
2000
ACCO
der
1063
Figure 11: Outlier patterns of first moment in transformer block layer- 12 of RoBERTa-Large at epoch
19

---

30
10
2000
1000
1901
1003
203
400
200
2000
4000
2003
Figure 12: Outlier patterns of first moment in transformer block layer-22 of RoBERTa-Large at epoch
3000
1000
1903
208
020
000
1000
1000
3003
00
30
303
400
300
2000
AECO
der
acs
1063
Figure 13: Outlier patterns of first moment in transformer block layer=23 of RoBERTa-Large at epoch
20

---

GPT-2 Medium E2E-NLG finetuning In Fig. 14 transformer blocks of GPT-2 Medium at different depths are shown. At layer and layer 2 (initial layers), patterns in w are obvious. At layer 13 and layer (intermediate layers), patterns in W W are obvious. At layer 21 and layer 22 (last layers), patterns W w w w obvious. First moment of W W are consistently noisy throughout layers. It is notable that the rows(or columns) that gather outliers are different across different layers.
19 the magnitude of first moment
VQ
VK
10
10
2000
1000
1361
359
2000
1000
4002
630
1000
200
Figure 14: Outlier patterns of first moment in transformer block layer-1 of GPT-2 Medium epoch 2.
0.4
04
2000
1000
1965
030
000
0.5
08
08
400
200
5000
JECO
1943
Figure 15: Outlier patterns of first moment in transformer block layer-2 of GPT-2 Medium at epoch 2.
21

---

30
10
2000
1000
1301
630
1000
1003
2000
203
200
2000
4000
2003
Figure 16: Outlier patterns of first moment in transformer block layer-13 of GPT-2 Medium at epoch
10
3000
1000
1903
620
000
1000
1000
3003
00
30
2000
303
300
2000
AECO
der
2003
Figure 17: Outlier patterns of first moment in transformer block layer-14 of GPT-2 Medium at epoch
22

---

30
10
2000
1000
1301
630
1003
2000
200
400
200
2000
4000
2003
Figure 18: Outlier patterns of first moment in transformer block layer-21 of GPT-2 Medium at epoch
2000
1000
1305
020
000
1000
3003
00
30
323
300
2000
AECO
1003
Figure 19: Outlier patterns of first moment in transformer block layer-22 of GPT-2 Medium at epoch
23

---

Quantization Quality via Histogram
C.1
Zero-point Problem
In Fig. 22 we show the effect of zero-point on quantization enror 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 great scale.
Figure 20: Histogram of second moment of the attention layer (W?, w w in transformer block-wise layer-2 of GPT-2 Medium at epoch In one horizontal line, the first figure is the original second moment. The second figure is the tensor after normalization. The third figure the quantized tensor. The last figure the dequantized object. Both the first and last figure is at log10 scale. Both the second and third take values in [0, 1]. All y-axis represents density. Good quantization methods try to make the third figure identical to the second figure and make the last figure identical to the first figure. Top: B B828/DE quantization. Bottom: B B128/DE-O quantization.
Figure 21: Histogram of second moment of the V in transformer block layer- 10 of RoBERTa-Large epoch Top: B128/DE quantization. Bottom: B12 B128/DE-O quantization.
24

---

Figure 22: Histogram of second moment the attention layer (W w W in transformer block
layers. 0. blocks. 0 of Swin-TT at epoch 210. Top: B128/DE quantization. Bottom: B B128/DE-0 quantization.
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, wise normalization are shown Fig. 23
cases where rank-1 normalization approximates better than block- Also, some cases where rank-1 normalization
25 27
approximates worse than block-wise normalization is shown in Fig. [4] 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.
Figure 23: Histogram of second moment of W in transformer block layer-23 of GPT-2 Medium at epoch 2. case where rank-1 normalization is better than block-wise normalization with block size 128. In this case, the tail in the right side of distribution is captured by rank-1 normalization but lost in block wise normalization. Top: B128/DE-0 quantization. Bottom: Rank- 1/DE-0 quantization.
C.3 Effectiveness Block Size in Block-wise Normalization
In Fig. 29 30 3 we show the effect of block size on quantization error for both first and second
moments. Fig. 29 30 shows that B2048 normalization quantizes significant portion 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, shows smaller block size improves quantization quality on second moment.
25

---

Figure 24: Histogram of second moment of the attention layer (W w W in transformer block layer-2 of GPT-2 Medium at epoch 2. A case where rank-1 1 normalization is worse than block-wise normalization with block size 128. Top: B128/DE-0 quantization. Bottom: Rank-1/DE-0 quantization.
Figure 25: Histogram of second moment of 2 in transformer block layer-4 of RoBERTa-Large at epoch 8. A case where rank-1 I normalization is better than block-wise normalization with block size 128. Top: B128/DE-0 quantization. Bottom: Rank- 1/DE-0 quantization.
Figure 26: Histogram of second moment of W in transformer block layer-2 of RoBERTa-Large at epoch 8. case where rank- normalization is worse than block-wise normalization with block size 128. Top: B128/DE-0 quantization. Bottom: Rank- 1/DE-0 quantization.
26

---

Figure 27: Histogram second moment 2 transformer block layers O. blocks of Swin-T at epoch 210. A case where rank-1 normalization is better than block-wise normalization with block size 128. Top: B128/DE-0 quantization. Bottom: Rank-I/DE-0 quantization.
Figure 28: Histogram of second moment of W in transformer block layers. bblocks. of Swin-T at epoch 210. A case where rank-1 normalization is worse than block-wise normalization with block size 128. Top: B128/DE-0 quantization. Bottom: Rank-1/DE-4 quantization.
1000
Figure 29: Histogram of first moment of W in transformer block layer-20 of GPT-2 Medium
epoch 2. Top: B128/DE quantization. Bottom: B2048/DE quantization.
27

---

1000
Figure 30: Histogram of first moment of w in transformer block layer-22 of RoBERTa-L at epoch
8. Top: B128/DE quantization. Bottom: B2048/DE quantization.
Figure 31: Histogram of second moment of W in transformer block layers 0. blocks. of Swin-T at epoch 210. Top: B128/DE-0 quantization. Bottom: B2048/DE-0 quantization.
D Experimental Details
D.I Quantization
There are several parameters 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 experiments we conducted, the rule is straightforward: tensors with size smaller than equal to 4096 will not be quantized. However, for larger models with a hidden size exceeding 4096, it advisable to exclude the bias and normalization layers from quantization. Regarding the quantization settings, stated in Sec. 5 we employ block-wise normalization with 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 will factorized while 1 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.
28

---

D.2 Hyperparameters and Training Details
In each benchmark, unless otherwise specified, we maintain the same hyperparameters for given optimize across different quantization schemes. Additionally, we use same optimizer hyperparameters across various optimizers, including our 4-bit optimizers, 8-bit Adam [[3]. SM3 [2, Adafactor [46] and the full precision counterpart AdamW [32] For Adafactor, we use 31 > O as default setting which same as the 81 value used in AdamW. Also, the case where 31
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 31 configuration, same as the 31 value used in AdamW.
Table 10: The hyperparameters for RoBERTa-L fine-tuning on GLUE.
Dataset
MNLI
QNLI
QQP
RTE
MRPC
SST-2
CoLA
STS-B
Batch Size
32
32
32
16
16
32
16
16
LR
le-5
le-5
le-5
2e-5
1e-5
le-5
le-5
2e-5
Warmup
7432
1986
28318
122
137
1256
320
214
Max Train Steps
123873
33112
113272
2036
2296
20935
5336
3598
Max Seq. Len.
128
128
128
512
512
512
512
512
Table 11: The hyperparameters for RoBERTa-L fine-tuning SQuAD and SQuAD 2.0.
Dataset
SQuAD & SQuAD 2.0
Batch Size
48
LR
1.5e-5
# Epochs
2
Warmup Ratio
0.06
Max Seq. Len.
384
Table 12: The hyperparameters for GPT-2 on E2E.
Dataset
E2E
Training
Batch Size
8
LR
4e-5
# Epochs
Warmup
500
Max Seq. Len.
512
Label Smooth
0.1
Inference
Beam Size
10
Length Penalty
0.8
repeat ngram size
RoBERTa We train all of our RoBERTa-L models with PyTorch Huggingfacd On GLUE benchmark, we mainly follow the hyperparameters in fairseq [36]. We use 31 0.9. 32 = 0.98, le-6, weight decay factor of 0.1 and linear learning rate schedule. Other hyperparameters On SQuAD benchmark, we mainly follow the reported hyperparameters in RoBERTa paper 30 We use 81 0.9, B2 0.98, 1e-6, weight decay factor of 0.01 and linear learning rate schedule. The other hyperparameters are listed in Tab. IT 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.
are listed in Tab. 10
29

---

On SQuAD 2.0, there may be 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] handle unanswered questions may differ from the solutions utilized in the BERT paper [7]. 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 31 = 0.9, 32 0.999, c le-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-Bate 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 H We completely follow the hyperparameters in the codebase. We report the mean and standard deviation results over 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 MAM,, LLaMA- 13B and LLaMA-33B with Alpaca codebasc We follow the hyperparameters in the codebase for LLaMA-7BB and LLaMA- and the hyperparam- eters 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 I-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 Adam W 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 A 1100 80G GPUs, RoBERTa-L using 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 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
nj == sign(z,) ( ((æ5])
Therefore, the unit interval for signed case [-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, mainly consider linear mapping and dynamic exponent mapping IK] See Fig. for illustration of quantization mappings.
httpe://ssthbb com/no cronoft/LoRA
httpe://esthbb. http:///tt...eeto///////////////
attpa://y wttp:///thhhhhc///hhht wwwphhhcoo/neereeeeeeeeee Transfomme
2ttps://gitan cataub..c//ttttttt Jab/stanford..ommm alpaca
30

---

syromk: cREsnOY
denanic SUBOTER
0.75
0.0
0.32
0.22
0.9
0.05
0.52
0.2
-102
0.0
Figure 32: Visualization of the quantization mappings for the linear and dynamic exponent at 4-bit precision. Left: Signed case. Right: Unsigned case.
Linear mapping It is notable that the linear mapping considered in work does not include zero both signed case and unsigned case. Actually, we only use linear mapping in unsigned case, which defined as torch. linspacceeeccc (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 F(i) fraction(t). In following paragraphs, we will define the dynamic exponent mapping formally based the binary representation.
In unsigned case, dynamic exponent mapping 13 is composed of exponent bits F,, one indicator bit and fraction bits F, where à = 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 (0.1, 1), which is formally defined as
1-0.1
j 0.1, 05i52
Pj
2F
Pk Pk+1
fraction(k] [k]
0 < 1
2
Therefore, number with E exponent bits and F fraction bits valued & has value of
10
fraction(k].
For signed case, the only difference is that dynamic exponent mapping additionally uses the first bit as the sign thus we have b + E + + F. Specially, at 8-bit case, we learn from the codebase FFF that dynamic exponent mapping assign 00000000; O 00000001, 1 in unsigned case and assign 10000000, and 00000000, with 1 and O1o. respectively. This means -1 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 length of 2", 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 say barely 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. In this section, talk about how integrate stochastic rounding into our formulation of quantization. When stochastic rounding is used, the definition of mapping M has some minor changes. Specifically, M still an element-wise function and defined as
M(n;) arg min {ng T() : ny T()) 2 0} arg max {ng T(i) nj T()) s 0} 05
0S+<29
httpo://gitaab. www.Haaaeeseeee/tttaaaddrrrrrere
31

---

In other words, M maps each entry n, to the maximal index set M(n,) such that for any i E M(n;) there is no other 0 & k v 2 - 1 with T(k) lying between T(i) and nj. Actually, M acts as a filter of 'T and give a more fine- iin-grained range of quantized output candidates. this definition, M(n,) has only one or two points since only stochastic rounding is considered in the final step.
Finally, define stochastic rounding R. When M(n,) only has point, R, just output this point. When M(n)) has two points and 92 with T())) < my (g2), stochastic rounding (R,)) defined as
22 with proba. T(q))-
T((1)
R, (nj, 91, 92) =
Ti @1 with proba. T(m)) T(n)
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.
Algorithm 2 Compression based Memory Efficient SGDM
Require: initial parameter Ou E RR, learning rate a, initial first moment mo = 0, total number of iterations T and momentum parameter B.
1: for /.= 1, 1,.,...,? " do
2: Sample minibatch and get stochastic gradient gr (of(0l-1,G)
3: m4-1 f- decompress(m::))
4:
m, f 3 m4-1 + gt
5:
Og - 0,-1
m,
6: Mi compress(m.)
7: end for
8: return 0r
Algorithm Compression- based Memory Efficient Adam
Require: initial parameter fl E IRP, learning rate a, initial moments mp 0, Ug 0, total number of iterations T and hyperparameters C2,c.
1: for = 1,2, ,...,, do
2: Sample minibatch Se and get stochastic gradient 9t. Vol(O--1,6)
3: mc-1; 1/-1 +- decomppess(mp)- decompress( (6,-1)
4:
m 81 -1-m++
{1) .
5:
Ur + 1g-0 81)
6:
m. mm 35)
7:
i + v1/(1 (35)
8:
O, + 0,-1 - a m, (V +
9:
m, 8. compress (mppresss(.,)..... compress(v.)
10: end for
11: return 0r
G Rank-1 Normalization
In this section, we present the detailed formulation of rank-1 normalization in Alg.
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.
32

---

Algorithm 4 Rank-1 Normalization
Require: tensor RRd X:=xxp; statistics My E Rd for p, permutation function A
mapping (1.....)} ,d] to indices of tensor 2; where d = dj x ... x dp.
for 1, 1........ do
2: for j ....... d, do
3: Pr.j max,
4: end for
5: end for
6: for = 1,2........ d do
7: M, gis (r))(().
8: end for
9: reshape 1-dimensional array M to the same shape as
10: return z/M
First, we make some assumptions. The first three are rather standard in stochastic optimization literature, while last two depict properties of stochastic quantizers.
1. (Convexity) The objective function is convex and has unique global minimum f(0"). 2. (Smoothness) The objective f()) is continuous differentiable and L-smooth;
3. (Moments of stochastic gradient) The stochastic gradient g unbiased, i.e., Blg(O)] (()).
and has bounded variance, i.e., E [l)((0) v[(0)]]] < v0 Rd.
4. (Unbiased quantizer) Vx Rdd E [x()]]
5. (Bounded quantization variance) Va E Rd E [10m(x) =x1?] s o
Then, we have following theorem:
Theorem 1. Consider the Algorithm with Assumptions 1-5. Let a (0, 211, then for all T
have
L
3
(lf(Or) - fa] -
1- +
110 - 1.
2T
a
a
ao? m
+
(1 - B) +
(2)
(1 (1- B)
where OT t ET-- Oz.
H.1 Proof of Theorem
To prove Theorem 1 we need some useful lemmas.
Lemma 1. In Algorithm 2 The conditional first and second moments of ge satisfies
Elga O=-1]= (0,-.)
(3)
E
|1 10-1 < [V](0,-))]1 +
(4)
Proof. By assumption, we easily have
E [9.10.111 -
With Assumption 3, it holds true that
E
[1][0)]]]
E
[lx(0) VS(0) |/
E [ls(0) [ = [
+ 2E (g(() f((), vf(O)]
E [lo(0) V](0)) + E [
< 02 + |V 2
which implies the second part.
33

---

Lemma 2. If Assumptions 3-5 hold, then sequence {2r} } satisfies
1
33
2++1 - it =
1 B (0c+1 - 0.)
1- 10x - 02-1)
(5)
B[2141 - 2x] = -a
1 v((0.)
(6)
B Bl1 - S I 1-3 )²
(Bll9++111 o
Proof. By definition of Zt, have the first equation immediately. Take expectation the first equation and we get
E [11+ - 21] = 1- E 10011 - 0o] 1 - -3 10 8(-1]-
Note that
E [0 06] E [0.+1 (0r amm+11] [am++1]
-aaE [m4+1]
-aE [8m, + 9t+1]
-a3E [m,] - f(t))
and
E 00 0 [0 (0 am [mm]]
= =aaE [me].
which gives the second equation.
-a
E 2111
21]
1 ((.)))
For the last equation, since
B
24+1
28
(02+1 - 01)
(0. 02-1)
1
1 - B
(m++1 3mm)
3
Take expectation and we have
2
|1111 - 24 ==
=
a
|
Bmell
B
2
< 2
||mm+1
(Bme + 9(+2)11] E [119441] r²])
I 3
< 2
a
E
[|]
+ amm)
1- 8
Proof of Theorem 7 From Lemma 2, we have
| - =012 < ( a ))
2 E [1 12] + +ool)
1 - 3
Substituting Lemma gives
a )) (IV +02 +0²+02) oz)
E [|24 11:41-22111
z0l1 2
1 (
3)
(8)
34

---

Suppose 0, is the optimal parameter and f. f(0+) is the minimal objective value. First, we have
|| 1. = || - 0.1 + 2 (2t - 0,, 2++1 - 2x) + || - =ell?
Take expectation over the randomness in the (t + 1) - step, we have
2a
B[|2+11 0.11% = 1| - -0111
(2+ 0,, v.((.)) + E[1 - =el1"]
1 - 8
2a
-l|ze- - -0
(0g 0., f(.)))
1
20/3
8)2 (8: -8 V (())) + El - ==11
Since f is continuously differentiable and L-smooth, we have the following inequalities.
34
(0. - 0,, v.((o))) 2 |V 8.
(9)
(8: -2- v.((.))) 2 f(O)) + 2L [V.)(0))]1
(10)
(0. - 11 ((,)) 2 f(0)) f(01-1)
(11)
Substitute them and get
20(11 - p)
20p
1++1 - -0.1171 $ 1| - 0,11
[V(00)]
1 - (f(()) - f=)
L(1 - 3)
ap [V))()))
(1- (f((t) f(((-1)) E|| zell"
2aB
L(1-3) - 3)
where p E (0, 1] is a parameter used to balance the first two inequalities. Denote M =
2
2 + o Substitute Eq. 8 into this incquality and collect the terms, get
2ap
20.
+
(f(0,) - (+) + E[[1=+11 - -0,113
T- B
(1-
(1 B)2 (f(0-11) fa) + 1| - -0.12+ + (1 --2 L(1 - ())) / [V)0))I + M
20,3
20
ax22 - p)
<
When satisfies the condition
2a²
al2
< 0, i.c. 0 a (1- (2-0)
the term about
31-37
L
[V.()))) is non-positive, thus have
2ap
2aB
(f(O)) - f.) + E[]l2++1 - 1&
1-3
(1-
2a3
(f(0)-1) - f.) + || 0-11 M
(1 - 3)²
Summing this inequality from 0 T and taking full expectation gives
T-1
T
2ap
Elf(O.) - f.]+
Σ
(1-3y² (l. f.l + Ell 0.119)
20,3
I
1=
2a,3 3)² Eff
f.1 + 0.119) + T M
which implies that
2ap T-1
20,3
E[f(0.) fal (1- 3)2 (f(w)) - fc) + 1100 0.11 T M
1 3 0
35

---

Since f is convex, we have Tf((r) - T Ei= f(6,)). Subsequently we have
[[S(Or) [.] - 7 p( (3) (5(00) -s-) 1-3
1
B
100 0-11
20x
- 3
+
M
2ap
Finally, when a E (0, 1-91 we can take p = 1, use L-smooth condition again and substitute M, which gives
L3
1 3
Elf(Or) - f.l -
+
110
0
2T 1- 3
ao²
a
+ (1-
(1 3)
36


Reviews:
[{"review_id": "Iniphcg75A", "reviewer": "Reviewer_rPDD", "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.", "questions": "1. What about the effect on large models\n2. Comparison with other optimization methods for large models (e.g., https://arxiv.org/abs/2306.09782)", "limitations": "This method can not achieve lossless accuracy.\n\nThe stability is not clear and sufficiently validated.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. aims at solving a key problem to reduce the memory occupation for model training\n2. conduct experiments on vision and language tasks", "weaknesses": "1. lack experiments on large models like 65B or 175B\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer rPDD for the constructive review. With respect to your questions:\n\n**Weakness 1 + Question 1: Experiments on large models.**\nWe have fine-tuned LLaMA-7B and LLaMA-13B with 32-bit AdamW and 4-bit AdamW on Alpaca and evaluated them on common sense reasoning tasks and MMLU. Results show that 4-bit AdamW will not destory the capability of pretraiend models while enabling them to obtain instruction-following ability. In MMLU and HellaSwag tasks, 4-bit AdamW fine-tuning outperforms 32-bit AdamW both on LLaMA-7B and LLaMA-13B. In other tasks, 4-bit AdamW is comparable with 32-bit AdamW.\n4-bit AdamW does not get worse than 32-bit AdamW when the model size grows. Also, we are trying to fine-tune LLaMA-30B, and will post updates when available. However, the amount of computation required for LLaMA-65B exceeds our accessible resources.\n\n|    LLaMA-7B | MMLU(5-shot) | HellaSwag | ARC-e | ARC-c | OBQA |\n|-------------|-------------:|----------:|------:|------:|-----:|\n| Original          |          33.1|       73.0|   52.4|   40.9|  42.4|\n| FT w. 32-bit AdamW|          38.7|       74.6|   61.5|   45.1|  43.4|\n| FT w. 4-bit AdamW |          38.9|       74.7|   61.2|   44.4|  43.0|\n\n|    LLaMA-13B| MMLU(5-shot) | HellaSwag | ARC-e | ARC-c | OBQA |\n|-------------|-------------:|----------:|------:|------:|-----:|\n| Original    |          47.4|       76.2|   59.8|   44.5|  42.0|\n| FT w. 32-bit AdamW|          46.5|       78.8|   63.6|   48.3|  45.2|\n| FT w. 4-bit AdamW |          47.4|       79.0|   64.1|   48.0|  45.2|\n\n**Weakness 2: About lossless accuracy and results on more settings.**\nIt is true that our 4-bit optimizer does not converge losslessly on all tasks. However, our 4-bit optimizer can achieve lossless results on all language model fine-tuning tasks, including fine-tuning RoBERTa-Large on GLUE and SQuAD, fine-tuning GPT-2 Medium on E2E-NLG, and fine-tuning LLaMA-7B and LLaMA-13B on Alpaca. This already demonstrates its practical applicability. \n\nWe are also running our optimizers on other tasks. We will post updates on these tasks within the discussion period. \n\n**Question 2: Comparison with other optimization methods for large models.**\nRecently, there have been many works focusing on optimization for LLMs.\n\nSophia [1] is a second-order optimizer designed for language model pretraining. \nSophia converges faster than the broadly used AdamW on GPT-2 pretraining task, \n*but it does not concern about memory efficiency and it has same memory cost as AdamW*. \n\nLOMO [2] is a memory efficient optimizer for fine-tuning language models. \nLOMO is essentially a vanilla SGD optimizer without momentum, and thus it naturally removes the memory consumption of optimizer states. Then, LOMO fuses backpropagation and optimizer update into one step to further reduce memory consumption of gradients and full-precision copies of weights. *However, as LOMO is just a vanilla SGD optimizer, its performance may be worse than Adam in some tasks.*\nInstead, our 4-bit optimizers focus on reducing memory usage of the first and second moments in stateful optimizers (e.g., Adam), and has a wider range of applicability.\nThe fusion technique in LOMO and our 4-bit optimizers complement each other, and could bring a more memory efficient Adam while maintaining performance. \n\nMeZO [3] is a memory efficient zeroth-order optimizer for fine-tuning language models. MeZO adapts the ZO-SGD algorithm by using in-place operation and resetting random seed. It only requires the same memory as inference (i.e., only forward parameters).\n*However, as MeZO does not utilize gradient, its performance is worse than gradient-based optimizers.*\nCompared with MeZO, our 4-bit AdamW achieves better results on language model fine-tuning (including fine-tuning RoBERTa-Large on GLUE and SQuAD, GPT-2 Medium on E2E-NLG, and LLaMA-7B and LLaMA-13B on Alpaca). On the other hand, MeZO has performance gaps in some tasks compared with full parameter fine-tuning with Adam and it needs more steps to achieve strong performance.\n\nIn summary, LOMO and MeZO are designed for fine-tuning language models with strong simplifications of the optimizer (vanilla SGD / zeroth-order optimization) and *they may not be applicable for pretraining.* In contrast, our 4-bit AdamW is essentially an AdamW, and should have similar range of applicability with AdamW. \n\nFinally, we notice that [1-3] are all contemporaneous work with our submission. \n\n[1] Liu, Hong, et al. \"Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training.\" arXiv preprint arXiv:2305.14342 (2023).\n\n[2] Lv, Kai, et al. \"Full Parameter Fine-tuning for Large Language Models with Limited Resources.\" arXiv preprint arXiv:2306.09782 (2023).\n\n[3] Malladi, Sadhika, et al. \"Fine-Tuning Language Models with Just Forward Passes.\" arXiv preprint arXiv:2305.17333 (2023).\n"}, {"review_id": "YRKVYuU8U1", "reviewer": "Reviewer_qs7T", "summary": "The 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.", "questions": "- 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). ", "limitations": "Limitations are fully and honestly discussed.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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", "weaknesses": "- This is outstanding work. I cannot find any weakness.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer qs7T for the valuable and careful review.\n\nConsider a weight matrix $x$ of size $m \\times n$. In Rank-1 normalization, we store the absmax of each row and the absmax of each column during the quantization stage, totaling $m + n$ elements.\nIn the dequantization stage, we calculate the normalization constant for each entry in the same manner as in the quantization stage, i.e., min(absmax row_i, absmax col_j) for entry $x_{i, j}$. \nThe subtle difference is that we do not need to calculate the absmax of the i-th row (and the j-th colum) but use the pre-stored values instead.\nAdditionally, thank you for pointing out the typos concerning the first/second moment. We will correct them in a later revision."}, {"review_id": "J5qUudu8S8", "reviewer": "Reviewer_d65g", "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.", "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?\n- Are there any other ways to demonstrate the convergence (from proof) rather than showing empirical results?\n- 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer d65g for the careful review. With respect to the figures, we will use larger ones in a later revision. \nWith respect to your questions:\n\n**Weakness + Question 1: How to pick the block size.** \nThere are some rules to choose the block size: \n\n- Firstly, quantizing with a smaller block size approximates the tensors better than a larger block size due to the finer granularity of quantization.\n- Secondly, the block size is limited by the memory. A smaller block size would lead to larger memory overhead (e.g., with fp32 scaling constants, a block size of 32 causes ab extra overhead of 1-bit per parameter).  \n\nEmpirically, we have shown the quantization error of different block sizes via histograms in Appendix C.3 (Effectiveness of Block Size in Block-wise Normalization), which shows that a block size of 128 exactly reconstructs tensors, better than a block size of 2048. \nIn general, we consistently choose a small block size of 128 to reduce quantization error and maximize accuracy, while keeping memory overhead under control. \nActually, the granularity of quantization affected by block size is independent of the model size, and we suspect that a block size of 128 works well for unseen (large) models. \n\n**Question 2: Theoretical analysis about convergence of low-bit optimizers.**\nThe convergence of low-bit optimizers can be guaranteed if their fp32 counterparts converge. Here we provide a theorem about the convergence of quantized SGD with momentum (Algorithm 2 in Appendix F) under some assumptions. We believe that the convergence of low-bit AdamW could be inferred from the convergence of AdamW. \n\nWe use $f(\\theta)$ to denote the objective function, $\\beta$ to denote momentum, $Q$ to denote quantizer on first moment and $\\alpha$ to denote learning rate.\nWith assumptions:\n\n- (Convexity) The objective function is convex and has an unique global minimum $f(\\theta^*)$.\n- (Smoothness) The objective $f(\\theta)$ is continuous differentiable and $L$-smooth;\n- (Unbiasedness and Bounded variance) $\\mathbb{E}[\\nabla\\hat{f}(\\theta)] = \\nabla f(\\theta)$ and $ \\mathbb{E}\\left[\\left\\lVert\\nabla\\hat{f}(\\theta) - \\nabla f(\\theta)\\right\\rVert^2\\right] < \\sigma^2$, $\\forall \\theta \\in \\mathbb{R}^d$, where $\\nabla\\hat{f}$ is stochastic gradient. \n- (Unbiased quantizer)  $\\forall x\\in \\mathbb R^d$, $\\mathbb{E}\\left[Q(x)\\right]=x$.\n- (Bounded quantization variance) $\\forall x\\in \\mathbb R^d$, $\\mathbb{E}\\left[\\left\\lVert Q(x)-x\\right\\rVert^2\\right]\\le \\sigma_m^2$.\n\nWe have the following theorem:\n\n**Theorem 1.** Suppose above assumptions hold. \nConsider the case where quantized SGD with momentum is performed on the objective function $f$, with momentum parameter $\\beta$.\nLet $\\alpha \\in (0, \\frac{1-\\beta}{3L}]$, then for all $T > 0$ we have\n\n$$\n    \\mathbb{E}[f(\\bar{\\theta_T}) - f_*] \\le \\frac{1}{2T}\\left( \\frac{L\\beta}{1-\\beta} + \\frac{1-\\beta}{\\alpha } \\right) \\left\\lVert\\theta_0 - \\theta_*\\right\\rVert^2\n    + \\frac{\\alpha \\sigma^2}{(1-\\beta)} \n    + \\frac{\\alpha\\sigma_m^2}{2(1-\\beta)}\n$$\n\nwhere $\\bar{\\theta_T} = \\frac{1}{T}\\sum_{i=0}^{T-1}{\\theta}_{i}$.\n\nThe theorem indicates that when $\\alpha\\rightarrow 0$ and $\\alpha T\\rightarrow \\infty$, the optimizer converge to a stationary-point.\nNote that the unbiased quantizer assumption is only for technical simplification. \nOur proof mainly follows the prior analysis of SGD with momentum [1, 2].\n\n[1] Ghadimi, Euhanna, Hamid Reza Feyzmahdavian, and Mikael Johansson. \"Global convergence of the heavy-ball method for convex optimization.\" 2015 European control conference (ECC). IEEE, 2015.\n\n[2] Liu, Yanli, Yuan Gao, and Wotao Yin. \"An improved analysis of stochastic gradient descent with momentum.\" Advances in Neural Information Processing Systems 33 (2020): 18261-18271.\n\n**Weakness + Question 3: Results on generative modeling and more vision tasks.** We are running our optimizers on generative modeling and more vision tasks. We will post updates on these tasks within the discussion period. "}, {"review_id": "KOKYy8kOBq", "reviewer": "Reviewer_f2sV", "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.", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(1) The authors thoroughly analyze the outliers of momentums, which makes the existing block-wise method not work well.\n\n(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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer f2sV for the constructive comments. With respect to your questions:\n\n**Weakness: Common sense reasonging and MMLU evaluations on large models.**\nWe have fine-tuned LLaMA-7B and LLaMA-13B with 32-bit AdamW and 4-bit AdamW on Alpaca and evaluated them on common sense reasoning tasks and MMLU. Results show that 4-bit AdamW will not destory the capability of pretraiend models while enabling them to obtain instruction-following ability. In MMLU and HellaSwag tasks, 4-bit AdamW fine-tuning outperforms 32-bit AdamW both on LLaMA-7B and LLaMA-13B. In other tasks, 4-bit AdamW is comparable with 32-bit AdamW.\n4-bit AdamW does not get worse than 32-bit AdamW when the model size grows.\n\n|    LLaMA-7B | MMLU(5-shot) | HellaSwag | ARC-e | ARC-c | OBQA |\n|-------------|-------------:|----------:|------:|------:|-----:|\n| Original          |          33.1|       73.0|   52.4|   40.9|  42.4|\n| FT w. 32-bit AdamW|          38.7|       74.6|   61.5|   45.1|  43.4|\n| FT w. 4-bit AdamW |          38.9|       74.7|   61.2|   44.4|  43.0|\n\n|    LLaMA-13B| MMLU(5-shot) | HellaSwag | ARC-e | ARC-c | OBQA |\n|-------------|-------------:|----------:|------:|------:|-----:|\n| Original    |          47.4|       76.2|   59.8|   44.5|  42.0|\n| FT w. 32-bit AdamW|          46.5|       78.8|   63.6|   48.3|  45.2|\n| FT w. 4-bit AdamW |          47.4|       79.0|   64.1|   48.0|  45.2|\n"}]

Author rebuttal:

Meta-review: Solid paper on training with "truly" 4-bit optimizers. Most existing work on low-precision training still uses higher-prevision accumulators; this paper proposes a recipe to bypass that limitation. The reviewers all agreed on acceptance, and I think this work will be of great interest to the community.
Decision: Accept (spotlight)