Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 3: reject, not good enough
review_content: 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.

The 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.
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.

Unfortunately, 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.

To 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.

To 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.

Now, 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.

The 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.

Thus, 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.

OTOH, 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.

The 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.**

### Conclusion

The 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. 



clarity,_quality,_novelty_and_reproducibility: See above.
summary_of_the_review: See above.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.

The 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. 

The paper also provides a proof that AdamA has similar convergence behavior as Adam, though I have some issues with that part.


strength_and_weaknesses: Strengths:

1. 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.
2. 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.

Weaknesses:

1. 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.

To 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.
clarity,_quality,_novelty_and_reproducibility: Clarity: The paper is quite well-written. It explains the previous approaches to reduce memory usage quite well, and places the proposed work properly in that context. Several figures, diagrams and detailed algorithm descriptions make the approach quite easy to understand.

Originality: The paper borrows quite a lot from previous approaches, but as the experimental results demonstrate the proposed modifications lead to significant improvements in practice.

Minor question:

Can the authors comment on the running time of AdamA, as compared to Adam, in some of the considered settings?
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.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: 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.  
strength_and_weaknesses: Strength:

- The idea is simple and direct, it is clear that the proposed method can reduce the memory footprint.

- The presentation of the paper is clear because of the simplicity of the methodology.

Weakness:

- 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?

- AdamA introduces 2X communication overhead compared with the standard data parallelism. 

- The experiments are inappropriately designed in the distributed scenario.  


clarity,_quality,_novelty_and_reproducibility: The presentation is clear because of the simplicity of the methodology.

This is a mismatch of AdamA's goal and the empirical study, especially for the distributed scenario. Why is the benchmark based on the combination with ZeRO-DP instead of more memory-efficient stages in ZeRO to support larger-scale models (where the memory reduction technique does matter)? For example, ZeRO-S3 (which is also known fully sharded data parallelism)? Can AdamA be combined with ZeRO-S3?

The source code is attached to provide good Reproducibility.
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).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 

From 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. 
strength_and_weaknesses: Strength:
1. 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.

2. the contributions are mostly well-supported.

a) [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.

b) [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.

c) [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.

3. clarity, reproducibility, quality are good (as mentioned in the next question).


Weakness (or unaddressed concerns):

The 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.

While the authors have shown some evidences, I still have some concerns unaddressed:

1. 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.

2. 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.

3. 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. 
clarity,_quality,_novelty_and_reproducibility: Clarity and reproducibility of the paper is good. The paper gives details and sources of the experiments and also provide source code to reproduce. Overall writing quality is good as well: easy to follow the main claims and their supports.
Novelty of the paper: to the best of my knowledge (however I am not expert in memory reduction in DNN), the method of intergrating the gradient accumulation into the optimier states is novel. The claimed contributions around this method are mostly well-supported (as above mentioned). So the novelty of the paper is good in my opinion.
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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper provides a memory-efficient version of Adam which allows for gradient updates to be applied online during gradient accumulation. The modified version (AdamA) updates optimizer states directly, instead of accumulating the gradients after each step, so that the memory footprint of gradients can be reduced.

Strengths:
1. The modification makes sense in somehow and why the modification delivers performance improvement is understandable.
2. The author provides experimental results over vision and NLP tasks to demonstrate the effectiveness of the modification in some certain cases.
3. In general, idea is clearly presented. The paper is well written and easy to follow. The presentation is good.

Weaknesses:
1. Although the proposed method provides value in some scenarios, it has limited applicability in general cases in practice. The method mainly focuses on the case of few-GPU (or single-GPU) scenarios with trainable model size. Those cases are not common enough and the benefits are not large enough. 
2. The method introduces more communication overhead. As we could have the scenarios where communication mainly affects the throughput, this is a drawback of the method.
3. The theory analysis provided by the paper for the convergence has flaws, which comes from the original Adam. This is not addressed by the authors.

justification_for_why_not_higher_score: Due to the following weakness, I would like to recommend a reject for this paper: 1) limited applicability, 2) drawback with more communication overhead, and 3) theoretical analysis should not be counted as a contribution.
Decision: Reject