ALAM: Averaged Low-Precision Activation for Memory-Efficient Training of Transformer Models

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q1) The complexity of ALAM is not analyzed. And the paper only compares training time for LLaMA-3B.

A1) We thank the reviewer for this valuable feedback. To address the reviewer’s concern, here we first analyze the time complexity of ALAM in more detail. Please note that the sensitivity calculation is conducted only twice (before training and at 10% of the training process), and hence its overhead is excluded from the analysis. Our ALAM utilizes average quantization, which involves average calculation and quantization including searching for min and max values, multiplication with a scaling factor, and stochastic rounding. These operations have a time complexity of $\mathcal{O(n)}$, where n represents the number of activations. The authors in [1] calculated the FLOPs and the number of activations of a transformer block with flash-attention as $24bsh^{2} + 4bs^{2}h$ and $34bsh$, respectively, where $b$, $s$, and $h$ denote the batch size, sequence length, and hidden size, respectively. Based on this and the fact that the size of the hidden dimension is generally larger than the sequence length in LLMs, the time complexity of training an LLM and the additional time complexity due to ALAM are calculated as shown in Table E.1 below.

Table E.1. Time complexity of training and activation compression required by ALAM.

Therefore, as the model size $h$ increases, the training time overhead due to ALAM becomes relatively smaller. To verify this, we trained models larger than 7B on the LIMA (1K) dataset for 5 epochs, and the results are displayed in Table E.2 below.

Table E.2. Time for fine-tuning LLaMA2-7B for 5 epochs on the LIMA datasets using QLoRA with different ACT frameworks with a single NVIDIA A6000 GPU. The micro batch size is set to 2.

Table E.2 shows that while there is a 23% time overhead for the LLaMA2-7B model, the overhead is reduced to 18% and 12% for the LLaMA2-13B and LLaMA-30B models, respectively. We once again thank the reviewer for this great suggestion, and we will include this analysis in the camera-ready version if the paper is accepted.

[1] Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., & Catanzaro, B. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.

Dear reviewer,

Thank you once again for carefully reviewing our submission. As the reviewer-author discussion period ends in less than 12 hours, we would be very grateful if we could hear back for any thoughts. If possible, please take a look at our responses and see if they adequately address your concerns. We are also open to any further feedback or queries you may have.

Thank you,

Authors of Paper #3660

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q1) The primary concern is the lack of novelty. The method relies heavily on existing frameworks, such as GACT and ActNN, using the same concept of quantization for activation compression. Additionally, the proposed quantization method is also a well-known clustering approach. The theoretical analysis, suggesting the minimization of quantization MSE (as shown in equation 6) appears rather straightforward. While the authors did introduce some tricks, such as gradient normalization variance, to improve the existing framework, the overall technical contribution seems to be relatively modest.

A1) We thank the reviewer for this valuable feedback. We agree with the reviewers that our AQ algorithm has been developed and proven to work in a similar way to ActNN [1] and GACT [2]. However, we would like to point out that the core idea of our algorithm, simply compressing a group of adjacent activations using a group average value without the need for extensive iterations to categorize similar values, is entirely new and has not been reported in the literature. In other words, to the best of our knowledge, our work is the first attempt to compress activations through simple averaging.

Unlike a previous method [3] that employs K-means clustering to group weights, which requires a time-consuming process to find weights with similar values, our novel activation compression method relies solely on simple averaging of adjacent activations without the need to identify similar values, which can be very efficiently implemented. In addition, this approach is validated mathematically, proving that simple averaging can minimize gradient variances.

We also agree that the minimization of quantization MSE in Eq. 6 looks straightforward, but actually deriving it was not very straightforward, as shown in Appendix A.2. More specifically, we mathematically proved that the gradient variance of batch normalization and layer normalization has the form of Eq. 6, which was not proven in prior work [1, 2, 4]. This was a key contribution that allowed us to prove that representing activation as a group average can minimize gradient variance for the first time.

Lastly, to the best of our knowledge, we are the first to demonstrate an effective activation compression algorithm that supports LLMs with more than 7B parameters. Unlike model compression algorithms like LoRA and QLoRA, previous activation compression algorithms such as ActNN, MESA, and GACT were not proven to work with LLMs. Our experimental results confirm that ALAM can train large language models such as LLaMA2-7B, LLaMA2-13B, and LLaMA-30B even at 1-bit, whereas conventional methods such as GACT fail to train the models.

Q2) Could the authors explain the overhead involved in doing the group-wise quantization?

A2) This is an important point, and we thank the reviewer for bringing up this issue. Average Quantization generates sub-1b values by first compressing adjacent activations to a group average value and then performing conventional quantization to the group average value. We suspect that the reviewer is wondering if there is any additional memory or time overhead involved in this process. First, there is no additional memory overhead for compressing activation. As explained in Sections 3.1 and 3.2, we group adjacent activations in the flattened activation vector during the forward pass. During the backward pass, they are restored simply by repeating the group average values. For example, when the group size is 2,

• Averaging in the forward pass: [4, 6, 7, 1, 6, 8] -> [5, 4, 7]
• Restoring in the backward pass: [5, 4, 7] -> [5, 5, 4, 4, 7, 7]

Therefore, there is no need to indicate specific locations, and hence our algorithm does not save any masks. Second, our Average Quantization requires additional time to compute the group average value. However, its time overhead is similar to that of the activation quantization process in conventional activation compression schemes, which is in line with experimental results shown in Table D.1 below.

Table D.1. Time for fine-tuning LLaMA2-7B over 5 epochs on the LIMA datasets using QLoRA with different ACT frameworks with a single NVIDIA A6000 GPU. The micro batch size is set to 2.

Dear reviewer,

Thank you once again for carefully reviewing our submission. As the reviewer-author discussion period ends in less than 12 hours, we would be very grateful if we could hear back for any thoughts. If possible, please take a look at our responses and see if they adequately address your concerns. We are also open to any further feedback or queries you may have.

Thank you,

Authors of Paper #3660

Q2) How to choose the parameter group is unclear. Is this parameter ‘group’ related to the sensitivity? What if the adjacent activations have little similarity?

A2) We thank the reviewer for raising an important issue, and we apologize for the confusion. As described in Sections 3.1 and 3.2, we group adjacent activations together after simply flattening the activations, without any other complicated processes. As the reviewer insightfully pointed out, varying the group size (g) and the precision of the group average (n) could produce different bit precisions (AQ-n/g bit). However, in all experiments using Average Quantization, we fixed the group size to 4 and performed 2-bit quantization on the group average value to obtain AQ 0.5-bit, which is an optimal configuration identified in experiments, as shown in Table C.2.

Table C.2. Performan comparisons of different AQ 0.5-bit configurations when training VGG-11 on CIFAR-100 (n: precision of group average, g: group size).

The group size is not related to sensitivity. More specifically, the sensitivity of each layer is first calculated using the GradNormVar algorithm, which determines the layer’s bit precision. In the layers with low sensitivity, we quantize activations into AQ 0.5-bit, where average quantization is performed with a fixed group size of 4 as discussed above.

The reviewer is certainly correct that adjacent activations may show little similarity. This could be mitigated by only grouping activations with similar distributions, but we did not implement this in our algorithm since this process could be time-consuming and require additional memory to save the mask for recovery. Our simple grouping scheme may lead to some loss of information in activations, but the impact on actual accuracy is minimal because ALAM assigns AQ 0.5-bit only to layers with very low sensitivity. Furthermore, the quality of AQ 0.5-bit activations is still higher than that of 1-bit quantized activations, as shown in Table C.2.

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q1) The evaluation part is not sufficient. This paper only compares the result with GACT, which is based on activation compression. However, for model compression, other techniques like pruning and low-rank compression are widely used.

A1) We would like to thank the reviewer for bringing up an important point. Following the reviewer’s advice, we conducted detailed comparisons between ALAM and other model compression methods. For comparisons, we selected LoRA [1], which is the state-of-the-art low-rank compression method. We also implemented QLoRA [2], which performs 4-bit quantization for weights and only computes weight gradient for LoRA parameters. QLoRA allows for a 4× reduction in the weights with comparable accuracy, making it widely used for fine-tuning LLMs in constrained environments. The results are presented in Table C.1. Please note that ‘average accuracy’ represents the average accuracy over all tasks in Tables A.1 and A.2 in the global response. Also, ‘param mem’ represents the parameter memory occupied by weights, gradients, and optimizer states during training.

Table C.1. Comparisons with model compression techniques in fine-tuning LLaMA2-7B and LLaMA2-13B on the LIMA dataset. The micro batch size is set to 8.

In experiments, LoRA noticeably reduces the memory allocated for gradients and optimizer states, which results in a 4× compression of parameter memory during training. Additionally, QLoRA exhibits a large compression rate of 16× by utilizing 4-bit weights. However, neither LoRA nor QLoRA compresses activation memory. Our ALAM, on the other hand, focuses on compressing activation memory, and hence achieves 10.6× and 10.3× compression rates on LLaMA2-7B and LLaMA2-13B, respectively. As a result, compared to when only LoRA and QLoRA were applied, additionally employing ALAM further saves 56% and 75% of the total memory with comparable accuracy for LLaMA2-7B and LLaMA2-13B, respectively. In summary, our ALAM is orthogonal to other model compression techniques, and they can be applied simultaneously for maximum memory savings.

Please note that only a handful of pruning schemes supporting very large models have been reported in the literature, with limited memory savings. For instance, LLM-Pruner [3] exhibits a significant drop in accuracy even when pruning just 20% of the parameters in LLaMA-7B, translating to similar saving only in the parameter memory.

Dear reviewer,

Thank you once again for carefully reviewing our submission. As the reviewer-author discussion period ends in less than 12 hours, we would be very grateful if we could hear back for any thoughts. If possible, please take a look at our responses and see if they adequately address your concerns. We are also open to any further feedback or queries you may have.

Thank you,

Authors of Paper #3660

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q1) It would be better to add QLoRA as a baseline for evaluation. Also, it is unclear how the proposed ALAM performs on larger models, such as LLaMa-13B, LLaMa-30B, etc.

A1) We thank the reviewer for this great suggestion. Following the reviewer’s comment, we applied QLoRA [1] and evaluated our algorithm on larger models: LLaMA2-13B and LLaMA-30B. Please note that we used the recently released LLaMA2 for a 13B model, which has the same structure as LLaMA but is pre-trained on a broader range of datasets, leading to improved accuracy. On the other hand, LLaMA2-30B is not publicly available, and hence LLaMA-30B was employed instead. In this experiment, we fine-tuned the model for 5 epochs on the LIMA dataset. The results are displayed in Table B.1. In the table, 'precision' refers to the target average precision of activation during training as detailed in the manuscript, not the precision of the weight.

Table B.1. Test accuracy, activation memory (ACT mem) with compression, and memory for calculating sensitivity (sens mem) in fine-tuning LLaMA2-13B and LLaMA-30B on the LIMA dataset. The micro batch size is set to 2.

Experimental results demonstrate that our ALAM (1-bit) compresses activation by 10.3× in LLaMA2-13B and 9.9× in LLaMA-30B while closely matching the baseline in accuracy for all tasks. We will include these updated experimental results in the camera-ready version if the paper is accepted.

[1] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, & Luke Zettlemoyer (2023). QLoRA: Efficient Finetuning of Quantized LLMs. CoRR, abs/2305.14314.

Dear reviewer,

Thank you once again for carefully reviewing our submission. As the reviewer-author discussion period ends in less than 12 hours, we would be very grateful if we could hear back for any thoughts. If possible, please take a look at our responses and see if they adequately address your concerns. We are also open to any further feedback or queries you may have.

Thank you,

Authors of Paper #3660

(2) ALAM with QLoRA for fine-tuning larger models such as LLaMA2-13B and LLaMA-30B

Following the comments of the reviewers HytY and Eyz9, we applied the widely-used model compression method, QLoRA [1], and performed experiments on LLaMA2-13B and LLaMA-30B models. Please note that since the LLaMA2-30B model is not publicly available, we conducted experiments on the LLaMA-30B model instead. In this experiment, we fine-tuned the models for 5 epochs on the LIMA dataset. The experimental results are displayed in Table A.2, where 'precision' refers to the target average precision of activation during training as detailed in the manuscript, not the precision of the weight.

Table A.2. Test accuracy, activation memory (ACT mem) with compression, and memory for calculating sensitivity (sens mem) in fine-tuning LLaMA2-13B and LLaMA-30B on the LIMA dataset. The micro batch size is set to 2.

Experimental results demonstrate that our ALAM (1-bit) compresses activation by 10.3× in LLaMA2-13B and 9.9× in LLaMA-30B while closely matching the baseline in accuracy for all tasks. We will include these updated experimental results in the camera-ready version if the paper is accepted.

[1] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, & Luke Zettlemoyer (2023). QLoRA: Efficient Finetuning of Quantized LLMs. CoRR, abs/2305.14314.