Training Neural Networks with Low-Precision Model Memory
Authors: Bingrui Li
Ziteng Wang
Jianfei Chen
Jun Zhu
Venue: Submitted to ICLR 2023 (2023)
Keywords: memory efficient deep learning
stochastic gradient descent
quantization

Abstract: The demand for memory to store model-related statistics ("model memory") is a major bottleneck for training large neural networks. A promising solution is low-precision optimizers, which reduce the numerical precision of the model memory. However, existing work only compresses the momentum, resulting in suboptimal memory efficiency. This paper proposes Low-Precision Model Memory (LPMM), an optimization framework with the entire model memory kept in low precision. LPMM compresses not only the momentum but also model parameters and gradient accumulators. We identify arithmetic underflow as the main problem in building low-precision optimizers and propose a stochastic quantization method and a microbatching technique to overcome this problem. We analyze the convergence behavior of LPMM and theoretically show how the proposed techniques could affect underflowing, which in turn affects the convergence. We apply LPMM to the SGD optimizer with momentum (SGDM).  On several realistic benchmarks, LPMM-SGDM can train neural networks with negligible loss of accuracy while reducing over 70% of the model memory compared to the full-precision SGDM.

Under review as a conference paper at ICLR 2023
TRAINING NEURAL NETWORKS WITH LOW-PRECISION MODEL MEMORY
Anonymous authors Paper under double-bling review
ABSTRACT
The demand for memory store model-related statistics ("model memory") is major bottleneck for training large neural networks. A promising solution is low- precision optimizers, which reduce the numerical precision of the model memory. However, existing work only compresses the momentum, resulting in suboptimal memory efficiency. This paper proposes Low-Precision Model Memory (LPMM), an optimization framework with the entire model memory kept in low precision.
LPMM compresses not only the momentum but also model parameters and gradi- ent accumulators. We identify arithmetic underflow as the main problem in build- ing low-precision optimizers and propose stochastic quantization method and a microbatching technique to overcome this problem. We analyze the convergence behavior of LPMM and theoretically show how the proposed techniques could affect underflowing, which in turn affects the convergence. We apply LPMM to the SGD optimizer with momentum (SGDM). On several realistic benchmarks, LPMM-SGDM can train neural networks with negligible loss of accuracy while reducing over 70% of the model memory compared to the full-precision SGDM.
INTRODUCTION
Many huge models (Knnton & Toutanova, 2019 Radford et al. 2019 Dosovitskiy et al. 2021
Brown et al. 2020 Fedus et al., (002)) emerge in the recent several years. Despite being powerful, training these models is challenging. Memory is the main bottleneck in developing large models, as the training memory footprint is typically proportional the number of model parameters.
During training, the device memory consumed by three types objects:
1. data-related objects ("data memory"), including data and each layer's activation. Their size is proportional to the data size, i.e., mini-batch size and image resolution / sequence length.
redel-related objects ("model memory"), including model parameters, momentum, and gradi- ent accumulators. Their size is proportional to the amount of model parameters.
3. temporary objects ("workspace memory"), such scratch memory used by computing kernels and memory fragments.
Among the three types, model memory is the main bottleneck in scaling up machine learning mod- els (Rajbhandari et al. 2020).
Quantization is promising way of reducing the model memory. Specifically, low-precision optimiz- ers (Ramesh et al. 2021 Dettmers et al. (021) represent their states with low-precision numerical
formats, such bit integers, which consume less memory. Particularly. [Dettmers al. (202T) propose an 8-bit optimizer, which quantizes the momentum to block-wise 8-bit format. However, existing works have two limitations. First, the convergence behavior of low-precision optimizers thcoretically not well understood. Second, they only quantize the momentum, while model parame- ters and gradients are left in full precision. Therefore, the overall memory saving is unsatisfactory.
In this work, we propose LPMM, novel framework for optimizing with low-precision model mem- ory. Unlike previous works, LPMM consider all model-related objects, including model parameters, momentum, and gradients, in low precision. We identify arithmetic underflow the major bottle- neck of building low-precision optimizers. Theoretically, we analyze the convergence behavior of low-precision optimizers. Our analysis reveals how the design of the optimizer could impact the

---

Under review as a conference paper at ICLR 2023
degree of underflowing, which link the convergence behavior. Algorithmically, propose stochastic quantization and gradient accumulation methods to reduce underflowing. These tech- niques are backed with our theoretical findings. We further discuss the quantizer design and system implementation for LPMM.
We evaluate LPMM the standard image classification benchmark. LPMM can quantize gradients and the momentum to 8 bit, and model parameters to 12 bit, with neglible loss of accuracy. In total, LPMM only requires 28 bits of model memory per parameter, compared 72 bits by Dettmers al. 2021 or 96 bits of the full-precision algorithm.
2
RELATED WORK
Training Methods With Compressed Model Memory Several works train neural networks with
quantized parameters, momentum, or gradients. QSGD (Alistarh al. (017) quantizes the gradient into lower bits for efficient communication. Low-Precision SGD Li et al. (2017) Li De Sa (2019); De Sa (2018); Yang et al. (2019) uses low-precision parameters for training weight-quantized neural networks. The 8-bit Optimizer (Dettmers et al., 202T) quantizes the momentum for SGDM and Adam. Many of these works are not designed for saving the memory. Moreover, they only consider different parts of the model memory, without unified framework and analysis.
Other Methods to Reduce the Model Memory Sharding (Rajbhandari et al. 2020) splits model parameters across multiple nodes, which is only applicable to distributed settings. Offloading (Ren et al. 2021 Rajbhandari et al.) 2021 stores model-related statistics in CPU memory. These methods have heavy communication overhead, which can be reduced by our method.
Other Memory Efficient Training Methods Gradient checkpointing (Chen et al. 2016) reduces the storage overhead of activations by recomputation. Activation Compressed Training (Chen al., 2021 Liu et al.) 2022) keeps low-precision activations in memory. These method reduces the data memory rather than the model memory, and are orthogonal our approach.
Quantized Training There is series of work focusing the acceleration of neural network training with low low-pree -ww-rrcciionn computations (Wang et al.] 2018b) Micikevicius et al. 2017 Zhu et al.] 2020 (unn at. 2020). However, they still store the model-related objects in full-precision or 16-bit.
TRAINING WITH NOW-PRECISION MODEL MEMORY
In this section, formulate the problem of training neural networks with low precision memory. We show that arithmetic underflow is the major bottleneck for reducing the numerical precision. To solve this problem, propose stochastic quantization and micro-batching techniques. Here, mainly consider stochastic gradient descent with momentum (SGDM) (Qian 1999; Sutskever et al.) 2013) as motivating example, but the proposed techniques also apply to other optimizers, such as stochastic gradient descent (Bottou, 2010) and Adam (Kingma & Ba 2015).
3.1 BASIC OPTIMIZATION FRAMEWORK
Consider the empirical risk minimization problem
min (())
((O),
where 0 is the model parameter. Since the dataset size is large, stochastic optimizers are adopted for solving the problem. SGDM is one of the most widely used optimizers. Starting with an initial model fp and momentum mg 0, SGDM performs the following updates:
m: 8m,11 f((0-1)) 0. + 0 am..
where Vf(0,-1) is an unbiased estimator of the gradient. We define ==(--))V(((LL..) = T Eiece (0,-), where Ç* is minibatch sampled uniformly from the dataset (1,2,...,71}
2

---

Under review as a conference paper at ICLR 2023
Gradient Accumulation
9tu
Network
Aru
6--s
P Parameter Upcate
%
Momentum pdatee
me-1
1%
Figure 1: Dataflow of the LPMM algorithm. Square nodes are low- low-precision and kept persistent in memory. Circular nodes are temporary full-precision data, which only consume small constant- size (dimension-free) buffer.
When training large neural networks, the dimensionality of 0 is huge. Storing the model-related statisties, including the model the momentum m, and the gradient Vf(0) can be major mem- ory overhead. We call these statistics collectively as the model memory. For example, the GPT- 2 (Radford al. 2019) model has 1.5 billion parameters. During training, each parameter requires (modell + 32 grad.) + 32 (momentum) bits or 12 bytes of model memory. Therefore, the model memory alone takes about 1.5 10 x 12bytes 17.6GB of GPU memory, which already exceeds the capacity of many mainstream GPUs.
To reduce training memory footprint, we can store the model-related statistics in low numerical precision. This involves quantizer Q(-), which encodes full full-precision (32-bit) values to lower- precision ones, such as 8-bit integers. The training procedure can be summarized as:
9: Q (Vi(0.... m, + 1o @Sm- Vf(0.-1)). + ((t-1- am,). (2)
Here, Sr is a gradient accumulator, which stores the temporary gradients during back-propagation. When the parameters and momentum are updated, they are first decoded to full-precision, computed in full- precision, and then encoded and stored in low-precision, as illustrated in Fig. [ further implementation details to Sec. 5.3
We defer
3.2 THE UNDERFLOOW PROBLEM
Quantization is lossy, which can only represent values up to a numerical precision. Consider the following symmetric linear quantizer Q Q5, R - RR as example:
Q5,((z)= = 8 Round (clip (z/ō, -B B)),
where (ound(.) (.) is a rounding function and clip(,, -B B) clamps the input into the range [-B, B]. The quantizer Q& (x:) converts to quantities within the range [-8B, up to the numerical precision The quantized value can be stored as a log, (2B)-bit integer. The quantizer can be generalized vector inputs and outputs by applying element to each entry.
Arithmetic underflow occurs when we try to sum up two quantities of different orders of magnitude. Consider the summation
0, = Qs (0,-1 + vz),
(4)
where O, is the parameter at the 6-th iteration, and 12 -am, is the update vector. Since 0,-1 itself is already in low-precision, we have (0,-1) 0,-1- Now, if the update is smaller than parameters' numerical precision, f.e. |vral < 8/2 for some dimension d, the small update will be "eaten" by the large parameter, treated effectively as zero:
0rd Q8. 10- vt,)) 0, 1,d-
(5)
3

---

Under review as a conference paper at ICLR 2023
12
4-
1.0
0.25
Bin stangss ca0
0.25
0.6
2-
0.4
0.19
0.05
0.8
0.09
Epochs
Epachs
(a) 12-bit
(b) 10-bit
(c) 8-bit
Figure 2: The magnitude of parameter update for convolution layer in ResNet-18. We vary the
parameter precision from 12-bit to 8-bit. Gradients and momentum are always in 8-bit. The "bin change ratio" is defined as 0 averaged over all dimensions. Larger is better.
The update is likely to be small, particularly at the late stage of training where the gradient (and momentum) is close zero, and the learning is also decreased around zero. If underflowing occurs, the actual update direction ö ;= 0, 0,- is far from the expected. the Extremely, there could be 0g = where the optimization completely stalls.
As shown in Fig. when the underflow problem is serious (fewer bits), the deterministic quantizer has the smallest actual update 0,/6. Actually, the deterministic quantizer hardly updates the param- eters. As shown in Fig. [3a] about 60% entries in the update vector is less than 6/2. These entries are ignored by the deterministic quantizer.
3.3 SOLUTIONS TO THE UNDERFLOW PROBLEM
We propose two methods for solving the underflow problem: unbiased stochastic quantization and gradient accumulation.
Unbiased Stochastic Quantization Stochastic quantization adopts an stochastic quantizer Q(-), satisfying E [Q(z)] = x. A simplest stochastic quantizer is the following:
(5.((z) = 5 SR (clip (z/6, -B, B))
(6)
where SR(-) is the stochastic rounding (Gupta et al. 2015
Courbariaux al. 2015) operation:
SR(x) [x: with probability [a]] otherwise [x]
The quantizer is unbiased as long as oB 2 [x]] Stochastic quantization mitigates underflow by accumulating small updates in a probabilistic manner. When the update is small in magnitude, each dimension only has a small probability to be updated. Yet after a large number of updates, the final parameters still have the same expectation with the full-precision ones. Besides this informal argument, unbiased stochastic quantization is also the key to establish the convergence theory, we shall see soon
Sec. 4.2
Microbatching Underflowing arises from the difference in magnitude of the parameter and the update. Therefore, it can be mitigated by enlarging the update. Microbatching (Huang et al.) 2019) works by summing up the gradient of multiple batches before applying the update to the momentum and parameters. More concretely, we can accumulate the gradient of N mini-batchos as
9t.0 t
St. + Q(g.,-1++ (f(0
and use It,N in place of 9: for updating the momentum in Eq. 2 Now, the accumulated gradient gr.v is N times larger than the original 9t. Subsequently, the update -am; becomes times larger, yet the magnitude of the parameter 0, is unchanged, underflowing is significantly alleviated.
As shown in Fig. Ba), about 20% entries of update vector 2 are even smaller than 8/8. this case, vanilla stochastic quantizer still suffer from high variance due to the repeated small updates.
4

---

Under review as a conference paper at ICLR 2023
0.25
1.0
0.33
3.5
05
0.15
0.30
0.05
52
0.00
9.9
0.03
0.
88 680
341
(0%, 80
10 191
(a) Micro-batch size
(b) Micro-batch size
Figure 3: The distribution of update vector magnitude across dimensions. The result is counted at the epoch 10 for single layer in ResNet-18. The blue bar (left Y-axis) represents the density and the red line (right Y-axis) indicates the cumulative distribution.
Microbatching significantly enlarge the size the update, and reduce the number of small updates (which underflow), as shown in Fig. Bb).
Micro-batching can be viewed equivalently large-minibatch SGD (Goyal et al., 2017). which simply multiply the batch size and the learning rate by N. Large-minibatch SGD known to converge similarly the original, mnill-minibatch SGD. However, enlarging the batch size is not memory efficient since the data memory will increase.
Putting all together, we propose LPMM-SGDM, an optimizer with low- oow-prccision gradient accu- mulators, momentum, and parameters. The entire training procedure is given by Algorithm [[] Fig, and show that our algorithm effectively mitigate underflow with significantly larger updates.
Algorithm L.MMM-GGMM
Require: initial parameter 6 E RP, learning rate a, initial momentum my 0 E R2, quantization function Q total number of iterations T, momentum factor B, gradient accumulation steps N for 1,,2.... 1.., do
2: 9t.0 4 0
3: for 1, ...,,
4: Sample minibatch ed uniformly from the training data;
5: gt,s + Qg + (f(0--1
6: end for
7: m. + -03mm-1++++ + (t,
8: 0, < Q0,-1 am,))
9: end for
return 0r
CONVERGENCE THEORY
In this section, analyze the convergence of the proposed LPMM-SGDM, algorithm and discuss the impact of the proposed techniques.
4.1 ASSUMPTIONS
First, make some assumptions. The first three rather standard in stochastic optimization literature, while last two depict properties of stochastic quantizers.
(Convexity) The objective function convex and has an unique global minimum f(
2. (Smoothness) The objective f00 is continuous differentiable and T-smooth;

---

Under review as a conference paper at ICLR 2023
3. (Bounded variance) E - ~ < 60 ER.
4. (Unbiased quantizer) Voc E TR -B(g(()) Fm (3)] =
[Qb()))
5. (Bounded quantization variance) Vx E R.. [1pp(z) - ==1"] S op. (1) - -x12² S ozz
EE [1)
x=1²] o
Here, Assumption depicts the variance of sampling mini-batches, and Assumption bounds quan- tization error of model-related statistics through the variance. Assumptions 4-5 can be realized by
linear quantizer Eq. (6), where we have [[)),n()) =1 S L Var E Rdd Yet, our anal-
ysis is applicable other non-linear quantizers such as those discussed in Sec. 5.1 We assume that the objective function is convex, which somewhat restrictive. However, this simple convex case is sufficient to reflect the underflowing challenge of training with low-precision model memory, and we leave the non-convex case as future study.
4.2 CONVERGENCE OF LPMM-SGDM
We first analyze the bias and variance of the quantized gradient accumulator. Lemma 1. In Algorithm 7 with Assumptions 3-5, the first and second momentum of 9t. satisfy
0 [00k-1]iV[[i (0,-))
(7)
E 10.-1] i + i(i + 2) (
(8)
Next, we construct an auxiliary sequence { following previous analysis of heavy ball meth- ods (Ghadimi et al., 2015) and SGDM (Liu et al. 2020):
B
24 0g. when 0, otherwise
1 B 0
81-1.
(9)
The sequence enjoys the following nice properties:
Lemma 2. If Assumptions 3-5 hold, then sequence {2t) satisfies
1
B
2++1 21 1
(0c+1 0i) 1 3 (0: 12-1)
(10)
E E[2441 - 22l = -a.N
1 f(i.)
(11)
3 )
(21111 (ellll++1, + oz + 1 2,32 (1 - B) og
B B[]lz+1 2el195
(12)
The quantity E[1 zel is composed of four parts: the batch gradient norm, the variance of stochastic gradient, momentum quantization and parameter quantization error. The latter three noise terms sum up independently. Our analysis applies to cases when only some of the model-related statistics are quantized by setting corresponding noise terms to zero. With this lemma, we derive the following theorem characterizing the convergence of LPMM-SGDM.
Theorem 1. Consider the Algorithm we have
1-8 I with Assumptions 1-5, Let a (0, (N+2jrl then for all
1
LB
1-
Blf(Or)
f+] <
1100 1100-0,11 -0,12
(13)
27'
1 - 3
aN
ag2
ao?
ao
(1 (1+282)02
(14)
(1 - 3)
(1 3)
2N(1
3)
2o.N(1 -
where 0r 1. ET-10 0.
Theorem[ the effectiveness of the proposed techniques in Sec. 3.3 First, LPMM can con- verge the global optimum when T oc, 0, and a.N co for any quantizer with non-zero

---

Under review as a conference paper at ICLR 2023
Op. Gg. and om- That is, the ultimate convergence is guaranteed regardless of the numerical preci- sion. In contrast, optimizers with deterministic quantizers could stall at some intermediate state, as discussed in Sec. 3.3 Therefore, stochastic quantization is essential for the convergence guarantee.
Second, we can study the impact of gradient accumulation. With sufficiently large T, the first term in Eq. (14) vanishes, and we therefore only consider the noise terms. We compare plain algorithm (N = 1) and a version with gradient accumulation (N > 1), with equal learning rate. Gradient accumulation reduces the parameter quantization error (the Tp term) by N times. Meanwhile, the gradient accumulator stores the gradient sum over N mini-batches, so og becomes larger. Therefore, gradient accumulation trades gradient quantization error for parameter quantization error. The error bound is tighter with a large N when parameter quantization error is more significant, which holds true according to our empirical observations in Fig. 2 and
IMPLEMENTATION DETAILS
In this section, we discuss our quantizer design and the implementation of the whole training frame- work.
5.1 GROUP AND NONLINEAR QUANTIZATION
For the quantizer, we use the nonlinear quantization on momentum and linear quantization on pa- rameters and gradients. All quantizers use the standard absmax symmetric quantization. The linear quantization defined as Eq. 6 The nonlinear quantization method (Dettmers) 2015) proposes a novel floating-poit format in which the bitwidth for the exponent is dynamic. Furthermore, given tensor to be quantized, we first use the absmax to normalize the tensor into [-1, 1]. Then, the quan- tizer finds the bucket for each element with a binary search on precomputed quantization map. Finally, the quantizer projects the value to one the adjacent quantization points by stochastic rounding. Since stochastic rounding is unbiased, this nonlinear quantization method is also unbi- ased, which indicates that all the analysis above holds except the quantization variance is intrinsically different.
Group quantization is also adopted to avoid the negative impact of outliers. We chunk every tensor into groups of size 2048, and every group has its own quantization precision 6. However, group quantization introduces more data to be stored and as group size decreases, the quantization is more accurate while the extra data cost will be higher. We choose the group size of 2048 balance the numerical precision and storage cost.
5.2 DYNAMIC QUANTIZATION
We experimentally observed that the magnitude parameters shows a trend of first increasing and then decreasing. If smaller quantization granularity â is adopted, then the overflow problems occur. Especially at the initial stage of optimization, the quantizer rudely truncated the parameters to be increased due to limited range, which leads biased quantizer and makes the solution suboptimal. However, if we quantize with an overly large 8, it results in poor numerical precision, and the underflow problem worsens. Therefore, take dynamic quantization policy to balance the overflow and underflow. Specifically, it calculates the quantization precision 5 with the quantizers' inputs at every iteration, which gives minimal lossless 6.
5.3 SYSTEM IMPLEMENTATION
Here we discuss how our algorithm is implemented. As illustrated in Fig. the forward, backward, and optimizer steps are all depicted straightforwardly. Before an iteration begins, we store all model memory in low precision. When the new input comes, the parameters needed for forward computa- tion are dequantized into a temporary full-precision "copy". Note that the parameters for each layer are not required simutacnously, so we only dequantize parameters right before its usage. As long as the computation related to the full-precision "copy" has finished, we delete these tensors. Therefore, the overhead of the temporary full- full-precision parameter is no larger than the largest single parameter tensor. In the backpropagation and optimizer step, a similar computation pattern happens. Note that the memory of gradient accumulators cannot be saved by not storing the gradient accumulator even
7

---

Under review as a conference paper at ICLR 2023
Table 1: Performance of ResNet-18 and ResNet50 on ImageNet.
Model
Method
Bits (param./grad//moment))
Accuracy(%)
ResNet-18
Full-Precsion
32/32/32
71.1
ResNet-18
LPMM
12/8/8
71.0
ResNet-50
Full-Precsion
32/32/32
77.2
ResNet-50
LPMM
12/8/8
77.3
if we do not apply microbatching (i.e. N = 1). On the one hand, the gradient may not be calcu- lated in one operation (or one layer). On the other hand, gradient computation can be composed of multiple accumulations, such as in convolution layer.
Since the computation in neural network training is always layerwise, then at most, only the full- precision parameters, gradients, and momentum associated with one single layer are kept in memory at the same time, while other model memory is inactive and in low-precision.
Based on this philosophy, implemented a training algorithm PyTorrc. We quantize the full- precision to low-precision representation by manually compressing them into byte stream using the CUDA kernel. Our system includes lot of operators and layers to perform desired forward and backward calculations. Given defined model, we quantize the parameter tensor individually, replace the original full-precision parameter and then move GPU, Then it can convert the layers of the model automatically by our customized layer. Besides, to get the gradient and perform com- putation more conveniently, we manually design gradient accumulator instead of using the original "grad" attribute.
6 EXPERIMENTS
Experimental Setup We compare the performance of LPMM to the full-precision counterparts on the image classification task. The full-precision baseline uses the standard SGDM optimizer (Qian 1999 [Sutskever et al. 2013). and we use our proposed LPMM-SGDD algorithm. We experiment with R8 and Re (He et al. 2016) on ImageNet (Deng et al. (009). We do not change any existing hyperparameters or precision of input and activations for each experimental run. The only change to replace full- full-precision parameters, gradients, and momentum with the low-precision ones. We conduct ResNet-18 experiments on RTX 3090Ti and ResNet-50 experiments on A100 GPUs.
The accuracy result is shown in Tab. ! LPMM-SGDM can quantize gradients and the momentum to 8-bit and model parameters to 12-bit, with no accuracy loss compared with the full precision SGDM.
Ablation Study To study the effectiveness of our proposed techniques, we conduct ablation analysis on ResNet-18 on LPMM under 12-bit, 10-bit, and 8-bit parameter quantization. We consis- tently use 8-bit gradients and 8-bit momentum. The result is shown in Tab. 2 The baseline uses deterministic dynamic quantizer (rounding to the nearest), and no microbatching is applied. We then apply stochastic quantization and microbatching sequentially. The performance with microbatching is the optimal result across different micro-batch size, which N = 8 for 12-bit and N = 16 for 10-bit and 8-bit. The results show that stochastic rounding and microbatching are critical for the performance of LPMM. The 8-bit runs collapsed without the unbiasedness guarantee. While as the bit goes up. the reduced quantization error compensates for this bias. Mierobatching helpful at all quantization level and make the 12-bit parameter LPMM match with full full-precision.
Role of Stochastic Quantization We compare the effect of stochastic rounding on LPMM-SGDM across 12-bit, 10-bit, and 8-bit parameter quantization. In these experiments, "deterministic" means all quantization operations use the "rounding to the nearest" quantizer, while "stochastic" means all quantization operations are stochastic and no microbatching is applied. The result is shown in Fig. 4 For the 12-bit setting, we can see that the training trajectories overlap in the early stage of optimization since the update is large enough compared to the numerical precision 8, and underflow- ing is not significant. Nevertheless, the deterministic method degrades from the middle stage due to the smaller update vector, which leads to underflowing. The degradation occurs earlier for 10-bit
8

---

Under review as a conference paper at ICLR 2023
Table 2: Ablation study of LPMM-SGDM on ResNet-18 on ImageNet.
Method
Bits (param//grad///oment))
Baseline
+Stochastic
+Microbatching
Precissre
32/32/32
71.1
LPMM
12/8/8
66.0
70.7
71.0
LPMM
1 10/8/8
53.4
69.7
70.6
LPMM
8/8/8
15.8
65.6
69.2
stchaskk
dctemiriski
stochasti:
datemerédk Rik.
6.0-
UNN
UPNH
Rddppett sonn
helppeccias
4.0
4.-
4.5-
stccasiic
LPNN
detarmiriat:
3.5
3.5
3.5
30
3.0-
255
25-
2.9
20
(a) LPMM w. 12-b param.
(b) LPMM w. 10-b param.
(c) LPMM w. 8-b param.
Figure 4: Convergence curves of compared algorithms.
and 8-bit settings since the quantization precision is larger, and underflow can happen more easily. This indicates that unbiased stochastic quantization is the crucial factor in the underflow problem.
Role of Micro-batching Here, we study the impact of the micro-batch size N. In these experi- ments, we use a micro-batch size of 256 and an initial learning rate 0.256. Other hyperparameters such warmup and weight decay set according to the actual batch size (i.e., 256NN. The result shown Tab. 3] When the parameter quantization error large, particularly in 8-bit setting, effectively while the increased gradient
microbatching mitigates the parameter quantization error is still tolerable. We can see that microbatching exactly plays a role on reducing parameter quantization error since the improvement 8-bit setting is largest. The improvement 12-bit is marginal since the parameter quantization error is already small.
Table 3: Effect of the micro- batto--bachhhoo size.
Micro-Batch Size(N)
4
16
Full-Precision
71.1
LPMM (12-b param. 8-b grad. / /8bb moment.)
70.7
70.0
70.5
71.0
70.9
LPMM (10-b param. / 88-b grad. 8-bb moment.)
69.7
69.0
70.1
70.3
70.6
LPMM (8-b param. / /-b grad. 8-b moment.)
65.6
65.9
68.1
68.8
69.4
7 CONCLUSIONS
We present LPMM, an optimization framework for training neural networks with the entire model memory kept in low-precision. We apply LPMM to the SGDM optimizer and give convergence guar- antees. We discover the critical underflow problem in low-precision training and propose two meth- ods: unbiased stochastic quantization and microbatching, the solution for this problem. Moreover, we conduct a comprehensive study on the effect of stochastic and microbatching theoretically and empirically. On the image classification task, LPMM achieves accuracy with negligible loss ing 12-bit parameters, 8-bit gradients, and 8-bit momentum, which reduce over 70% model memory compared with the full-precision training.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Communication-effcienn sgd via gradient quantization and encoding. 2017.
Qsgd:
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. International Conference on Machine Learning, pp. 560-569. PMLR, 2018.
Léon Bottou. Large-scate scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT pp. 177-186. Springer, 2010.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Blerbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in neural information processing systems, 2020.
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Actnn: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, pp. 1803-1813. PMLR, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.661444444 74, 2016.
Xi Chen, Xiaolin Hu, Hucheng Zhou, and Ningyi Xu. Fxpnet: Training a deep convolutional neural network in fixed-por tried-ponntioe representation. In 2017 International Joint Conference Neural Networks (JJCN), pp. 2494-2501. IEEE, 2017.
Matthieu Courbariaux, Yoshua Bengio, and networks with binary weights during propagations. In Advances in neural information processing systems, pp. 3123-3131, 2015.
David. Binaryconnect: Training deep neural
Christopher De Sa, Megan Leszezynski, Jian Zhang, Alana Marzoev, Christopher Aberger, Kunle Olukotun, and Christopher Ré. High-accuracy low-precision training. arXiv preprint arXiv: /8X3:0800333383 2018
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, 2009.
Tim Dettmers. 8-bit approximations for parallelism in deep learning. arXiv preprint arXiv:15/100461. 2015.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An im- age is worth 16x 16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Representations, 2021.
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.
Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global convergence of the baav method for convex optimization. In 2015 European control conference (ECC), pp. 310-315. IEEE, 2015.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: training imagenet in 1 hour. arXiv preprint 17X.v 2017.
10

---

Under review as a conference paper at ICLR 2023
Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In International conference on machine learning, pp. 1737-1746. PMLR. 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, 2016.
Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. Bert:
Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLI, pp.pHH.ppppp.1111111 4186, 2019.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Goldstein. Training quantized nets: A deeper understanding. 2017.
Zheng Li and Christopher M De Sa. Dimension-fie free bounds for low-precision training. In Advances in neural information processing systems, 2019.
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.1 1357, 2022.
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. In Advances in Neural Information Processing Systems, volume 33, pp. 18261- 18271, 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, al. Mixed precision training. arXiv preprint arXiv:1710.037400 2017.
Y. Nesterov, Introductory lectures convex optimization: A basic course (applied optimization). Université catholique de Louvain, Center for Operations Research and Econometrics (CORE), (1), 2013.
Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks, 12 (1):145-151) 1999.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-16. IEFF, 2020.
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 Inter- national Conference for High Performance Computing, Networking, Storage and Analysis, pp- 1-14, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pp. PMLR, 2021.
11

---

Under review as a conference paper at ICLR 2023
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Min- jia Zhang, Dong Li, and Yuxiong He. (ZeRO-Offloadl) Democratizing (Billion-Seae model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Xiao Sun, Naigang Wang, Chia-Yu Chen, Jiamin Ni, Ankur Agrawal, Xiaodong Cui, Swagath Venkataramani, Kaoutar El Maghraoui, Vijayalakshmi Viji Srinivasan, and Kailash Gopalakr- ishnan. Ultra-low precision 4-bit training of deep neural networks. Advances in Neural Infor- mation Processing Systems, volume 33, pp. 1796-1807, 2020.
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, pp. 1139-1 1147. PMLR, 2013.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 074611 2018a.
Naigang Wang, Jungwook Choi, Daniel Brand, Chia- Yu Chen, and Kailash Gopalakrishnan. Train- ing deep neural networks with 8-bit floating point numbers. In Advances in neural information processing systems, volume 31, 2b.
Guandao Yang, Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning, pp. 7015-7024. PMLR, 2019.
Yukuan Yang, Lei Deng, Shuang Wu, Tianyi Yan, Yuan Xie, and Guoqi Li. Training high- performance and large-scale deep neural networks with full 8-bit integers. Neural Networks, 125:70 82, 2020.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Train- ing low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint 100..6600.660 2016.
Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang, Zhelong Li, Xiuqi Yang, and Junjie Yan. Towards unified int8 training for convolutional neural network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1969-1979, 2020.
12

---

Under review as a conference paper at ICLR 2023
A PROOF OF LEMMA AND THEOREM
A.1 PROOF OF LEMMA 1
Proof. When i 1, by Assumption 4 and the due to V f(0) is unbiased, get
[0r,1 = Q(V]((--1))
=E vf(0--1)
And when i >
we have
19.,1 00:]]= =E Q(9 +
=E (0r.i-1 Vi(0. 1))
= -E 1 + E V (((
[91,5-1 8: + (8(1))
By reduction it gives the first part of Lemma
E [gc, -V]((
With Assumption 3, it holds true that
Vi(0) |||
(V(()) VS(()) + |
E
E
E
Vi(00- V
+ E |V - 2E
V - V5(0), vS >
E
||(( (
+ E |
< +0
For the second part of Lemma 1, when 1,
9t.1 - v)(0,-1) + vf(0 2 10
[ | = E
E
9.1 - 9.1-
v)(0-1)
17
+ E
V (((0-1)
2E
9t,1 77(0--1) >
= E
91.1 v((0.-1)
+ E
[f((
< 39 [0
+
13

---

Under review as a conference paper at ICLR 2023
And when i >> 1, if this equation holds true for 1 to i 1, then for i we have
E [1s 100-1]
IE
[|00. 91,6-1 100-1-v(00-1)
11] +E 1-BB|oo.--
+Vi0-+++((0+)))) 19]
< 91, 9t,i-1 (1.00-1),00.1111 2E
vf(0--1)
11/
+E |1,i-1 vf(0)-1) 1|]
EE
00.i -1.1-1 Vi(0,-1)
o + E
|19..-111 11] E
V f((01-1)
+ < 91,1-1, vf(0--1)
of E |19.5-11 [1]] E v (( 11
12]
+ ([((0-1))
S - + 1)0g + (1 1) + 1)
+ 1)o² + +1V5(0-=1)1+ +
i (6 + 2: - [V/ +00
s i + i(i+2) 2) ( +
So by reduction we complete the second part.
A.2 PROOF OF LEMMA 2
Proof. By definition of zt, we have the first equation immediately. With it we get
1
8 EE 10 0.-1]
E [241 21]
18 [0+11 Oc]
1-3
1- 3
By Algorithm 1 we have
E [0++1 - 01] = E [0c+1 - (0L am1++1] [cmm++)]
-aE [m++1]
-.aE [8m, - -9+11N]
--aBE [me] NNN
and
E 00 - 01-1] = [0, (0,-1 am,) EE [am]]
- [me]]
Substitute them derive the second equation
ag
aN
+ a3
E [2441
zz] =
E [me]]
Vf()))
E [m)]
1-
1- 3
1- 3
-aN
(((e.)
33
As for the last part, since
1
Z++1 2t 1- 9
(0 - -0)) 1- (0: - 0,-1)
(A++1 (0: amm+1))
(0: (0,-1 - (0-11-am:)) 1 -B (m++1 8mm)
B
14

---

Under review as a conference paper at ICLR 2023
Take expectation and we have
27
||2441 - 20ll?
= E
1 B
(0 - (0: - am)+1)) -
1 1- I-3( (0, - (01-1-ami -am)) - -1 3 (mi+1 - 3mz)
-
a
E
|10++1 (0: am+++))l
+ (1 -32
E [10: - (00-1 101-1-000 -ama)l
(1 - B)²
(1 1--3 3
2
||mm+11 - Smel? 2
2
E K< 0s+1 -(0: am1+1), 1,, (0-11),,,0000-00000 (0s-1 ams) >]
(1 1-3
2a
E K 0,11 - (0. am 1), M 33m4+ >]
(1 (1-8)²
2003
E E| 0, -O0t11 omm)), m441 3mm >]
(1 -)3
1 1 B8
(
<
(1 B)2 oz +
11 3 (B [llm-1 (Bme + 9++1,))1
+ E [10+1
2T [< m (Bmm + I++1,N) 0C+1, >))
a 2
[P[ + +om) + 1 32
(1 3)2 ot
B
which is the last part of Lemma 2.
A.3 PROOF OF THEOREM 1
Proof. From Lemma 2, we have
||=+11 - 2: - (1 1 3)
a ) (2E [I0++IN]]+++
+ +0²)+ + 1 2,2
E
(1 - =3)2 o
And by Lemma 1, we know
E [[IttINI]]] S NoO + + 2 2V No²
Substituting it gives
2 ( (-3) ) )² (2N(N+2) 2V) + 2No + 2No + +
+ 1 2,3
|| - 20117] S
(1 B)2 o
(15)
Suppose 0, is the optimal parameter and f. = f(0.) is the minimal objective value. First, we have ||ž 0_ 2 || 0= + 2 (z+- - Or, 2111 - 21) + || - zell?
Take expectation over the randomness in the (t+1)-th step, we have
2a
El 0.11 = 1| -0_11
1- 3 (z+ 0., v.((O))) + E[[l2++1 - 2x
NaN
-11E,- - 11 1 B
(0, 0., vf(,))
2aB3N
(1 - 9-)) (0, - 0,-1, v ((O))) + E|| - -
Since f is continuously differentiable and L-smooth, we have the following inequalities. Nesterov 2013
(0. 0., vS((,)) 2 [V
(16)
1
(0, 10,-0, ((())) 2 f(0.) - -f. + 21 [Vf(22)]
(17)
(0. 6, V f(())) 2 (0.) f(0(-1)
(18)
15

---

Under review as a conference paper at ICLR 2023
Substitute them and get
20 - p) |1 [ 20,N (f(0t) f.)
E| 2b+1 - -0.119 1|22 - 0.11
L(1 (1-33 3)
1 3
a.N L(1 B) [Vf(0))) (1 81- (f(Ot) - -(((-1)) + 2 - =ell"]
2
where p E (0, 1] is a parameter we use to balance the first two inequalities. For the last term, denote
M- (2No² + 2No + No +
1+28² (1-87T0p- Substitute 15 into the last inequality and
collect the terms, we get
2 + 1-8 (1- (1-8) 3)2 (f(O.) - f.)) + +.)+BBllZe B1 - 0.1121
No.NB
(1 8)2 (f(0)-1) (.) + || - -0,12+ 22a N(N + 2)
2
a2N(2-P) p)
a2 -0-
|V + M
(1 B)2
1(1 - (3)
When a satisfies the condition N(N+(
a2(2-p) I(1-3)
< i.e. 0 a (1-3)(2-))
the term about
[I-p)²
2(N+2T
|V is non-positive, thus we have
2aNp p +
2aN3
31 3/2 (S(O.) f-) + B[ 0.11
1-
2aNB 8
< (1 - B)2 (f(0,-1) fa) + |2 - -1 + M
Summing this inequality from 0 to T - 1 and taking full expectation gives
NaN
p
T-1
E(f(0.) - S.] +
T 2a 3
((S S.] + B|| - -0.119)
1
(1 - B)²
1=
1=
T-1
2aN3
By² [[5(0=-1) - [-] Billz 0.119) +T.M
which implies that
2aN
E(f(0.) - f.] < 2aN3
(f(Oo) - f.) + 1100 - 1111 + T M (1- 3)2
1=0
Since f is convex, we have (T) < T
-1' 1= ((0:)). Subsequently we have
1
B
1 3 1100 0.11 B) (f( f-)
B]f(Or) f.1
p(1
2aN p
3 M
+
2ar
Finally, when a E (0, which gives
1-8
can take p use L-smooth condition again and substitute M.
(2
L
1 - 18
B(((rr) - fa] S
2T
I 1-3
+
aN
110 10.1
ag²
ag²
a m
(1 23+)
3)
(1
3)
2N(1 - 3) 2aN 3)
It's noted that when N = 1, we get the basic quantized SGDM without gradient accumulation, whose convex convergence result is
1
LB
1-9
B(((Tr)
f.1 <
1100 0.11
2T
8
ag²
a0²
a0
(1 + 28 20 8) +
+
+
(1
3)
(1
B)
2(1 - -8)
16

---

Under review as a conference paper at ICLR 2023
B EXPERIMENTS ON TRANSFORMERS
Tab. contains the accuracy of RoBERTa-Large Large et al. 2019) under LPMM framework on GLUB Wang al. 201 dataset. Each column is the median performance metrics (accuracy, Fl, etc.) across 5 random seeds. Similar to LPMM-SGDM, the LPMM-Adam adopt linear quantization on parameters and gradients, nonlinear quantization on momentum and square momentum. The results show that consistent 8-bit LPMM-Adam could achieve (nearly) lossless accuracy compared with full-precision reference.
Table 4: Performance of RoBERTa-Large GLUE Tasks. LPMM(12/8/8/5) means that 12-bit parameter, 8-bit gradient, 8-bit momentum, 8-bit square momentum are used.
Model
MNLI
QNLI
QQP
RTE
SST-2
MRPC
CoLA
STS-B
Avg.
Full-precision
90.2
94.7
92.2
86.6
96.4
90.9
68.0
92.4
88.9
LPMMM(12/8888)
90.5
94.6
91.9
84.8
96.4
90.7
67.1
92.3
88.5
LPMM((/8/8/8)
90.4
94.6
91.9
85.6
96.4
90.0
64.6
92.1
88.2
C MEMORY SAVING AND COMPUTATIONAL OVERHEAD
In this section, we investigate memory saving and the overhead of LPMM. Tab. shows the memory saving and throughput on several benchmarks. The model memory here is the peak memory, where all parameters, gradients and optimizer states stay in memory. All measurements are performed under standard benchmarks, i.e. ImageNet classification for ResNet-50 and GLUE for RoBERTa- Large, and standard training configurations for all benchmarks. Our LPMM-Adam saves up to 3.8 GB GPU memory (about 72%) for RoBERTa-Large and saves 62% for ResNet-50, which make the training(or finetuning) of large models more accessible. Besides, LPMM would lead to a compuation overhead especially when the model large, and the 8-bit Optimizer Dettmers et all, 2021) has higher throughput than the baseline. This is due to Dettmers al. (2021) use operator fusion techniques to integrate the optimizer update, quantization into one single operator as CUDA kernel, while LPMM CUDA only for quantization and PyTorch for other operations. This problem could be mitigated if the LPMM framwork is completely implemented in CUDA.
In Tab. 66] investigate the largest trainable model in single GPU with LPMM. We use a batch size of one and vary the depth and width respectively to explore the limit size. The GPU memory is consistently 11 GB for this comparison. Tab. and both show LPMM could attain a higher reduction rate for language tasks than vision tasks. This is because activation accounts more memory vision tasks.
memory under
Table 5: Saved Memory and Throughput on Various Tasks.
Method
Model
Model Mem
Saved Mem
Throughput(
32-bit Momentum
ResNet-50
299.1 I MB
0.0 MB(0%)
347,86
8-bit (Dettmers et 2021
ResNet-50
224.4 MB
74.7 MB(25%)
348.71
Momennumm 27878)
ResNet-50
112.6 MB
186.5 MB(62%)
347.71
32-bit Adam
RoBERTa-Large
5.3 GB
0.0 GB(0%)
186.2
8-bit (Dettmers et al.
2021
RoBERTa--Larg
3.3 GB
2.0 GB(38%)
244.1
LPMM-Adaml (2/8/8/8
RoBERTa-Large
1.5 GB
3.8 GB(72%
61.5
D COMPARISON OF LPMM WITH QUANTIZED TRAINING TECHNIQUES
To clarify the difference contributions and techniques between LPMM and other work, including quantized training, communication efficient training, memory efficient training, etc., summarize the computation bits and storage bits existing works. First, we would like to point out that there types of numerical precision in low-precision training system:
17

---

Under review as a conference paper at ICLR 2023
Table 6: Largest trainable(or finetunable) Model with same GPU memory.
Maximum Value
Dim
FP
Dettmers et al.
2021
LPMM
D
712
804
1016
ResNet-
296
340
472
D
40
72
148
RoBERTa-Large
W
1312
2048
3232
Storage bitwidth characterizes the numerical precision of parameters, gradients, and mo- mentums stored the memory, where we use
parameter storage bitwidth (PSB) refer the bitwidth of the stored parameters (a.k.a., primal/master parameters or latent weights);
gradient storage bitwidth (GSB) to refer the bitwidth of stored (parameter) gradients, which are fed into the optimizer;
momentum storage bit (MSB) to refer the bitwidth of stored momentums.
Computation bitwidth characterizes the numerical precision of the computing kernels (GEMM, conv, etc.), where we use
parameter computation bitwidth (PCB) to refer the bitwidth of parameter inputs to forward operators;
gradient computation bitwidth (GCB) to refer the bitwidth of (activation) gradient inputs to backward operators;
momentum computation bitwidth (MCB) to refer the bitwidth of momentum in opti- mizer operators.
We list the computation and storage bits of existing works Tab. 7 Existing now-precision train- ing methods (DoReFa, WAGE, FP8, INT8, GradScale) mainly consider reducing the computa- tion bitwidth. The parameters stored in high-precision (16 32) and then converted to low- precision (4 or 8) for computing forward and backward propagation. Among these methods, LPMM can achieve lowest storage bits. Mcanwhile. those memory efficient quantized training methods (WAGEUBN, FXPNet) suffer from major accuracy loss, as shown in Tab. g
Table 7: Difference between LPMM and other low-precision training works. p-parameter, g-gradient, m=momentum. "_" stands for not considered or mentioned.
Computation Bits
Storage Bits
Type
Method
p.
m.
m.
DoReFa-Net et al.
2016)
1
32
32
WAGEUBN Yang el al.
2020)
8
8/16
24
15
Quantized
FP8 training Wang et al. [NTu et al. 2020)
20181 8b)
8
16
16
16
16
Training
8
32
EXPNetíChen et 2017
I
12
32
12
12
12
GradScale Sun et al.
2020)
4
16/32
Communication
QSGD Alistarh et al. 2017
32
32
32
4/8
Efficient
signSGD Bernsteen et al.
2018
32
32
Training
LP-SGD et al.
2017)
32
32
1
32
Memory
HALP
De Sa el al. 2018)
32
32
Efficient
SWALP Yang et Yang
al. 2019
32
32
32
9/32
8
Training
8-bit Optm
Dppum.DDettneee (Dettmers et al. 2021
32
32
32
8
LPMM((uus)
32
32
32
12
8
Next, we compare LPMM with existing methods in more detail at the aspects of contributions, experiments and how stochastic rounding is used.
18

---

Under review as a conference paper at ICLR 2023
Difference Contributions First, talk about the quantized training. DoReFa-Net, INT8 and GradScale mainly focus on the acceleration of the training process. While they also can reduce the training memory footprint via low-bit kernel, the reduction is nearly independent to the model size. On the other hand, FP8 training, WAGEUBN and FXPNet focus on both training acceleration and memory footprint reduction (of the model storage). For communication efficient training, they quantize the gradient into lower bits for more efficient communication in distributed scenario. While the goals are different, both communication efficient training and LPMM use lower GSB, which is necessary and beneficial to persistent storage of gradients.
For memory efficient training, we mainly refer the work focusing on the quantization on post- update parameters and other values. LP-SGD only stays in theoretical analysis. HALP quantized the post-update parameters though, they store the main parameters in full-precision and in fact are intended to use low-precision computations. SWALP extended the HALP but in fact stored high- precision parameters as well. Unlike others, 8-bit Optimizer focused on reducing the storage of optimizer states instead of primal parameters, in fact the two are proportional (even equal), Actually, WAGEUBN, FXPNet and SWALP did not report the memory saving which probably indicates they only conducted simulation experiments, and to our best knowlegde, are the first to store low-bit parameters with < 16-bit in realistic scenarios.
Difference Experimental Results Here we mainly talk about the experiments related low- precision model memory. As shown in Tab. X the performance of WAGEUBN and SWALP are lossy compared with full-precision baseline on ImageNet dataset, while FP8 training and GradScale are nearly lossless on both CIFAR and ImageNet under their mainly claimed setting. The FXPNet has lossless accuracy on CIFAR but no experiments on ImageNet. Actually, there an intrinsic complexity gap between the two datasets. From our observation, 8-bit primal parameters are enough for CIFARIO 10 attain lossless accuracy, but 12-bit parameters are needed for ImageNet.
In other words, to our best knowledge, there are no comparable performance with full-precision baseline under the primal parameters less than 16-bit in previous work. However, LPMM attain lossless accuracy using 12-bit parameter. For other values, QSGD and 8-bit Optimizer can achieve lossless performance compared with full-precision using 4/8-bit gradient and 8-bit optimizer states, respectively.
Table 8: Comparison of experiments between LPMM and comparable previous work for ResNet on ImageNet
Storage Bits
Method
Model
p.
m.
Acc.
FP Baseline
WAGEUBN Yang et al. 2020
ResNet-18 18
24
8
15
67.4
68.70
FP8 training Wang ct al. 2018b
ResNet-18
16
16
16
66.95
67.43
EXPNet(Chen et al. 2017
12
12
12
GradScale ct al.
2020
ResNet-18
16/32
16
68.27
69.40
SWALP WWnLP et al. (PMM((urs)
2019)
ResNet-18
9/32
32
32
65.11
69.51
ResNet-18
12
8
8
71.0
71.1
LPMM(ours)
ResNet-18
8
8
8
69.2
71.1
QSGD (liistar et al. 8-bit Optim.(Dettmer (Dettmers al. 2021
2017)
Re
32
4/8
8
74.76
74.68
ResNet-50
32
32
8
77.2
77.1
LPMM((urs)
ResNet-50
12
8
8
77.3
77.2
How SR is used Then, we discuss about how stochastic rounding (Gupta et al. 2015) is used in various work. Firstly, FXPNet and GradScale did not use stochastic rounding technique in their main contribution. DoReFa-Net, WAGEUBN, INT8 only use stochastic rounding in the quantizer before forward and backward operations to generate a low-bit input to the low-bit kernel, which is completely different from our setting. The FP8 training, HALP, SWALP uses stochastic rounding in parameter update, same as LPMM. However, we dive deeply into why stochastic rounding is nec- essary from network training dynamics and parameter distribution, while they either give accuracy gap with or without stochastic rounding, or not concern about the SR technique itself. Lastly, 8-bit Optimizer and signSGD did not stochastic rounding, and QSGD used stochastic rounding on
19

---

Under review as a conference paper at ICLR 2023
gradients, which think possibly make different roles from that parameters since the differ- ence on distributions. Again, QSGD also did not concern about the why SR is necessary. In short, empirically analyze that why the unbiasedness of SR is necessary from the training dynamics and parameter distribution, which is not given previous work.
20


Reviews:
[{"review_id": "d0acpNReGzF", "reviewer": "Reviewer_wG9u", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes low-precision model memory (LPMM) as a quantization training concept. In LPMM, all parameters, including weights, activation, gradient, and momentums are quantized to low-bit, which achieves theoretical low-bit training that saves significant memory footprints on the hardware. ", "summary_of_the_review": "To sum up, the clarity and quality of this paper need to be improved. The author of the paper did some interesting works on model quantization but fails to demonstrate them with thorough discussion and experiments. Please refer to strengths and weaknesses for more information.\n\nI think this paper needs major revise, both on the technical contribution and experiments. \n", "strength_and_weaknesses": "Strengths\n\n1. This paper proposes the low-bit version SGD optimizer, which is a good approach towards low-bit training on the device that support low precision computation.\n\n2. The batch size experiments in figure 3 is very interesting. The author also provide convergence analysis on the training algorithm.\n\n3. The paper organization is clear, which makes it easy to read. \n\n\nWeaknesses\n\n1. There are many existing work on all low-bit training. And the author didn’t compare the proposed method to them. Such as WAGEUBN (Training High-Performance and Large-Scale Deep Neural Networks with Full 8-bit Integers), DoReFa-Net (Dorefa-net: Training low bitwidth convolutional neural networks with low bitwidth gradients), FXPNet (Fxpnet: Training a deep convolutional neural network in fixed-point representation), GradScale (Ultra-Low Precision 4-bit Training of Deep Neural Networks) and (Towards Unified INT8 Training for Convolutional Neural Network). They are very important research on all low-bit training and has been published in top conference/journals. Although some of them are cited in related works, but the author didn’t discuss nor compare with them.\n\n2. Because there are many similar works, then the novelty of this work is not that significant. And I don’t see the author discuss about the problems of existing works, which makes the motivation of this work relatively weak.\n\n3. One important part is missing, which is how batch-norm layers are used. During training, batch-norm is keeping the training, but it will change the low-bit number into fp number. Could  the author discuss this problem? \n  \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "NXLHpuqY0O", "reviewer": "Reviewer_UjMp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose low-precision model memory (LPMM), which quantizes a model's parameters, gradients, and optimizer state to low precision. Prior works have focused only on params/gradients or optimizer state separately. The main challenge is that params, gradients, and momentum may have very different dynamic range, leading to underflow (typically of the gradient updates). LPMM addresses this via unbiased stochastic rounding and increasing the batch size to increase the magnitude of the updates. The paper uses the term \"microbatching\" but this is just accumulating multiple batches of gradients before computing a weight update (some frameworks call this gradient accumulation).\n\nThe authors experiment with ResNet-18 and ResNet-50 on ImageNet, showing that LPMM with 8-bits for params, grads, and momentum can get within 2% of the full precision accuracy. If the params are kept in 12 bits then they can match full precision.\n\nEDIT: raised score from 3 to 5 after rebuttal.", "summary_of_the_review": "The paper applies stochastic quantization and microbatching to quantize parameters, gradients, and optimizer state. These techniques are well known in DNN quantization and thus the paper is a bit lacking in novelty. The experiment section is weak with only results for ResNets on ImageNet. Lastly, the techniques may be difficult to implement and/or incur run time overhead, which is currently not discussed.", "strength_and_weaknesses": "Strengths:\n - Appears to be the first paper to quantize params, gradients, and optimizer state jointly\n \nWeaknesses\n - Experiment section is weak. ImageNet is no longer a SOTA benchmark and ResNet is an 8-year old model that is known to be easy to quantize and train. Experiments focus only on SGD with momentum, but currently ADAM is necessary for Transformers.\n - The techniques applied are well known. For example, the Deepmind Gopher paper [1] used stochastic rounding during parameter updates. Microbatching is already implemented as a toggle option in HuggingFace and Megatron Transformer frameworks. These techniques are not really used in a novel or insightful way.\n - The technique seems difficult to implement and may incur run time overhead. The authors had to rely on non-linear quantization for the momentum, which is quite complex. Unbiased stochastic quantization requires generating random numbers and comparing them to a high-precision rounding probability for each element. These two issues may cause run time overhead on GPU and present difficulties for dedicated hardware. I would like the authors to discuss the actual memory saved on GPU (not just bitwidth) and whether LPMM reduces wall-clock time for training.\n\n[1] https://arxiv.org/pdf/2112.11446.pdf", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "N0Shaym7eSI", "reviewer": "Reviewer_wNyQ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose low-precision training techniques to lower the number of bits to represent parameters, gradients, and moments. SGDM optimizer is assumed while low-precision training can improve the model accuracy by using two main principles: unbiased stochastic quantization and microbatching. Overall, memory consumption can be reduced by over 70% when ResNet-18 and ResNet-50 (on ImageNet) are trained.", "summary_of_the_review": "- Ideas are incremental (or not new at all)\n- Experiments need to be more extensive. Only two models are not enough to prove the claim.\n", "strength_and_weaknesses": "*Strength\n- Convergence of the proposed low-precision training is analyzed.\n- The impact of microbatching is investigated with details.\n\n*Weakness\n- Stochastic rounding has been introduced by lots of previous works. What is the major difference compared to those related works? For example, the paper titled \"Training deep neural networks with 8-bit floating point numbers\" already introduced while the first paper introducing stochastic rounding is probably \"Deep Learning with Limited Numerical Precision\" which is missing in the reference.\n- Properties of microbatching seem to be very similar to chunk-based additions of the above reference.\n- Comparisons with previous low-precision training are missing in the experiments.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2PVtJ4MwyLd", "reviewer": "Reviewer_Dafx", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper addresses the issue of large model memory as a bottleneck to model scaling. It proposes a framework to compress model parameters, gradient accumulations, and momentum, during training. Achieves iso-accuracy on ResNet/Imagenet models down to 12/8/8 bits, respectively, and minimal degradation down to 8/8/8 bits. Underflow is addressed by combining stochastic rounding with microbatching (accumulation of subsequent gradients). Convergence is mathematically guaranteed in the presence of stochastic quantization.", "summary_of_the_review": "The paper applies known techniques to successfully compress model memory (i.e., model parameters, gradient accumulations, and momentum) and achieve strong compression and iso-accuracy in two CNN models. I would like to see more details about the algorithm implementation and its runtime (possibly an additional appendix). ", "strength_and_weaknesses": "Strengths:\n- combination of stochastic rounding (SR) and microbatching achieves remarkable model memory compression with no (or limited) accuracy degradation\n- good theoretical support on LPMM-SGDM convergence\n\nWeaknesses:\n- the paper heavily relies on known techniques (SR and microbatching) which are applied as-is to the particular case study\n- the main insight is claimed to be identification of underflow as the core issue to be solved. However, the impact of delays in the parameters update (or lack thereof) is a well known quantization-related issue, both in the case of having access or not to a full precision copy of the quantized parameters (see for example the already cited [1], which discusses the conceptually-similar swamping issue as well as SR, or [2], which discusses mitigation strategies for updates delayed by quantization). So, it's hard to accept the claim that it has been \"discovered\" in this particular scenario\n- demonstration is limited to two CNN models so it's unclear if findings may generalize (e.g., is this compression level suitable for other models?)\n- what is the runtime (or training time overhead) of this algorithm? how practical is it to apply this training technique beyond the two examples shown?\n- there is a mention of a training algorithm implementation at the level of CUDA kernel in PyTorch but no details about it\n\nOther corrections:\n- section 3.2: part of sentence missing: \"...is far from the Extremely, there could be\"\n- algorithm 1: should the quantization in line 5 be carried out _after_ all the micro-batch gradients have been accumulated? Doesn't applying the quantization at every step defeats the purpose of accumulation if gradients are individually small?\n\n[1] N. Wang et al. \"Training Deep Neural Networks with 8-bit Floating Point Numbers\" (section 3.2)\n[2] M. Nagel et al. “Overcoming Oscillations in QAT”", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject