Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training
Authors: Yijia Zhang
Yibo Han
Shijie Cao
Guohao Dai
Youshan Miao
Ting Cao
Fan Yang
Ningyi Xu
Venue: Submitted to ICLR 2023 (2023)
Keywords: Large model training
Memory reduction

Abstract: Running out of GPU memory has become a main bottleneck for large-scale DNN training. How to reduce the memory footprint during training has received intensive research attention. We find that previous gradient accumulation reduces activation memory but fails to be compatible with gradient memory reduction due to a contradiction between preserving gradients and releasing gradients. To address this issue, we propose a novel optimizer accumulation method for Adam, named Adam Accumulation (AdamA), which enables reducing both activation and gradient memory. Specifically, AdamA directly integrates gradients into optimizer states and accumulates optimizer states over micro-batches, so that gradients can be released immediately after use. We mathematically and experimentally demonstrate AdamA yields the same convergence properties as Adam. Evaluated on transformer-based models, AdamA achieves up to 23% memory reduction compared to gradient accumulation with less than 2% degradation in training throughput. Notably, AdamA can work together with memory reduction methods for optimizer states to fit 1.26x~3.14x larger models over PyTorch and DeepSpeed baseline on GPUs with different memory capacities.

Under review as a conference paper at ICLR 2023
ADAM ACCUMULATION TO REDUCE MEMORY FoOT- PRINTS OF BOTH ACTIVATIONS AND GRADIENTS FOR LARGE-SCALE DNN TRAINING
Anonymous authors Paper under double-blind review
ABSTRACT
Rumning out of GPU memory has become main bottleneck for large-scale DNN training. How reduce the memory footprint during training has received intensive research attention. We find that previous gradient accumulation reduces activation memory but fails to be compatible with gradient memory reduction due to a contradiction between preserving gradients and releasing gradients. To address this issue, propose novel optimizer accumulation method for Adam, named Adam Accumulation (AdamA), which enables reducing both activation and gradient memory. Specifically, AdamA directly integrates gradients into optimizer states and accumulates optimizer states over micro-batches, so that gradients can be released immediately after use. We mathematically and experimenaally demonstrate AdamA yields the same convergence properties as Adam. Evaluated on transformer-batt models, AdamA achieves up to 23% memory reduction compared gradient accumulation with less than 2% degradation in training throughput. Notably, AdamA can work together with memory reduction methods for optimizer states fit 1. 126x-3.1m 14x larger models over Py Torcch and DeepSpeed baseline on GPUs with different memory capacities.
INTRODUCTION
The past few years have witnessed the remarkable achievements of large-scale DNN models across domains from computer vision to natural language processing (Devlin et al.] 2018 Radford et al.. 2019 Dosovitskiy et al.] 2020) Smith et al. 2022). Training such big models requires massive powerful GPUs with indispensable large memory capacity, which is prohibitively expensive and inaccessible to most researchers. Even for fine-tuning large pre-trained model where computational power is less critical factor, running out of memory is increasingly becoming the first and foremost serious limitation (Ren et al. 2021 Rajbhandari al. (021).
Recently, there has been an explosion of interest around methods to reduce the memory footprint during model training (Sohoni al. 2019 Rajbhandari et al. 2020 Pudipeddi et al. 2020 Chen et al. 2016 Shazzerr & Stern 2018). However, there is hardly a one-size-fits solution to address the out-of-memory issue for two reasons. Firstly, many memory reduction methods usually come at the cost of sacrificing convergence (Mostafa & Wang 2019. Micikevicius et al. 2017) or training throughput (Chen et al. 2016 Pudipeddi et al. 2020).. It remains unclear how significant the cost of one method or combination of methods would be for different models before testing. Secondly, the ratio of the memory footprint of various parts (e.g., weights, gradients, optimizer states, activations) varies with the model and training configurations. No single method always performs best in different cases.
Among memory reduction methods, gradient accumulation and gradient release are two effec- tive methods to reduce activation memory and gradient memory, respectively (Huang et al., 2019) Pudipeddi et al. 2020). Both methods have no negative impact model convergence and training throughput. Unfortumately. these two methods are inherently mutually exclusive. Gradient accumula- tion reduces the activation memory by splitting mini-batch into sequence of micro-batches and accumulating the gradients of all micro- batro-b Gradient release reduces the gradient memory by freeing up the pradient-occupied space in a layer-by-layer manner. The contradiction preventing

---

Under review as a conference paper at ICLR 2023
the two from being used together is must preserve accumulated value of gradients until the last batro, but the other releases the gradients immediately after use. Saving activations or gradients, previous works prefer the former as activations usually consume the most memory during training, while the gradients memory can be ignored when models are small. However, with the ever-increasing model size, the gradient memory consumption cannot be ignored.
Micro-batch,
Micro-batch,
Gradient accumulation
FWD, BWD,
Accumulate to
FWD, & BWD;
Accumulate
Upsone
acoumulered
Update
Micro-batch,
Micro-batchy
Optimizer accumulation with AdamA
FWD, & BWD,
Acoumulane & with &o
FWD, & BWD;
Accumulate & with g,
...
Update w
Figure 1: AdamA to integrate gradients into optimizer states and accumulate optimizer states over micro-batches.
To enable saving memory footprints of both activations and gradients, we propose a novel opti- mizer accumulation method for large-scale DNN training with Adam, named Adam Accumulation (AdamA), which can still maintain the convergence and training throughput. Specifically, instead of accumulating gradients, AdamA integrates gradients into optimizer states (m and in Adam) immediately after the gradients are produced, and accumulates optimizer states sequentially over micro-batches, as shown Figure [ This subtle change of directly integrating gradients optimizer states makes the memory space for whole model gradients no longer needed, eliminating the afore- mentioned contradiction between preserving gradients and releasing gradients. Consequently. AdamA can reduce the gradient memory to I/M M of the original (M is the number of layers). and the activation memory to 1 1/N of the original (N is the number of micro-batches). We further mathematically and experimentally demonstrate AdamA performs the same as standard Adam in terms of the convergence properties and final model accuracy, although the optimizer update of AdamA deviates a little from standard Adam. Notably, AdamA is complementary to previous methods that reduce weights and optimizer states, providing a possibility to achieve an even higher memory reduction rate.
We evaluate AdamA on both language and vision tasks, with the typical transformer architecture and convolution architecture. Our experimental results show that AdamA performs the same convergence properties Adam. Compared with gradient accumulation baseline, AdamA can reduce memory footprint up to 23% with less than 2% degradation in training throughput. We further combine AdamA with DeepSpeed ZeRO-DP POs. which aims reduce optimizer states in distributed data parallel scenario. Training with AdamA, DGX system can fit model x -3.14x larger over PyTorch and DeepSpeed baseline can do.
Our contributions can be summarized as follows:
We propose AdamA, novel optimizer accumulation method to enable reducing memory footprints of activations and gradients simultaneously. Compared with gradient accumulation baseline, AdamA save up to 23% memory footprint.
We conduct convergence analysis for AdamA. Mathematical and experimental results on real workloads show AdamA performs the same convergence propertics as Adam.
We implement the training pipeline of AdamA with Pytorch and DeepSpeed. The system easy to use and incurs less than 2% effect on training throughput.
2 BACKGROUND AND RELATED WORK
The memory footprint during model training be categorized into four parts: weights, gradients, optimizer states and activations. As different models, optimizers, or batch sizes lead to different ratios of the four parts, many works have been proposed to reduce them accordingly.
Reducing weight and optimizer state memory. In model training iterations, weights and optimizer states inherently have the temporal dependency, i.e., the values at time step / update the basis of
2

---

Under review as a conference paper at ICLR 2023
their values at time step t 1. Hence, the training system must maintain weights and optimizer states for updates between two consecutive iterations. To reduce the weight and optimizer state memory. many compression- -aaddd methods (e.g., sparsification, quantization and matrix approximation) have been proposed, but often sacrifice the convergence or end model accuracy (Mostafa & Wang 2019) Micikevicius al. 2017 Shazeer & Stem 2018).
Reducing activation and gradient memory. Activations and gradients are computed and used only inside each training iteration, indicating a potential release the memory occupation after finished. Gradient accumulation and gradient release are effective methods to reduce activations and gradients. respectively (Huang et al.) 2019 Pudipeddi et al. 2020). The key idea behind gradient accumulation is to split mini-batch into several micro-batches. This method computes the gradients of micro-batches sequentially and accumulates them to reduce the memory footprint of activations as well as to keep the same convergence properties as the original mini-batch. Gradient release executes the backward process in a layer-by-layer manner, which immediately releases the gradient-occupied memory after the weight updating is finished, so that the memory allocated for gradients can be reduced from size the whole model size to the size of the maximum layer.
The contradiction between gradient accumulation and gradient release. Unfortunately, gradient accumulation to save activation memory and gradient release to save gradient memory are mutually exclusive. Because while the other frees up the gradients immediately after use. Our proposed AdamA resolves this contradictory and enables saving both activation and gradient memory. Please note that method is complementary to previous memory reduction methods (e.g., checkpointing (Chen et al. 2016).
must maintain all gradients for accumulation until the last dicro-batch.
Adafactor (Pudipeddi et al., 2020), offloading (Rajbhandari al., 2021 Ren al., 2021; Pudipeddi et al. (0200), and can be applied together with these methods to achieve even higher memory reduction rate.
3 METHODS
3.1 ADAM ACCUMULATION (ADAMA)
As mentioned, gradient accumulation to save activation memory contradicts gradient release to save gradient memory. The core reason is that gradient accumulation accumulates gradients till the last micro-batch, so that the gradient memory of the whole model must be preserved. Intuitively, as gradients are eventually used to update the optimizer states (m and v in Adam), we can integrate gradients into optimizer states in advance, the gradients memory can be released, thus resolving this dilemma. Inspired by this insight, we for the first time propose an optimizer accumulation method, namely AdamA, that integrates gradients into optimizer states immediately after produced and then accumulates optimizer states sequentially over micro-batern batche..
Algorithm 1 Adam v.s. AdamA with micro-batches
Initialize Oo. mo <- up <- l 0, N # of microbalches
while 0, not converged do
1. <-- [+]
for each micro-batch i mini-batch do
9t.1 N/ (of(,((00-) end
m,
= 31m4-1 + +1 -1 (1-)))
N-1
== gt,i
820, (1-82)((E) Di= I1.i) 2 U.ss.
N-
Update
mh 1- 14
02 x 0--1
amá
I-33
Vott
end
Algorithm| iillustrates the difference between standard Adam and our proposed AdamA, in the of micro- batcoon Standard Adam first accumulates gradients of all micro-batches, then updates my with the accumulated gradients and with the square of the accumulated gradients (as shown in
3

---

Under review as a conference paper at ICLR 2023
the blue text). Different from the update mechanism in Adam, our proposed AdamA updates ea through accumulating the square of gradienns generated from each micro-batch. This slight change in AdamA allows that gradients can be used and released immediately once they generated, leading to a significant reduction on gradient memory during training. In order to analyze AdamA's impact on model convergence,
mathematically prove that AdamA yields the same convergence rate as Adam (shown in Section 3.2). and experimentally demonstrate AdamA performs the same as Adam in vision and language tasks (shown in Section
In Algorithm 2] we show the detailed training pipeline of AdamA to reduce both activation and gradient memory. Similar gradient accumulation, AdamA divides a mini-batch of training data into several micro-batchan to reduce activation memory to N of the original without During the backward pass of each micro- batch, once the gradients of layer (gr,.)) are produced, j and gi.., will be accumulated to the optimizer states of this layer (m,, and Vt.j), respectively, In this process, It,i. memory is released after the accumulation procedure. As result, the peak memory allocated for gradients can be reduced to only I of the full model gradient size.
37
Algorithm 2 The training pipeline using AdamA to reduce both activation and gradient memory
Initialize Oo. mo < 0, 1o 6 0, l + 0, N + # of micro - batches, M <- # of layers
while O, not converged do
I < 1+ 1, m, 31 - U4 + 32
for each micro-batell i in mini-batch do
// Reduce activation memory 1/N of the original without micro-batches
for each layerj j in backward computing do
9t.i.j + N. Vo.t....(00 (0:-1)
Assign memory for St.i.j // Reduce gradient memory to 1/M of full model gradients
maj f- Mt.j + (1 B1)gr...
Ve.j f Vt + (1 32)92,1
Release memory for It,i, The gt.i values are not needed any more
end end
Update
mh
m. 3
Ur
f-
T-
02 t- 0,-1
Vie+e
end
3.2 CONVERGENCE ANALYSIS
In this section, we demonstrate the convergence properties of AdamA. Adam (Kingma & Ba 2014) is optimization method that adaptively rescales the updating vector with second moment of the gradient. Compared with Adam, AdamA has the same updating direction (i.c., m)), but different adaptive scaling length (i.e., ).) ). We refer Adam's 's proof methods to show that AdamA has the same theoretical convergence properties as Adam.
Following analysis method in the online learning Saunewokk(Zinkeviii 2003). we define fi as the convex cost function at time t, and 0, as the parameter we predict. We evaluate the convergence properties of AdamA using the regret R(T)) = Ex= (f. - f:(0)), which the sum of all the previous difference between our prediction fi ((e) and the best fixed point parameter f(()) (Kingma & Ba 2014). In Theorem 1, guarantee that AdamA has the same regret bound O(VT) with Adam and the detailed proof is given in the appendix. We define the vector 91.T.s, E TR² as the ith dimension of gradients from the 8th micro-batch in one mini-batch till T. Following Adam paper, Theorem I holds when the learning rate is decaying at a rate of t- 1 and first moment running average coefficient 31,t decay exponentially with X. (Kingma Ba 2014
Theorem 1. Assume B1 B2 [0,1) satisfy 7
1. N is the number of miero-batches in a batch. The function ft has bounded gradients, V (())l < G. (0) < Go for all 0 E Rdd. For any m, E (1,..., 7), the distance between any generated by AdamA is bounded,
vi
4

---

Under review as a conference paper at ICLR 2023
which can be presented as |00,- 6m l22 S D, 110, mml I Doo. AdamA achieves the following guarantee, for all T 2 1.
D R(T) S 20(11 - 31)
N Torst a(31 + mz
D Gxv 20 31) A)2
ΣΣ
19rT, 6112+
(1- (1-3
In Corollary 1, we show the average regret of AdamA is O ( which is the same as Adam. It is obvious the limit of the average regret is O when T gets larger.
Corollary Assume that the function ft has bounded gradients, (0)] - S. (V)( I Gec for all A E R For any m, E {1.... T, the distance between any 8g generated by AdamA is bounded, which can be presented as 10, 0m S D, ||0,, mm Do Combining Theorem >> hml [1 all S oG V we can get the average regret of AdamA:
with the upper bound >
R(T) 0
T
3.3 SYSTEM IMPLEMENTATION TO REDUCE MEMORY FOOTPRINT
Our system design is intended to reduce memory and incur little impact on training throughput at the same time. The challenge achieve efficient memory allocation and low communication cost.
In the single device scenario, the model forward computing is done as normal. When doing backward computing. interleave back propagation and optimizer states update. For each model layer, gradients are accumulated to the corresponding optimizer states and immediately released. We implement this mechanism in Py'Torch with backward hook to insert the optimizer states accumulation and layer gradient release operations. It should be mentioned that frequent memory allocation and free operations are costly. Nevertheless, the deep learning framework would maintain a memory pool for tensor memory assignment and release. This prevents the heavy overhead of system calls.
In the distributed data parallel scenario, the same operations apply except that gradients need to all- reduced among distributed devices. When training with AdamA, the straightforward implementation is to insert layer gradients all-reduce operation each micro-batch to update optimizer states. And yet, compared with standard Adam procedure, where all-reduce only needs to be done once after all miiro-bbtccess the communication cost would be increased from 0(1) to O(N) in one mini-batch (N is accumulation steps). To reduce the communication volume, choose to all-reduce optimizer states instead of gradients. In this way, local optimizer states are updated on each device and synchronized at the end of the mini-batch with PyTorch all-reduce API. Therefore, the communication volume stays constant in one mini-bus batch. The details about our system design can be found in Appendix.
4 EXPERIMENTS
In this section, we evaluate AdamA on both vision and language tasks. We show that AdamA does no harm the convergence properties compared with Adam. Then, we demonstrate the memory reduction result of AdamA and its impact on training throughput. Finally, we include case study where AdamA is combined with DeepSpeed ZeRO-DP to further explore the ability train large-scale models and push the limit to reduce memory of gradients, activations and optimizer states.
4.1 CONVERGENCE BEHAVIOR
To verify the convergence properties AdamA, we experiment on both NLP (transformer-bascd) and CV (convolution-basee models. We set the same mini-batch size when training with Adam and AdamA. We set the accumulation steps N to 2,4,8 when training with AdamA.
For NLP model, follow RoBERTa method (Liu al. 2019) to pre-train BERT-Large (L 24, = 1024, = 16, 340M) on a DGX A100 with sequence length of 128 and mini-batch size of 1024. We use the implementation of BERT model from Microsoft DeepSpeed (Rasley et al. (020).
5

---

Under review as a conference paper at ICLR 2023
Adam
AdamA(N=2)
Trainling loss
AdamA(N=4)
AdamA(N=8)
200M
Samples
Figure 2: Sample-wise convergence properties for BERT-Large pre-training with sequence length 128 using DGX A100. AdamA has almost the same training loss curve with Adam.
For the pre-training corpus, we use the English Wikipedia and BooksCorpus downloading from NVIDIA GitHub (NVIDIA). should be noted that the corpus we is different from that used in original BERT (Devlin et al. 2018) because the BERT corpus is not available to the public at this time. For the pre-training hyper-parameters, we follow the method in RoBERTa.
Table Similar accuracy can be achieved by fine-tuning on the BERT-Large model pre-trained by Adam and AdamA. F1 scores are reported for MRPC, Spearman correlations are reported for STS-B, and accuracy scores are reported for other tasks.
Setting
MNLI-M MNLI-MM SST-2
MRPC STS-B
QNLI
QQP
RTE
CoLAA
Adam
80.62
80.96
90.48
84.03
86.56
87.46
86.53
58.48
43.25
AdamA (N=2)
80.50
80.72
90.83
86.27
84.18
87.59
86.49
57.04
43.27
AdamA (N=4)
80.38
80.93
90.48
85.39
85.71
87.30
86.52
56.68
43.11
AdamA (N=8)
80.39
80.97
89.79
85.67
85.62
87.44
86.51
61.01
43.86
Figure presents the sample-wise convergence results when training BERT-Large with Adam and
AdamA. No matter how many micro-batches one mini-batch, we find the convergence curve of
AdamA coincides with that of Adam. To further evaluate the convergence of the BERT-Large model trained by Adam and AdamA, we fine-tune the models on all tasks from GLUE benchmark (Wang al. 2018). We fine-tune for epochs for all tasks and select the best fine-tuning learning rate (among 2c-5, 3c-5, 4c-5, 5e-5) on the Dev set. Table sshows shows the fine-tuning results. Obviously, the model pre-trained with AdamA provides similar accuracy with that pre-trained with Adam.
80
0
Traininn loss
AdamA (N-2)
AdamA (N-4)
[op] Accuracy
60
AdamA (N=6)
50
40
30
AdamA (N=2)
20
AdanA (N=4)
10-
AdamA (N=8)
20k
40k
50k
Bok
100k
120k
Steps
Epoches
(a) Training loss
(b) Test Top-1 accuracy
Figure 3: The training loss curve and the test accuracy of ResNet-50 on ImageNet.
For CV model, train Re with A100 GPUs ImageNet (Deng et al. 2009) dataset to evaluate the convergence properties of AdamA. Following the training setting provided by MMClas- sification (mmlab), train ResNet-50 with mini-batch size of 1024. For the learning rate, initial

---

Under review as a conference paper at ICLR 2023
it to 1e-3 and cosine decay it to 1e-5. Figure presents the training loss curve and the test accuracy of ResNet-50, from which we can jump to the conclusion that AdamA has almost the same convergence properties with that of Adam in CV tasks.
Considering that Batch Normalization (BN) (lofte & Szegedy 2015) is used in ResNet, we also pay attention the effect model accuracy which may be brought by the difference of the micro- -btch normalization statistics and that statistics of the entire mini-batch. Mentioned the influence of BN on model convergence tends be constant if the micro-batch size increases above a certain extent. Therefore, we do not pay efforts to keep exactly the same BN algorithm between chicro-batched training and con-micro-bbtthed one. In the experimental results shown in Figure we also find the impact on convergence can be ignored.
(Sohoni et al., (019).
To help understand the difference between AdamA and Adam during training process, we make the following statistics. From the update equation 0, < 01-1
antis clear that our adaptive scaling
length differs from the standard Adam in coefficient Vol/Vl'' We track the coefficient training ResNet-50 on CIFAR-100 dataset. In Figure 4 we plot the mean value of V 0 during each training step and its value range. It shows generally the coefficient keeps around 1.0 and the deviation value range is within 1%. We think the minimal deviation in each iteration might contribute the same convergence properties of Adam and AdamA during training.
deviation coefficeett value
1.015
1.010
1.005
1.000
0
10000
20000
30000
40000
50000
steps
Figure 4: Statistics on the value of generated from ResNet-50 CIFAR-100. The real line
L
Vo
shows the coefficient mean during iterations and the area shows its range. shows that the value AdamA deviates from the standard Adam is within 1%.
4.2 MEMORY REDUCTION
Allocated Memory (G8)
Gradient Accumulation
AdamA
10.3
10
7.1
7.2
5.6
N N
=
Accumulation Steps
Figure 5: The memory reduction of AdamA compared with gradient accumulation when training BaRg-L
As AdamA eliminates the contradiction when combining gradient accumulation and gradient release, we first show the improvement of AdamA compared with gradient accumulation. Compared with gradient accumulation, AdamA save the memory footprint of both activations and gradients. We measure the memory footprint when training BERT- with AdamA on a DGX A.00 (8 A100 GPUs) with the mini-batch size of 256 and the sequence length of 128. As shown in Figure AdamA
7

---

Under review as a conference paper at ICLR 2023
save 1.6GB more memory than gradient accumulation matter how many the accumulation steps are set in a mini-batch.
Table 2: When training BERT-Large. AdamA achieves less memory usage than Adafactor (Shazeer & Stern 2018) and SM3 (Anil et al. 2019).
Optimizers
Reduction target
Mini-batch size per GPU Memory usage per GPU (GB)
Adam (baseline)
N/A Optimizer states Optimizer states
8
6.15
Adafactor
8
4.83
SM3
8
4.90
AdamA(N=8)
Activations and gradients
8
4.18
To further show the memory saving effect of AdamA, expand BERT model to BERT-4B with billion weights using the scaling method of GPT-3 (Brown et al. 2020). We set the mini-batch size to 64 and accumulation steps to 8 in this experiment. In Figure ((a), train BERT-4B with gradient accumulation and AdamA using Py PyTorch framework. It can be found that AdamA can save 23.2% memory footprint compared with gradient accumulation when the weights number of model get to 4 billion.
Compared with other memory-efficient optimizers, e-g. Adafactor (Shazeer & Stern 2018 and SM (Anil et al. (211)), the memory reduction of our proposal is bigger under the same experiment setting. The comparison with Adam baseline is shown in [able[]] The reason AdamA can reach more significant memory reduction is AdamA targets at reducing the memory usage of both activations and gradients, while other works only aimed reduce optimizer states memory. the same time, AdamA can work well with these previous works to get further memory reduction.
Allocated Memory (GB)
70
69.3
62.9
Weights
Allocated Memory (GB)
70
69.9
Optimizer Activations Weights
Opt maler
60
40.3
42.2
40
29.7
30
30
22.1
20
20
Baseline
Gradient Accumulation
Adami
Pa ZeRB-OP Pacyy ZeRO-BP P.,
(a) PyTorch
(b) DeepSpeed
Figure 6: The memory reduction of AdamA when training BERT-4B using PyTorch and DeepSpeed.
To show the compatibility of AdamA with existing methods, combine AdamA with ZeRO-DP (Rajbhandari et al., 2020). a popular memory reduction method for optimizer states. ZeRO-DP Po. partitions the optimizer states different GPUs when training with data parallelism. In Figure 6 we combine AdamA with ZeRO-DP Pos to further reduce gradients and activations. It shows that AdamA with ZeRO-DP Pos can save 20.1 GB more memory footprint than only ZeRO-DP Pos. Even compared with ZeRO-DP Postg. which partitions both optimizer states and gradients, our combined method can reduce 7.6 GB more memory.
Table 3: The largest model size can fit on different DGX systems with AdamA. DeepSpeed
Py PyTorch
Gradient accumulation
AdamA
ZeRO-DP Pos
ZeRO-DP Pos AdamA
DGX-1 DGX-2
1.4B
1.8B
1.1B
3.3B
3.0B
4.0B
2.5B
6.8B
DGX A100
7.6B
9.6B
5.8B
18.2D
8

---

Under review as a conference paper at ICLR 2023
In Table 3 explore the largest transformer based model can fit DGX systems with various memory capacity with AdamA. At present, the mainstream DGX systems on the market include DGX-1 (8 V100-16GB GPUs), DGX-2 (16 V100-32GB GPUs), and DGX A100 (8 A100-80GB GPUs). In order to keep the same experimental settings, we set the number of GPUs to 8. The mini-batch size and accumulation steps set to 256 and 8, respectively. With PyTorch framework, the largest model AdamA can train is 1.26x to 1.33x larger than gradient accumulation can train. Combined with DeepSpeed ZeRO-DP, AdamA can train model with 18.2 billion weights in a DGX A100, which is 3. 3. larger than the model the system can train with only ZeRO-DP Pos.
4.3 TRAINING THROUGHPUT
Training, samplesi
3200
T 3000
2400
Adam
Adam
2,220
Adam
3000
AdamA
E 2600
AdamA
2000
AcamA
@2600
2800
B2200
2400
Trainif 19 1400
1600
2400
Tra leini
2000
1200
1800
N
1000
N
Accumulation Steps
Accumulation Steps
Accumulation Steps
(a)
(b)
(c)
Figure AdamA has less than 2% effect on the training throughput compared with gradient
accumulation using Adam: (a) training ResNet-50 with single GPU; (b) training BERT-Base with
A100 GPUs; (c) training BERT-Large with 8 A100 GPUs.
In this section, we show AdamA has negligible impact on training throughput. During training, is reasonable to set the micro-batch size as large as the device memory can contain, in order to saturate GPUs to achieve maximal training throughput. Therefore, the micro-batch size is fixed in this section.
Single-GPU Scenario As mentioned Section 3.3 our system design for AdamA Single-GPU
implementation is intended to incur extra throughput overhead. In Figure 7a), conduct throughput comparison with standard Adam training ResNet-50 with one A 1100 GPU. We keep the miero-batch size to 256 and switch accumulation steps to 4 and 8. We can conclude that training with AdamA has little throughput impact in single-GPU scenario.
Distributed Data Parallel Scenario Explained in Section 3.3 our system design keep the com- munication number to be constant by synchronizing the optimizer states. Although it may incur more communication volume compared with standard Adam that synchronizes the gradient, the impact on throughput minimal. In Figure (b)C). we conduct multi-GPU experiments with two models: BERT-Bas with 4 A100 GPUs and BERT-Laage with 8 110 GPUs. The micro-batch size of all the models is set to 1024. The experiments show that the training throughput difference within 2%. The throughput gap between AdamA and Adam is gradually decreasing with the increase of gradient accumulation steps. This because the communication volume is constant in mini-batch, the communication overhead proportion becomes smaller in a mini-batch with larger gradient accumulation steps.
CONCLUSION
This paper presents Adam Accumulation (AdamA), a novel optimizer method for large-scale DNN training. It enables saving memory footprints of activations and gradients simultaneously. Besides, AdamA yields the same convergence properties as Adam. Compared with gradient accumulation, AdamA can reduce the memory footprint up to 23% with less than 2% degradation in training through- put. Combined with memory reduction methods for optimizer states, AdamA can fit 1.26x-3.14x larger models Py Torch and DeepSpeed baseline different DGX systems.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.06/74, /4, 2016.
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, pp. 248-255 Icee, 2009.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/8 /0lo:/0880005 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale, In International Conference on Learning Representations, 2020.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. PMLR, 2015.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6958 2014.
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: arXiv: 2019.
Paulius Micikevicius. Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint orXiv:/710.0337 2017.
Open mmlab. Open mmlab mmelassification. Imcllass ficat on/ Accessed May 17, 2022.
lab/
Hesham Mostafa and Xin Wang. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In International Conference on Machine Learning, pp. 4656-4655. PMLR, 2019.
NVIDIA. Nvidia deep learning examples for tensor cores.
https://github.coom////D
Deepleeaniing@xxmmess////e Accessed May 17, 2022.
Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, and Sujeeth Bharadwaj. Training large neural networks with constant memory using new execution algorithm. arXiv preprint arXiv:2002.05665, 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya 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. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEF, 2020.
10

---

Under review as a conference paper at ICLR 2023
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, pp. 1-14, 2021.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimiza- tions enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 3505-3506, 2020.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang. Minjia Zhang, Dong Li, and Yuxiong He. (ZeRO-Offlodd): Democratizing (Billion-Scale) model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604, PMLR, 2018.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et Using deepspeed and megatron to train megatron- -ettrtrnn- nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.1199 2022.
Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszezynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report. arXiv preprint arXiv: 10631, 2019.
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 arXiv:/804.0746/. 2018.
Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference on machine learning (icml-03), pp. 928-936, 2003.
11

---

Under review as a conference paper at ICLR 2023
A CONVERGENCE PROOF
In this section, we show the convergence analysis of AdamA. Compared with proof in the original Adam paper Appendix, it's easy AdamA follows almost the same analysis proof with Adam. The most obvious difference is that Adam doesn't i take micro-batch into consideration. Here we use symbol N as the number of micro-batch in AdamA and subscript for micro-batch index.
Here we highlight those conclusions that differs between Adam and AdamA from Adam paper Appendix.
First, we construct our proof based on the claim convex function can be lower bounded by hyperplane, which is Mathematically expressed Lemma
Definition 1. A function f Ra + R is convex if for all x, y E Rd for all X e [0,1].
+1 2 f(Ax + (1 A)y)
Lemma 2. Ifa function f: R -> R is convex, then for all x, Rd
f(v) 2 f(x) Vz(2) x)
Lemma 3 and Lemma are proved to support the proof of Theorem
Lemma 3. Let g: (.(((e) and 9ı.t be defined as above and bounded, /|9/l22 S G. 119ello S G... Then.
L.. < 2G all t
Proof. The proof is the same Lemma 10.3 in Adam Appendix(Kingm) & Ba 2014) and hence is ommitted here.
B For 31, 82 E [0, that satisfy VB
# 1 and bounded 9t, 119.11 < G,
Lemma 4. Let y
v3,
lggellss lo I Goo, and the micro-bate batch number equals to N, the following inequality holds
T
m
2
1
N
-
1gr. all2
Vth,,
1 - v1
32
12

---

Under review as a conference paper at ICLR 2023
Proof.
mi,i
T
mi
V1 82
CE (1 - 31181 - )x,
v
v t1
(1 -1-34 32
| (1 - 322 -22
m
T
-1),8 gx,))
Vin
(1 B8
T(1 B2)
-k Sk.
1=1
mi.i
v1 83 (1- 31)2
T
B
Σ T(
Vliis
(1 - 3/ VT(1 - B2)
V
k=1
VI- (1-3))
T
37
N
T
T-
gh.i. bll2
Vibri
(1 - B7) VT(1 - B2) k=1
VB2
b=1
mi.
v1 - 83
(1-8,)²
3
T
T
k
1l bll2
(1 8488
VT(1 - B2)
mi.e
T
T
N
19 bll2
VT(1 - 12)
k=1
T Eae 11g1,T, 6112
T-1
Σ
trys
vt(l - (2)
j=0
T Em 11g1,T,1, 6/12
T
Σ
tyd
vt(l -32)
1=1
j=0
Applying Lemma 3,
T m
N
2G 00
S
/ / (32)
b=1
Theorem 5. Assume that the function fr has bounded gradients, [ < G. (.I S Gos for all 0 E Rd and distance between any 9, generated by AdamA is bounded, 1|0, 6mll S D.
||0,, |0 I Doo for any m, n {1...., Th. and 81 [0, 1) satisfy a micro-batch number equals to N. Let a; = Vt and BI, 8 A E 10,
8 < 1, and the vB
AdamA achieves the following guarantee, for all T 2 1.
D2
a(81 + +1 11-81vT- B2
d N
d
Go Go 32
R(T) S 2a(1 - B1) 1=1
Torst
loarn 6112+
Σ
2a 31)(1 - X)2
Proof.
Following the proof in Adam Appendix Theorem 10.5, by applying Lemma 4, we can get the final convergence bound by substituting Ed
ΣN 2b=1. [1g al|2 for Σi 119..,
Thus, it can be claimed that the convergence rate
R(T O VTT the same as Adam does.
B MORE DETAILS ABOUT ADAMA IN THE DISTRIBUTED DATA PARALLEL SCENARIO
In the distributed data parallel scenario. we pay efforts to make the update effect of m, and U, of AdamA (number of GPUs M, number of microbatches per minibatch N) consistent with the
13

---

Under review as a conference paper at ICLR 2023
update effect of AdamA (number of microbatches per minibatch NM) in single device scenarios. To achieve the effect, we propose a new update method for m and v among different GPUs.
As mentioned in Section 3.3, choose to all-reduce optimizer states instead of gradients at the end of each mini-batch. In this way, the value of m, and V are shown below before optimizer states all-reduced. M equals to the number of GPUs, and N equals to the number of microbatches per minibatch. Other symbols follow our definition Algorithm I, and Gt.i f ( Notice that we will multiply vt-1 by M,32 instead of 32 before the start of each minibatch.
N-1
N
Vofr, Voft,((0-11)
= -1m--1 + (1 (1--80 B1)
Σ
It. = 31m--1 + (1 11--0 31)
N
i=0
N-1
Voft.(00-11)
0 = -3200- (1- (2)
Σ
Φ.. - + (1-11)
N
During all-reduce operations for my we take the average of my from each GPU (add them together and divide by M). For it, we divide by M instead of M after summing i from each GPU. After that, the value of ma, and v, each GPU are:
NM-1
(o.1.(00-11
m, = -1 + B1)
Σ
NM
NM-1
Voft.(00-1)
0 = 320 (1-)
Σ
NM
is easy to find the m, and 2/ keep consistent with the my and V2 from Algorithm 1, as long as we replace N with NM in line 5 It.; (O...( In this way, make the update effect of m and 24 in distributed scenarios (number of GPUs= M, number microbatches per minibatch = N) consistent with the update effect of AdamA in single device scenario (number of microbatches per minibatch = NM). As the convergence properties of AdamA has been proven the same with Adam in single device scenarios, its convergence properties can also be guaranteed in distributed scenarios.
14


Reviews:
[{"review_id": "oUHY5kfrcwb", "reviewer": "Reviewer_6Jc2", "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": "The paper considers the problem of reducing memory usage for training neural networks. It proposes a memory-efficient version (AdamA) of the popular Adam optimizer. The motivation behind AdamA is to enable gradients to be released after computation over each micro-batch (micro-batching computes the gradient sequentially over a mini-batch to reduce the memory footprint of storing the activations over the mini-batch, while still maintaining the convergence properties of the original mini-batch). To enable the gradients to be released, AdamA incorporates the gradients into the optimization state (m_t and v_t for Adam). Allowing for gradient release enables AdamA to use prior work which computes the backward pass in a layer by layer manner, releasing the gradients for previous layers. By using both micro-batching and gradient release, AdamA can lead to considerable memory reductions.\n\nThe paper empirically validates that AdamA has similar convergence properties as Adam on several benchmarks, while being able to reduce the memory usage by order of a few GBs. \n\nThe paper also provides a proof that AdamA has similar convergence behavior as Adam, though I have some issues with that part.\n\n", "summary_of_the_review": "In summary, though the theory in the paper appears to be flawed and this should be addressed, I don't see that as being a fatal flaw. I think the paper could lead to improvements in neural network training and I tend towards acceptance.", "strength_and_weaknesses": "Strengths:\n\n1. The paper provides a simple and clean approach to reduce memory consumption for neural network training. The idea is compatible with several other approaches for reducing memory usage, and could be composed with other methods to yield even more savings.\n2. The experimental evidence appears to be quite thorough. The experiments are systematic, and performed over a variety of vision and NLP tasks. They demonstrate that AdamA has very similar convergence to Adam, while providing substantial memory savings.\n\nWeaknesses:\n\n1. The paper proves convergence of AdamA by following the proof of convergence for Adam. Unfortunately, the proof of convergence for Adam has fatal flaws as pointed out by \"On the Convergence of Adam and Beyond\". I tried to go through the proof, and I think the proof in the paper inherits the same flaws. In particular, \"On the Convergence of Adam and Beyond\" shows that Adam in fact does not always converge, and the same counterexample should apply to AdamA as well. It is possible that some of proposed variants to rectify this convergence issue (such as AMSGrad) can be adapted to this setup as well, but it would have to be analyzed.\n\nTo be clear, I do not believe that the theoretical properties are the selling point of the approach. Instead, the selling point is the strong empirical results, which appear to be sound as far as I can see.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "brsBYy3HH5c", "reviewer": "Reviewer_3RgH", "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 AdamA, a memory-efficient version of Adam when combined with gradient accumulation. The idea is simple: instead of accumulating the gradients after each forward and backward step for the micro-batch, AdamA directly updates Adam optimizer states so that the memory footprint for the gradients will be reduced.  ", "summary_of_the_review": "The idea is simple, straightforward but with limited novelty; some discussion about more direct alternatives is missing (e.g., why not directly use a smaller batch?). Additionally, AdamA actually introduces more communication overhead (running allreduce over m and v introduces 2X communication compared to allreduce over g). ", "strength_and_weaknesses": "Strength:\n\n- The idea is simple and direct, it is clear that the proposed method can reduce the memory footprint.\n\n- The presentation of the paper is clear because of the simplicity of the methodology.\n\nWeakness:\n\n- There is a lack of more direct alternatives, e.g.,  why not simply use a smaller batch? Would AdamA be better than small batch in any case?\n\n- AdamA introduces 2X communication overhead compared with the standard data parallelism. \n\n- The experiments are inappropriately designed in the distributed scenario.  \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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "QgAFZsz4gy", "reviewer": "Reviewer_a8MG", "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 focus on the problem of reducing memory usage in parallel large-scale DNN training. The problem itself is indeed an important research problem that can benefits both the research community and industry. The authors proposed a method (AdamA) to improve upon previous memory reduction methods. The method is that: during gradient accumulation over multiple micro-batch, we can directly integrates gradients into optimizer states and release the gradients. \n\nFrom emipirical study, the authors demonstrated 1) AdamA can further reduce memory compared to previous work, on various settings. 2) AdamA shows similar convergency as Adam baseline. Authors also provide a mathmatical analysis for the convergence. In addition, the authors provide a easy to apply training pipeline that can be used by other researchers. ", "summary_of_the_review": "As mentioned above, I think the quality of the paper, the novelty, and most of the evidence for supporting the claimed contributions are quite good. My major concern lies on the convergence part. For the proposed method to be pratically helpful to the research community and industry, this convergence part needs to strongly supported. If the author can provide explanantions and evidence to address my concerns, I would be happy to accept the paper.", "strength_and_weaknesses": "Strength:\n1. novelty of the proposed method: to the best of my knowledge, this method is novel and adds quite significant value on top of the baseline gradient accumulation method, in terms of memory reduction. It enables combining benefits of different memory reduction methods.\n\n2. the contributions are mostly well-supported.\n\na) [well-supported] the claimed contribution of the proposed AdamA can enable reducing memory footprints of activations and gradients simultaneously, leading to good amount of memory reduction compared to previous work. This claim is well-supported by multiple empirical evidences, including directly comparing memory allocation and the max model size on both common vision and language datasets/models. Examples include ResNet-50 and Bert-large.\n\nb) [well-supported] the proposed method in parallel training setting could have suffered from heavy communication overhead due to the micro-batch all-reduce operation. The authors proposed to address this by updating local optimizer states without micro-batch all-reduce and update once at the end of mini-batch. This method made the throughput of the propoesd method to be similar to original Adam method. The empirical results supports the claim.\n\nc) [still have concerns] the authors demonstrate some evidence that the proposed method has similar convergence as Adam: mathmatically analysis about the convergence and empirically showing similar performance on ImageNet and GLUE. However these evidence may not be sufficient, plesae see the weakness part.\n\n3. clarity, reproducibility, quality are good (as mentioned in the next question).\n\n\nWeakness (or unaddressed concerns):\n\nThe proposed method made two major changes compared to Adam: one is the 2nd momentum accumulation in the optimizer states by replacing the square of the accumulated gradients to accumulating the square of gradients. the other one is in parallel setting, replacing the all-reduce per micro-batch with local update per micro-batch. These two intuitively may have impact on model performance/convergence, and thus need evidence to show their impact is negligible for the proposed method to be useful.\n\nWhile the authors have shown some evidences, I still have some concerns unaddressed:\n\n1. the convergence analysis is done for the first change, but not the 2nd change (even though it only happens with parallel setting, but I think large-scale DNN basically means we use parallel setting all the time). Also the analysis is based on non-pratical assumptions, such as convex cost function. I understand non-convex scenarios are usually not provided in such analysis. But I do want to call out this weakness to emphasize that we would need strong empirical evidence to support the \"similar convergence\" claim.\n\n2. The empirical evidence of convergency versus memory reduction misses some important results. The authors separately shows the convergence results in Sec 4.1 and memory reduction in Sec 4.2. However for some experiments settings, we can only see one but not the other (e.g. we have memory reduction results of Adafactor etc but not its convergence/performance). Could the authors provide the full comparison (note this should not require new experiments, just need all the metrics that should have been obtained during authors' original experiments)? Specifically, examples could be a table of <method, memory reduction, accuracy metrics> and a table of <method, max model size, accuracy metrics>, on top of table 2 and table 3.\n\n3. The evidence shown on the vision domain is not very strong. The authors used ResNet-50 with Adam as baseline. However we know that Adam is much weaker that SGD on ResNet-50 and can only achieve sub-optimal performance on ImageNet (fig. 3 shows <75% but SGD gives more than 75% I remember). Thus ResNet-50 is not a good model to compare AdamA and Adam. If the author would like to strengthen the claim on vision domain, I would recommend to use some models that originally uses Adam as its strongest optimizer, like transformer based models. ", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0hXX3W4pb_", "reviewer": "Reviewer_u5im", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "This paper provides a modification to Adam that allows for gradient updates to be applied \"online\" during gradient accumulation instead of after gradient accumulation has finished. The main selling point of this approach is that it allows for the memory savings from what the authors call \"gradient release\" (also equivalent to the optimization in Zero-2 or Zero-DP P_{g}) to be combined with the memory savings from gradient accumulation.\n\nThe authors demonstrate that this can lead to increased memory savings compared to Zero-DP P_{os+g}, and negligible throughput drop compared to gradient accumulation.", "summary_of_the_review": "See above.", "strength_and_weaknesses": "Overall, the optimization makes sense, and why the optimization delivers improved performance in the author's benchmarks is understandable. It's also useful to know that this modification to Adam *can* be done without impairing its convergence qualities.\n\nUnfortunately, from my understanding of this paper, it seems to be of limited applicability. In particular, AdamA *cannot* be combined with Zero-DP P_{os} (i.e. optimizer state sharding) and gradient accumulation without significantly increasing communication costs.\n\nTo summarize my understanding of AdamA, the primary idea is that it modifies the optimizer state update to allow for online updates of the optimizer state. Thus, instead of needing to compute all of your gradient minibatches before updating the optimizer state, you can update the optimizer state (and thus release the gradient minibatches) before all of your gradient minibatches are completed. *Note* that this does require a change to how data-parallel training is typically done. Typically, data-parallel training accumulates your gradients and then all-reduces your gradients across all devices. However, as AdamA would like to release your gradients as soon as possible, AdamA must all-reduce the *optimizer state* across all devices after finishing microbatches. This is a typically an increase in communication volume, as gradients are only one value while optimizer states (i.e. of Adam) are two floating point values.\n\nTo review Optimizer State Sharding, the idea is that instead of replicating *both* your parameters and your optimizer state across all devices (like typical data parallelism), you only replicate your parameters, and shard your optimizer state across your devices. Then, after you compute your gradients, you can perform a reduce-scatter communicate your gradients to the devices that need them (as only the device with the optimizer state of layer N needs the gradients for layer N), perform your parameter update, and then perform an all-gather to replicate your parameters across all devices again. This, notably, is the *same* communication volume as regular data-parallelism (which only performs an all-reduce), as both reduce-scatters and all-gathers are half the communication volume of an all-reduce. Also, optimizer state sharding does not interfere with gradient accumulation.\n\nNow, the issue here is that, from my understanding, one cannot combine all 3 of optimizer state sharding, AdamA, *and* gradient accumulation without a significant increase in communication volume.\n\nThe way that AdamA avoids communication after each minibatch in gradient accumulation is by doing an online update of the optimizer state. *However*, this can only be done as long as your optimizer state is present on every device, which is not true in Optimizer State Sharding.\n\nThus, although the paper presents memory reduction in Figure 6b through combining AdamA with optimizer state sharding, I suspect that the training throughput is substantially decreased.\n\nOTOH, from Figure 6a, we can see that although it reduces the gradient memory, the bulk of the memory is *optimizer states*. So in that setting, I suspect that using optimizer state sharding would reduce the memory more compared to AdamA.\n\nThe rest of the paper seems clear enough to me, but this is my primary issue with the paper. **As my review heavily relies upon my understanding of this point, if I've misunderstood the paper (or am missing other considerations), please let me know and I would be happy to adjust my rating.**\n\n### Conclusion\n\nThe authors claim that \"AdamA can reduce the memory footprint up to 23% with less than 2% degradation in training throughput [compared to Adam baseline without other memory optimization techniques]\" and that \"AdamA can fit 1.26x larger models over ... Deepspeed baselines\". However, from my understanding of the paper, it seems that AdamA cannot *also* fit 1.26x larger models over Deepspeed baselines without a substantial degradation in training throughput. Thus, it seems to me that I don't see a situation where one would use AdamA over optimizer state sharding. \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject