ALAM: Averaged Low-Precision Activation for Memory-Efficient Training of Transformer Models
Authors: Sunghyeon Woo
Sunwoo Lee
Dongsuk Jeon
Venue: ICLR 2024 poster (2024)
Keywords: Memory efficient training
Activation-compressed training
Average Quantization
NLP
Transformer
Primary area: general machine learning (i.e., none of the above)
TL;DR: ALAM compresses activations to their group average with a lightweight sensitivity calculation, achieving up to a 10x activation memory reduction in LLMs.

Abstract: One of the key challenges in deep neural network training is the substantial amount of GPU memory required to store activations obtained in the forward pass. Various Activation-Compressed Training (ACT) schemes have been proposed to mitigate this issue; however, it is challenging to adopt those approaches in recent transformer-based large language models (LLMs), which experience significant performance drops when the activations are deeply compressed during training. In this paper, we introduce ALAM, a novel ACT framework that utilizes average quantization and a lightweight sensitivity calculation scheme, enabling large memory saving in LLMs while maintaining training performance. We first demonstrate that compressing activations into their group average values minimizes the gradient variance. Employing this property, we propose Average Quantization which provides high-quality deeply compressed activations with an effective precision of less than 1 bit and improved flexibility of precision allocation. In addition, we present a cost-effective yet accurate sensitivity calculation algorithm that solely relies on the L2 norm of parameter gradients, substantially reducing memory overhead due to sensitivity calculation. In experiments, the ALAM framework significantly reduces activation memory without compromising accuracy, achieving up to a 10$\times$ compression rate in LLMs.

Published as a conference paper at ICLR 2024
ALAM: AVERAGED LOW-PRECISION ACTIVATION FOR MEMORY-EFFICIENT TRAINING OF TRANS- FORMER MODELS
Sunghyeon Woo, Soonwoo Lee, Dongsuk Jeon Scoul National University, Scoul, Korea {wsh0917, ori915, djecni](snu: ac..k ac.kr
ABSTRACT
One the key challenges deep neural network training the substantial amount of GPU memory required to store activations obtained in the forward pass. Vari- ous Activation CrwvattoonCoopresseee Training (ACT) schemes have been proposed to mit- igate this issue; however, it is challenging to adopt those approaches in recent transformer-based large language models (LLMs), which experience significant performance drops when the activations deeply compressed during training. In this paper, we introduce ALAM, a novel ACT framework that utilizes aver- age quantization and lightweight sensitivity calculation scheme, enabling large memory saving in LLMs while maintaining training performance. We first demon- strate that compressing activations into their group average values minimizes the gradient variance. Employing this property, we propose Average Quantization which provides high-quality deeply compressed activations with an effective pre- cision of less than bit and improved flexibility of precision allocation. In addi- tion, we present a cost-effective yet accurate sensitivity calculation algorithm that solely relies on the L2 norm of parameter gradients, substantially reducing mem- ory overhead due to sensitivity calculation. In experiments, the ALAM framework significantly reduces activation memory without compromising accuracy, achiev- ing to a 10x compression rate in LLMs.
INTRODUCTION
In recent years, deep learning has demonstrated human-like or even better performance on language- related tasks using large language models (LLMs) (Chen et al., 2023] [mith et al.] [022] OpenAI 2023 Touvron et al., 2023a b). Behind this success lies various efforts to increase the model size since it directly improves model performance due to scaling laws Kaplan et al. (2020); [Sorscher et al. (2022a b). However, the amount of memory required for training also increases proportion- ally, hindering their practical implementation. One of the reasons behind the memory bottleneck during training is that the backpropagation algorithm (Kelley) 1960) needs to store all intermedi- ate activations generated during the forward pass in memory for later use in the backward pass for calculating parameter gradients, which leads to extensive memory requirements. For instance, as shown in Fig. [ GPT-3 (Chen et al. (023)) with 175B parameters and MT-LNG (Smith et al. 2022) with IT parameters require 67.3 GB and 132.7GB activation memory, respectively. exceeding the memory occupied by the parameters and optimizer state (Korthikanti et al. 2022). This issue even worsens with larger micro-batch size or sequence length because the activation memory also creases proportionally whereas the memory occupied by the parameters and optimizer state remains unchanged. Therefore, is crucial to reduce activation memory when it comes to training L.LMs.
Activation rematerialization (Chen et al. 2016 Jain et al. [019 Feng & Huang 2021 and reversible networks (Gomez et al.) [0177] Kitaev et al., 2020) Sander et al. 2021 Cai et al. 2023) store only part of the activations, recomputing the rest during the backward pass. These methods demand extra computation the backward pass and reduce training speed. Reduced- Reducce-preccsionn training (Micikevi- cius al.) 2018 Wang al.) 2018 Chen al. 2020 Sun et al. 2020) aims reducing computation precision and training memory by representing each variable in training (e.g., weight, error, activa- tion, and gradient) with low-precision data formats such FP8. However, training accuracy drops

---

Published as a conference paper at ICLR 2024
addivation memory parameters and optimizer state memory
Forward propagation
Wegger
Memory (GB)
120
1. Calculate Sensitivity by GradNormVar
Alloocee bits with 0.5-bit
67.3 GB
114 GB
132.7 GB
80
0.5-bit
0.5-bit
2
4-bil
6-bit
3
40
1
save for backward
31.5 GB
31.5 CB
Memory
1758
530B
Decompress Activation
GPT-3
MT-LNG
Figure 1: Memory usage breakdown of GPT-3 (22B/175B) and MT-LNG (530B/IT) with data and model parallelism. The dashed red line de- notes the 80GB capacity of the NVIDIA A100 GPU.
load
Backward propagation
Figure 2: ALAM training framework.
quickly with lower precision, and optimized kernels for low-precision operations are required to fully utilize its advantage in training speed. On the other hand,
Training
(ACT) (Chakrabarti & Moseley 2019; Chen et al.) 2021 Liu et al. 2022b Pan et al.] 2021 Liu et al. 2022a) solely aims to reduce activation memory usage by compressing the activations before storing them in the forward pass. Conventional ACT methods, however, exhibit large performance drops for high compression rates when applied LLMs. For example, MESA (Pan al.) 2021)
compresses all layers uniformly with an identical compression rate, noticeably undermining training performance for activation precision under bits. GACT (Liu et al. 2022a) determines compression rates (i.c.. allocated bit precision) of each layer based on its sensitivity, achieving state-of-ther of-tth--rt compression rate of bits transformer models. However, deeply quantized activations noticeably affect training performance, and the lowest bit precision it can allocate to less-sensitive layers is bit, which limits the flexibility of bit precision allocation and hinders further compression.
In this work, we propose new ACT framework, ALAM, which greatly improves compression rates while maintaining training performance. We first propose to represent a group of activations with their group average value as a means to compress them deeply and mathematically show that this scheme minimizes gradient variance due to activation compression. Based on this observation, we propose Average Quantization (AQ), which replaces group of activations with their group average to realize sub-1b precision effectively. We experimentally demonstrate that the sub-1b activations obtained by AQ maintain high quality, even surpassing that of conventional 1-bit quantized acti- vations. Moreover, by assigning sub-1b precision to the layers with very low sensitivity, we can further enhance the overall compression rates. We also introduce the GradNorm Var algorithm that determines the sensitivity of a layer based on the variance of the L2 norm of parameter gradients. This approach significantly reduces the memory required for sensitivity calculation, when compared to GACT (Liu et al.) (202a)) which requires the variance of parameter gradients from multiple seeds. Experimental results confirm that our ALAM outperforms the ofate-of-the-arr method various transformer models including LLaMA2-7B/13B and LLaMA-30B, In summary, the main contribu- tions of our paper are:
We propose Average Quantization, which generates high-quality sub- bb activations, abling further compression without sacrificing accuracy. To the best of our knowledge, this is the first attempt compress activations through simple averaging.
We propose GradNorm Var, lightweight sensitivity calculation algorithm that employs the variance of the L2 norm of parameter gradicnts, eliminating the need to retain all parameter gradients and substantially reducing memory usage.
2

---

Published as a conference paper at ICLR 2024
The proposed ALAM framework is demonstrated to successfully compress activations in various transformer models, achieving a 9.9x compression rate LLaMA-30B.
ACTIVATION-COMPRESSEE TRAINING
2.1
FORMULATION
We represent the loss function of forward propagation in a neural network as Xp), given input Xy and parameters 0. When the stochastic gradient descent (Ruder 2016) is applied, the parameters are updated as 0 = 0' V L(0", Xo), where denotes the iteration number and v) is the learning rate. Then, the gradients are calculated as follows:
C Ve G, 8, X1-1)
Here, X₁ and 0n denote the activation and parameter of the (-th layer, respectively. Gj denotes the gradient function of the (-th layer, generating the input gradients V Vx,-1 and parameter gradients V. from the output gradient Vx, C. parameters fj, and activations X₁-1- For the sake of simplic- ity, denote L(0, Xp) as L. Given Eq. 1 all layer activations X;-1 must be stored in memory during the forward pass to compute gradients, leading to substantial memory use. The ACT frame- works are designed to tackle this memory consumption issue by compressing activations follows:
vx,,C. G X1-1)
where X represents the compressed activation, and V represents the gradients obtained by ACT. For instance, if we compress activations from FP32 to INTI, a 32x reduction in activation memory usage can be achieved.
2.2 CONVERGENCE OF ACTIVATION- TRAINING
If L is continuously differentiable, VLo is L-Lipsehitz for some constant I,, L is bounded below by Linf, and there exists o² > 0 satisfying Var VoC] S o² for all parameters, then we have the following convergence theorem (Chen et al. 2021):
E [ < 2(L - Linf) + ao
(3)
at
As the iteration & increases, the first term of the bound in Eq. 3 decreases minimizing the gradient variance Var [aL] critical for the convergence of ACT. For stochastic quantizer, Liu et al. (2022a) demonstrated that the gradient variance can be expressed as:
Therefore,
Var [ar[Vol] [oLc
S
Σ a
-12
where by denotes the bit precision (i.e., the number of bits to present activation) of the I-th layer. The sensitivity o is defined how sensitive it is activation compression. If the sensitivity is high, the gradient variance is larger at the identical bit precision, undermining training performance and requiring higher bit precision to mitigate it. Therefore, we can determine the optimal bit precision of each layer based on its sensitivity to maximize memory savings with minimal training performance degradation by reducing gradient variance Eq. 4 Juuiing a greedy algorithm.
GACT (Liu et al. 2022a) calculates sensitivity qi by observing the variance of parameter gradients from multiple runs of actual training using different seeds for stochastic rounding. More specifically, it initially compresses activations using specific seed and calculates parameter gradients. Then, the algorithm changes only the seed for compressing the 1-th layer's activations and recalculates parameter gradients. Finally, it treats the variance of all parameter gradients between these seeds as the sensitivity of the 1-th layer, indicating how changes in the [-th layer's 's activation impact gradient variance. This process is detailed in Algorithm
3

---

Published as a conference paper at ICLR 2024
Algorithm 1 Sensitivity calculation in GACT
Require:
9: parameter gradient of i-th element in I-th layer.
P1, P2: two random seeds
L: number of layers, ry: elements in 1-th layer
b (bi)f21: compression scheme
Ensure:
s , sensitivity of /-th layer
for - 1 to L do
Set rı for compressing all activation X, run backprop and store ho + (g." nu, S 1}} 12 only for compressing (-th layer's activation æ1, recalculate h1 10 ni, L}
Change
(()) = 2 |/ko-
ku (2 1)²
end for
AVERAGE QUANTIZATION
3.1 MINIMIZING GRADIENT VARIANCE
In this section, we propose a new approach to activation compression and show that it minimizes the gradient variance in ACT. Chen et al. (2021) proved that the gradient obtained by ACT is unbiased (i.e., E [oo] = E Vo))
Following this property, the gradient variance of ACT can be expressed as:
Var [Voc] = Var [VC] + +[[IAV<211]
(5)
where A AVoL VoL VoL. The details can be found in Appendix A On the right-hand side of Eq. 5 the first term is the gradient variance of SGD, which is not related to activation com- pression. Consequently, the second term, representing the gradient difference caused by activation compression, plays an important role in minimizing the gradient variance in ACT. Evans & Aamod! (02T)) derived boundary functions for gradient differences in commonly used layers in neural net- works, such as cully-connected, convolution, batch normalization, and layer normalization layers. We further extend these boundary functions in Appendix A and prove that the boundary function of gradient difference can be simply expressed as
o ( 2
6
where AX = X X, and d is positive constant. Now suppose divide activations into several groups and represent each group using a single representative value, where its optimal value can be determined using the theorem below.
Theorem 1. The value of the boundary function of gradient difference is minimized when each group of activations is approximately represented using its group average value.
Theorem can be proven by showing that the derivative of d|| with respect to an approximate value is equal to 0 when using average as the approximate value because D a convex func- tion with respect the approximate value. Detailed proof can be found in Appendix(A Furthermore, we experimentally demonstrate this property by training a 512-512-10 MLP (Haykin 1994) model on MNIST (Deng. 2012). We approximate each group of activations using various values: average, median, minimum, maximum, and randomly selected activation in a group, where we group adja- cent elements in the flattened activation vector, and the group size is set to 32. Other experimental details are provided in Appendix D Experimental results in Table confirm that averaging is the most effective compression method, which achieves high training accuracy close to full-precision training.
Table 1: Test accuracy of MLP trained on MNIST with different activation approximation methods.
FP32
Average 96.09
Median
Minimum
Maximum
Random
Accuracy (%)
98.72
84.82
fail
fail
90.31
4

---

Published as a conference paper at ICLR 2024
3.2 AVERAGE QUANTIZATION
Average (/g)
n-bit quantize
(n/g)-bit
AQ 0.x-bit
Figure Average Quantization for high-quality sub-lb precision.
In Section 3.1 we show that averaging is an effective method to compress activation with mini- mal impact on training performance. By adopting this property, we propose Average Quantization, which generates high-quality compressed activation with sub- 1b precision. We split activations into multiple groups and only store a single activation for each group, which is an average value of all the elements in the group. For instance, if we take an average for group of four activations and quantize it to 2 bits, we only need to store bits for group. This achieves an identical compression rate to when compress all of four activations separately to 0.5-bit precision in conventional ac- tivation compression, and hence we denote this as AQ 0.5- Note that our Average Quantization flattens and groups adjacent activation elements during the forward pass, and then recovers them by repeating the group average in the backward pass without requiring an additional mask or spatial index for recovery.
We experimentally show that our approach produces high-quality compressed activation. In experiments, we train VGG-11 (Simonyan & Zisserman 2015) on CIFAR- 100 (Krizhevsky al.) 2009) using various types of compressed activations and compare test ac- curacy and validation loss. Fig. ddisplas the results for 2-bit uniform quantization, 1-bit uniform quanti- zation, AQ I-bit, and AQ 0.5-bit. Here AQ 1-bit and
70
60-
Test Accuracy (9)
5o-
FP32
40
2-bit
1-b
AQ 1-bit
AO 0.5 bit
20-
AQ 0.5-bit activations are obtained by averaging a group of two and four activations, respectively, and
10-
then quantizing the average value into 2 bits. Ex- perimental details are provided in Appendix [ 1-bit uniform quantization fails to train the network, but AQ I-bit succeeds in training, closely matching 2-bit uniform quantization. Even AQ 0.5-bit successfully trains the network without divergence, outperform- ing 1-bit uniform quantization. Therefore, we expect to replace 1-bit quantization with AQ 0.5-bit without sacrificing accuracy and allocate more bits high- sensitivity activations for given memory budget to improve training performance.
kpoch
Figure 4: Comparisons of various com- pressed activations in VGG-11 training on CIFAR-100. For a fair comparison, activa- tions of linear, convolutional, and batch nor- malization layers in VGG-11 are uniformly compressed in all experiments.
EFFICIENT SENSITIVITY CALCULATION
In ACT frameworks, it is crucial to accurately determine the sensitivity of layer, as it greatly af- fects training performance. As discussed in Section (.2iionm222aaaa and Algorithm [] GACT (Liu et al. (2022a) determines sensitivity by calculating the gradient variance from multiple runs of actual training with compressed activations using different seeds, and hence it must temporarily store all parameter gra- dients for each seed. Consequently, the peak memory usage substantially increases when calculating the sensitivity in models with a large number of parameters. For instance, GACT requires 700 GB memory for calculating sensitivity from two seeds when training GPT-3 with 175B parameters. To alleviate this issue, we propose the GradNorm Var algorithm, which determines the sensitivity of layers as the variance of L2 norm of parameter gradients obtained from multiple seeds, as detailed in Algorithm? 2 More specifically, our algorithm starts by compressing activations using random seed. Instead of storing all parameter gradients, the algorithm calculates and retains only the L2 norm of parameter gradients for each layer. Subsequently, different seed is used for compressing the (-th layer's 's activations, and the L2 norm of parameter gradients recalculated. Finally, the algo- rithm employs the variance in the 1.2 norm of parameter gradients between these seeds as sensitivity
5

---

Published as a conference paper at ICLR 2024
of the 1-th layer. Since GradNorm Var only stores the L2 norm of parameter gradients for each layer, rather than all parameter gradients from multiple seeds, it can significantly reduce the amount of memory required for caleulating sensitivity.
Algorithm 2 Proposed GradNorm Var algorithm
Require:
GO; L2 norm of parameter gradients g() in 1-th layer.
Fi, P2; P2: two random seeds
L: number of layers, ny: elements in I-th layer
b= (bi)f compression scheme
Ensure:
8 ( sensitivity of 1-th layer
for 1 to T. do
Set ri for compressing all activations X, run backprop and store ko +- {G 11gg)) A 1} Change r1 to 12 only compressing activation zı, recalculate h1 t {G()) = l19 ||2l & L}
((4)) 2 l ha (2 1)a
end for
To validate the effectiveness of our GradNorm Var algorithm, we evaluated the sensitivity and bit precision allocation using the conventional method in GACT that employs gradient variance and our algorithm during the fine-tuning BERT-Large (Devlin et al. 2019) on the MRPC dataset (Wang al. 2019). As illustrated in Fig. S the sensitivity determined by our algorithm represents a similar trend to that of the conventional approach with correlation of 0.9981, leading to comparable al- location. This is consistent with actual training performance; our GradNorm Var algorithm achieves 86.9% training accuracy, closely matching the performance of conventional sensitivity calculation scheme in GACT (86.5%).
30
Proditon (abs)
10"4
221
SANS-OMS
10
0.5
0.5
20
40
05
30
100
320
340
180
60
20
350
1.00
140
300
Layer
Layer
(a) Gradient variance (GACT)
(b) Gradient normalization variance (ours)
Figure 5: Sensitivity and bit preicision allocation comparisons when the target average precision is 2 bits in BERT-Large fine-tuning MRPC datasets. Note that adopt AQ 0.5-bit instead of conventional 1-bit quantized activation in both experiments.
5 ALAM: AVERAGED LOW-PRECISION ACTIVATION FOR MEMORY-EFFICIEN TRAINING
Combining the proposed techniques, we introduce ALAM, a new ACT framework that achieves a significant amount of memory saving while maintaining training accuracy. When the user specifies the target average precision, which represents the average precision of activations across all layers, the proposed ALAM framework automatically carries out ACT using the Average Quantization and GradNorm Var algorithms. The detailed process as follows:
1. (Calculate sensitivity) The sensitivity of each layer is calculated using the GradNorm Var algorithm.
2. (Allocate bit precision) The bits allocated to layers based their sensitivity in way that it minimizes the gradient variance in Eq. 4 using a greedy algorithm. AQ 0.5-bit is employed in low-sensitivity layers instead of conventional 1-bit quantization.

---

Published as a conference paper at ICLR 2024
3. (Training) During the forward pass, activations compressed based on the allocated bit precision, and then decompressed the backward pass to update parameters.
Note that calculating sensitivity and allocating bit precision oceur only before training and after 10% of total training epochs because activation sensitivity saturates early (Liu et al. 2022a).
RESULTS
In this section, we evaluate our ALAM framework by training diverse transformer models on various tasks. More details on the experiments can be found in Appendix D
We fine-tuned BERT models with different sizes, including DistiBBERT-HE Base Sanh et al. 2019), BERT-Baas, and BERT- (Devlin et al.) 2019), on the datasets in the GLUE benchmark (Wang et al. 2019) using three approaches: without activation compression (baseline), with GACT, and with ALAM. We evaluate test accuracy, activation memory, and the memory re- quired for calculating sensitivity various target average precisions for activations. Note that GACT and ALAM only compress float tensors, while they do not compress the other values such as masks and indices. Therefore, reducing the target average precision from 32 bits to bits does not always result in a 16x compression of activation memory, and the actual compression rate can vary by model. Table shows that our ALAM achieves a significantly better trade-off between accuracy and memory savings compared to GACT in all cases. In DistilBERT- -ase, ALAM achieves a compres- sion rate of up to 16.4x with <0.5 accuracy degradation in the WNLI dataset. Improvements from GACT expand further when fine-tuning BERT-Base on more challenging datasets such as CoLA, closely matching the uncompressed baseline in terms of accuracy. In BERT-Large, ALAM achieves activation memory savings of up 22.5x with accuracy comparable to the baseline. This improve- ment stems from the enhanced compression capability of ALAM at low-sensitivity activations.
Text Classification
Fig. comppres the sensitivity and bit precision allocation obtained by GACT and ALAM for BERT- Base on the MPRC dataset when the target average precision is set to bit. GACT uniformly assigns 1-bit precision to all layers regardless of their sensitivity the lowest precision supported in GACT is 1 bit, resulting in training failure. In contrast, ALAM allocates bit precision to layers differently based on their sensitivity, successfully training the model with 84.8% accuracy. Furthermore, while
Table 2: Test accuracy, activation memory (Act mem) with compression rate, and memory for cal- culating sensitivity (sens mcm) in fine-tuning diverse language models on the datasets in the GLUE benchmark. The bolded text represents the test cases with < 0.5 score degradation compared to the non-compressed baseline, as well as the largest memory saving among the candidate algorithms with 0.5 score degradation.
Model
Dotaset
Menric
Reacline
2i
2-001
GACT
1.5
1:b0
2:N0
Bit
(Ours)
Fiit
Conr
87.1
66.3
$1.2
Anl
R6.7
K3.6
$4.
Disii
STS-B
Act mem
$34 MB
50 MB (13.9xc)
48 MB (17.4%)
Joil
64 MI (13.0>)
50 MB (16.7x)
36 MB (25.2x)
BFRT:
Sans mam
766 MR
7:5 MB
feil
1.2
1.28B
1.2K8
Base
Are
36,3
55.8
51.3
19.5
56.3
SKT
55,5
WNLI
MB
53 MB (33.2x))
48 MB
MB
MB
36MB
753
MR
Coet
Juni
Joo
SSB
54.5
CoLA
Act mem
1678
130 (12.9xx)
Jini
fool
135 MB (12.4x)
105 MD (16x)
72 MB (2333x)
Seus mem
1240 MB
fii
2.5 KB
2.5 KB
2,5
BERT-
Are
68.1
65.5
Joo
WK.5
67.8
66.3
Base
RTB
Act mem
1678 MD
130 MD (12.9x))
Fori
135 MD (12.4H)
105 MB (16.0x)
'22 MB (23.3x)
(110M)
Sens men
1240 MB
fun
2.3 KB
2.38
2.5
Ace
86,8
65.4
Jun
foll
86.9
85.8
84.3
MRN:
1678
(12.7)) ()
1o0 MR (16.3%)
MB (22.4x)
mem
MD
2.3
Ace
VX11
Jun
SST-2
Act mem
4434
360 (12.3xx)
junt
fool
368 MB (12.0%)
282 (15.7)
197 (22.5k)
Sess mon
3817 MR
4.6 KR
4.6.
4.6
Aze
92.4
91.5
30n
And
92.1
92.1
90.7
(345M)
BERT:
Large
MNLI
QNLI
Sens mom
Act mem
4434
4434 MB
36.3
360
35% MB (12.496)
3817 MR
85.2
funi
fun
Jun
Koll
366 MB (12.1x)
4.6 AB
86.4
(12.0s))
IPM MB (15.9))
MB (15.9))
195
195 MB (22.7x)
ST.
4R
mom
MB
4.6
Ase
JUN
VI.I
QQP
Act mem
4434
363 (12.2
funt
Jool
MIH (12.4%)
282 (15.7%)
197 (22.5x)
Sens mem
3817 MD
fani
fool
4.6 AB
KB
4.5

---

Published as a conference paper at ICLR 2024
50
Peciiton (bir)) (Bi3))
30
10
23
350
140
160
Layer Index
Layer Ind.x
GACT
ALAM (Ours)
Figure 6: Sensitivity and bit precision allocation of GACT and ALAM when target average precision bit in BERT-Base fine-tuning on the MRPC dataset.
GACT requires substantial amount of memory for calculating sensitivity (e.g., 3817 MB on BERT- Large), ALAM consumes only negligible amount of memory, less than 4.6 kB in all experiments.
For evaluations using LLMs. we fine-tuned LLaMA2-7B (Touvron al.
Large Language Models
2023b) on the Alpaca dataset (Toori et al. 2023). and
Reselin
LLaMA2-133 13B (Touvron et al. 2023b) and LLaMA-
ALAM (3-o1)
30B (Touvron et al., 2023a) on the LIMA dataset (Zoou et al. (023)). Then, we observed 5-shot accu- racy for MMLU (Hendrycks et al. 202T) and O-shot accuracy for common sense reasoning Bhakthhvat-
18-
ALAM (2 br)
1.6-
(4-a0)
Veliidation
1.4
salam et al. 2021 Bisk al. 2020 Zellers al.
2019 Win 2019). reading comprehension (Clark
1.0
al. 2019 and truthfulness tasks (Lin et al. 2021
0.8
In experiments, we applied parameter-efficien fine- tuning (PEFT) such LoRA (Hu al.) 2022) and
0
1000
2000
3000
4000
3000
6000
Training Steps
QLoRA (Dettmers et al., 2023). Table 3 shows ALAM achieves a similar accuracy the baseline at 1-bit precision and provides up to 10.6x compres-
Figure 7: Learning curves when fine-tuning LLaMA2-7B on the Alpaca dataset.
sion rate, while GACT diverges even at 3-bit preci- sion. Our ALAM (1 (1-bit) compresses activation by 10.3x in LLaMA2-13B and 9.9x in ILaMA-30B while closely matching the baseline accuracy for all tasks. ALAM also significantly reduces mem- ory for sensitivity calculations; however, the savings are not substantial because PEFT requires only a small number of parameter gradients. Fig. 7 displays learning curves obtained during the fine- tuning of LLaMA2-7B. demonstrating that our ALAM closely matches the baseline even at 1-bit precision while 4-bit GACT does not converge to the loss as rapidly as the baseline. Furthermore, the model fine-tuned with ALAM (1-bit) generates appropriate responses to the given prompt, as shown in Appendix
Table 3: Test accuracy, activation memory (Act mem) with compression rate, and memory for cal- culating sensitivity (sens mem) in fine-tuning LLMs. The bolded text represents the highest score and smallest activation memory while successfully training without divergence. The micro-batcl size is set to 8 for LLaMA2-7B and 2 for LLaMA2- 13B and LLaMA-30B.
Model
PEFT
Scheme
Precision
Act mem
Seas mem
MMLU
Arc-c
PIOA
Hellaswag
WinoCrrnnde
Boal()
TruthuulQA
Baseline
16-bit
21.1 GB
0
47.6
79.5
75.7
58.8
77.9
41.8
GACT
4-bit
(3.6x)
MB
LLaMA2
LoRA
4-bil
3-hit
4 GB (4.4%)
539 GB (3.8K)
0.4K6
MR
46.2
fuil
47.9
79.3
full
75.6
full
8933
full
77.5
ful
42.1
folt
ALAM
3-bit
2-hit
3.4 GB (6.2
0.488
45.9
47.1
79.7
75.6
75.4
68.8
77.3
42.3
I-bii
2.0 GB (10.6x)
0.4KB
45.7
47.0
79.3
75.6
69.1
77.6
41.1
L.IaMA2
ORA
39.1
-130
ALAM
F-bii
6.8 GB (0033
0.565
54.8
50.0
802
78.8
71.6
81.0
38,9
T.T.MA
82.1
ALAM
1-ii
G.6 (9.%)
0.76
56.7
52.7
803
813
42.
8

---

Published as a conference paper at ICLR 2024
Table 4: Comparisons with model compression techniques in fine-tuning LLMs using single NVIDIA A6000 GPU. The micro-bater size is set to 8.
Model
Triining Stratogy
Average accumey
Model state
Memory usage
Activation
Tocal
Full fine-tuning
Out of memory
52 CBB
21 GB
75 GB
LLaMA2-7B
LoRA
LoRA + (1-b6
13 GB
2 GB
GB
Full Enc funing
Out of snemory
104 GB
33 GB
LLaMA2-138
QLaRA
GB
40
(LORA ALAM
10
We also compared the training memory usage of the model compression techniques and ALAM in T The results shows that LoRA and QLoRA only reduce the model state memory which includes parameters, gradients, and optimizer states compared to baseline. ALAM, on the other hand, compresses activation memory, achieving 10.6x and 10.3 compression rates LLaMA2- 7B and 13B, respectively. By applying ALAM to LoRA and QLoRA, we can reduce total memory by 56% and 75%, while maintaining comparable accuracy for LLaMA2-7B and 13B, respectively.
Usability To demonstrate the utility of ALAM, we first observed how much our ALAM could increase the micro- batch size single NVIDIA RTX-3090 Ti GPU when fine-tuning LLaMA2-7B using LORA. Fig. g shows that ALAM allows for using the batch size of 16, im- proving the baseline and GACT by 8x and 2x, respec- tively. We also simulated how much memory could be reduced for fixed micro-batch sizes when fine-tuning LLaMA2-7B using LoRA, as shown in Table [ The size of activation memory increases with the miero-batch size and sequence length, whereas the model state mem- ory storing parameters, gradients, and optimizer state re- mains unchanged. Consequently, the memory savings of ALAM is more noticeable with larger micro-batch sizes, as demonstrated in Table 5 For instance, when micro- batch size is 8, ALAM reduces the total memory by 56% and 21% % compared to the baseline and GACT, respec-
Maximm Batch Size
16
10
8
2
Esseline
GACT (4-b2)
ALAM (1-B8)
Figure 8: The maximum micro-batch size of different training schemes for fine-tuning LLaMA2-7B.
tively, but the saving increases to 84% and 52% when the micro-batch size 64.
Table 5: Activation memory and total memory required for fine-tuning LLaMA2-7B using LoRA with different training schemes.
Scheme
Procisioon
Micro batch
Micro batch 32
Micno- batch
Reseline
16-bit
GR
GBB
84 GB
98
169 GR
182 GB
GACT
4-bit
6GB
G9
24 GB
57 GB
47 GB
GB
DISCUSSION
In this paper, we introduce ALAM, a new activation-compressse training framework for transformer models. Our ALAM enables deeper compression of activations with low sensitivity through Aver- age Quantization and drastically reduces the memory required for sensitivity calculations using the GradNorm Var algorithms. As result, ALAM reduced the activation memory of LLaMA-30B by 9.9x while outperforming GACT, the state-of-the-art ACT framework, with comparable accuracy. We hope that our ALAM will boost the training of LLMs, overcoming the challenges of limited GPU resources.
Limitations ACT frameworks require additional computation to compress and decompress activa- tions. This complexity contributed to a 11 1-23% increase in training time for both GACT and ALAM (see Appendix C). In next study, we will address these challenges for efficient training.
9

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
In Sections 3 and we elaborated in detail on the core components of ALAM, such as Average Quantization and GradNorm Var. Section J describes the operation of the ALAM system, while the experimental results in Section 6 can be replicated by referring to the hyperparameters and other experimental settings detailed Appendix D Additionally, we provide the reproducible code for ALAM at https//ggghbbb.comooooSSSnnynnn
ACKNOWLEDGMENT
This work supported by the National Research Foundation of Korea (Grant NRF- 2022RICICI066880 and the Institute of Information & communications Technology Planning & Evaluation (Grants 2021-0-01343-004 and HTT-2203-RS--022--005500
REFERENCES
Winogrande: An adversarial winograd schema challenge at scale. 2019.
Sumithra Bhakthavatsalam, Daniel Khashabi, Tushar Khot, Bhavana Dalvi Mishra, Kyle Richard- son, Ashish Sabharwal, Carissa Schoenick, Oyvind Tafjord, and Peter Clark. Think you have solved direct-answer question answering? try arc-da, the direct-answer A12 reasoning challenge. CoRR, abs/2102.03315, 2021.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. pp. 7432-7439, AAAI Press, 2020.
Yuxuan Cai, Yizhuang Zhou, Qi Han, Jianjian Sun, Xiangwen Kong, Jun Li, and Xiangyu Zhang. Reversible column networks. In International Conference on Learning Representations, 2023.
Ayan Chakrabarti and Benjamin Moseley. Backprop with approximate activations for memory- efficient network training. Advances in Neural Information Processing Systems, pp. 2426- 2435, 2019.
Jianfei Chen, Yu Gai, Zhewei Yao, Michael W. Mahoney, and Joseph E. Gonzalez. A statistical framework for low-bitwidth training of deep neural networks. In Advances in Neural Information Processing Systems, 2020.
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael W. Mahoney, and Joseph Gonzalez. Acton: Reducing training memory footprint via 2-bit activation compressed training. volume 139 of Proceedings of Machine Learning Research, pp. 1803-1813. PMLR, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, (0s./604.06117 74, 2016.
Zekai Chen, Mariann Micsinai Balan, and Kevin Brown. Language models are few-shot learners for prognostic prediction. CoRR, abs/2302. 12692, 2023.
Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACI, 2019.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficiett exact attention with io-awareness. In NeurlPS, 2022.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 1916)14 2012.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. CoRR, abs/2305. 14314, 2023.
10

---

Published as a conference paper at ICLR 2024
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre- pre-training of deep Association for Compu-
bidirectional transformers for language understanding. pp. tational Linguisties, 2019.
R. David Evans and Tor M. Aamodt. AC-GC: lossy activation compression with guaranteed conver- gence. In Advances in Neural Information Processing Systems, pp. 27434-27448, 2021.
Jianwei Feng and Dong Huang. Optimal gradient checkpoint search for arbitrary computation graphs. pp. 14433-14444 Computer Vision Foundation / IEEE, 2021.
Aidan N. Gomez, Mengye Ren, Raquel Urtasun, and Roger B. Grosse. The reversible residual network: Backpropagation without storing activations. In Advances in Neural Information Pro- cessing Systems, pp. 2214-2224, 2017.
Simon Haykin. Neural networks: a comprehensive foundation. founddaaatn Prentice Hall PTR, 1994.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the Interna- tional Conference on Learning Representations (ICLR), 2021.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth Interna- lional Conference on Learning Representations. OpenReview.net, 2022.
Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, and Joseph E. Gonzalez. Checkmate: Breaking the memory wall with optimal tensor rematerial- ization. CoRR, abs/1910.02655 2019.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, 0112000108361, 2020.
Henry J Kelley. Gradient theory of optimal flight paths. Ars Journal, 30(10):947-55,, 1960.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations. OpenReview.net, 2020.
Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. CoRR, abs/2205.05 .05..0599... 198, 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods, 2021.
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Jocy Gonzalez, Michael W. Mahoney, and Alvin Cheung. GACT: ac- tivation compressed training for generic network architectures. In International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 14132-14152. PMLR, 2022a.
Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. EXACT: scalable graph neural networks training via extreme activation compression. In The Tenth International Conference on Learning Representations. OpenReview.net, 2022b.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 1608 03983, 2016.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory F. Diamos, Erich Elsen, David García, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In International Conference on Learning Representations. OpenReview.net, 2018.
11

---

Published as a conference paper at ICLR 2024
Marius Mosbach, Maksym Andriushchenko, and Dietrich Klakow. On the stability of fine- tne-tunnngg BERT: misconceptions, explanations, and strong baselines. In International Conference on Learn- ing Representations. OpenReview.net, 2021.
OpenAL. GPT-4 technical report. CoRR, abs/2303.0877 2023.
Zizheng Pan, Peng Chen, Haoyu He, Jing Liu, Jianfei Cai, and Bohan Zhuang. Mesa: A memory- saving training framework for transformers. CoRR, abs/2111.11122 2021.
Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint 100.006609.0447,,.20 2016.
Michael E. Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Momentum residual neural networks. In International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 9276-9287, PMLR, 2021.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. ArXiv, abs/1910.01 108, 2019.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations, 2015.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu. Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zheng, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, and Bryan Catanzaro. Using deepspeed and megatron train megatron-turing NLG 530b, A large-scale generative language model. CoRR, abs/2201.11990. 2022.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruming. In Advances in Neural Information Processing Systems, 2022a.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In Advances in Neural Information Processing Systems, 2022b.
Xiao Sun, Naigang Wang, Chia-Yu Chen, Jiamin Ni, Ankur Agrawal, Xiaodong Cui, Swagath Venkataramani, Kaoutar Maghraoui, Vijayalakshmi Srinivasan, and Kailash Gopalakrishnan. Ultra-low precision b-bit training of deep neural networks. In Advances in Neural Information Processing Systems, 2020.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-followin llama model, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, abs/2302. 13971, 2023a. doi: 10,48500ARXXI.2202..555 13971.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu. Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CoRR, abs/2307. .92888. 2023b.
12

---

Published as a conference paper at ICLR 2024
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations. OpenReview.net, 2019.
Naigang Wang, Jungwook Choi, Daniel Brand, Chia-Yo Chen, and Kailash Gopalakrishnan. Train- ing deep neural networks with 8-bit floating point numbers. In Advances in Neural Information Processing Systems 31, pp. 7686-7695, 2018.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-a natural language processing, 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? In Association for Computational Linguistics, pp. 4791-4800. Association for Computational Linguistics, 2019.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. LIMA: less is more for alignment. CoRR, abs/2305. 1106, 2023.
13

---

Published as a conference paper at ICLR 2024
APPENDICES
A PROOF OF THEOREM
In this section, we first establish the relationship between gradient variance and gradient difference in activation compression training (ACT) depicted in Eq. [ We then derive the boundary function of gradient variance for fully-connected, convolution, batch normalization, and layer normalization layers. We confirm that all of these boundary functions can be formulated as Eq. 6 Finally, we prove Theorem
A.1 RELATIONSHIP BETWEEN GRADIENT VARIANCE AND GRADIEN DIFFERENCE IN ACT
Chen et (2021) proved that the gradient variance of ACT is unbiased follows:
E B| = E V E V
(7)
Assume two random vectors X (2,) INN and Y (Vn) E RN B Vn, then [||X || can be expressed as
|||
E
Σ n + ynn?
Σ
din + 230 + Un
E [|| + E 11 + 2
E| E [[ym]
E [ + E Y ||3]
(8)
By adopting this property, we can derive the gradient variance of ACT as below.
Var JoL
E
V 4|
1 II?
E Vo + AV [ oL 1
E
V B | V
= EE | + + - |E
= Var[VeL] + [[I
A.2 BOUNDARY FUNCTION OF GRADIENT VARIANCE
Evans Aamodi (022) suggested that compressing activations would increase the gradient vari- ance, and minimizing the gradient variance between actual and approximated activations can lead to convergence of a network trained using approximated activations. Therefore, they derived the boundary function of the gradient variance for widely used layers such as fully-connected, convo- lution, batch normalization, and layer normalization layers. Here we summarize their findings and extend them generalize the boundary function as Eq. 6
Fully-Connected Layer Assuming weight 0 (Onghs) E R2xxH input activation X
(Inh.) E N XHH, and gradient of output activation V = (Df/0yna, E N NxH), Evans &
Aamodt 2021 derive the boundary function of | V fl as
(10)
14

---

Published as a conference paper at ICLR 2024
Approximating activations does not modify the gradient of output in fully-connected layers, and hence we treat | as constant. Therefore, the boundary function can be expressed as D(AX) d || where d is constant.
Convolution Evans Aamodi (2021) also find the boundary function of convolutional layer with filter size R x S, stride T, imput activation X (Fncchw) RNNxcxHH*U and the gradient of the output activation Vy f = (Oyczhu E R NxCxxXX
RS T Ty
Av &
(11)
The gradient of the output in convolutional layer can be treated constant, similar to that of the fully-connected layer. Hence, we can simplify the boundary function to D(AX) = d A|?,, where d represents constant.
While approximating activations (X + + AX) does not directly impact 0f does not have any terms related to x,) in fully-connected and convolu- tional layers, it does modify error propagation in batch normalization (i.e.. 8a, zz)). Therefore, Evans Aamodt (202)) analyzed the gradient variance of error propagation, where the weight is (7c) E Re, the input activation is X (Inchw) E IR xxxCXxx and the gradient of the output activation Y /0ynchu E IR NxCxHx ). The boundary function is
Batch Normalization
error propagation (i.c.. 8as
Øf has terms related to
2 M ge Ae
|AVx/1 S
Σ n,C,k.ww
(12)
where sa is the the inverse of the standard deviation of the input activation across the batch dimension (oā + 10-11-15111 501/, M is N HW, and ge is the upper bound of the gradient variance by approximating 8f
activation (i.e., (A 0% ) + (oro ja < g.).
We simplify the inequality in Eq. further by utilizing the inequality below:
Σabi
a
by
Va,. bb: 0
(13)
By applying the inequality above to Eq.
12
we obtain
25 ge
|Av\/12 <
Σ
Au.
nokww
M
n.c.b,w
M
Σ
2s Nagee
2
2
Σ
2
nchw
M
1
2s 2
|
(14)
n,,,hh,w
In addition, we derive the boundary function of the gradient variance of parameters below.
|AV 2 Σ
Axnehw) }
yynekw
0s
Ar A
Oynchu
of
Au
Oynehw
11,
n.hi,
Σ
2 nchw
=
Σ
22
Vy
f TAXI
(15)
Oynohon
n.C.b.a
15

---

Published as a conference paper at ICLR 2024
We applied the inequality in Eq. 13 followed by the Cauchy- Cauchy-Schwarz inequality, and then applied the inequality in Eq. 13 again to find the boundary function. Finally, we can write the boundary functions of batch normalization in Eqs. 14 and 15 as d |X| where d is constant.
Layer Normalization Modifying activations affects the gradient variance of error propagation in layer normalization well (Evans & Aamodi, (021). Assuming weight of layer normalization (rt) E RC, input activation x= (Iine) E ,े*x and the gradient of the output activation Vy = (0f /dync) R NCC the boundary function is
2 Ar 8n
AVxfI>>
~
Σ T2C
nc
(16)
where Sn. denotes the the inverse of the standard deviation the input activation across the feature dimension (o 10-5///2 We additionally applied Eq. 2021 for simplicity as follows:
13 to Eq. derivee by Evans & Aamodt
A
AV
C
R2C
Inc
N
Σ S Σ 22)
||
(17)
n,C
We further derive the gradient variance of parameters, as shown below.
0f AInc)²
AV
Sp. Oyne
Of
Ae
Oynn
< C
V 1
(18)
The Cauchy-Schwarz inequality and Eq. 13 are applied find the boundary function of gradi- ent variance of parameters in layer normalization. Consequently, the boundary function of layer normalization in Eqs. and 18 can be represented as || where d constant.
In summary, the boundary function of gradient variance in general layer can be expressed as Eq, 6 We do not consider layers such as ReLU and MaxPool that only require a 1-bit mask for backward pass, as these layers do not require activation compression.
A.3 MINIMIZING GRADIENT VARIANCE
Consider activation X = 11 i Ng, < < G} where G and N, denote the number of groups and the number of elements in group g. respectively. The elements in each group approximated to arbitrary single value as X + AX {8(0), z(9), ..... ((0) S G}. In this case, we can prove Theorem usinng Eq. [6, as detailed below.
Theorem 1. The value of the boundary function of gradient difference is minimized when each group of activations is approximately represented using its group average value.
Proof. We define activation X [ < 5iS N S G} and compressed activation as X AX (f(0) ((()
_ S g S G). Then, Eq. 6 be expressed as
16

---

Published as a conference paper at ICLR 2024
D(A d|
G N, (5l0) I; )
Therefore, D(AX) is a convex function with respect to 8699) and we can minimize the value of ((A when the condition (D(AX)/05l) is satisfied. Then,
N
@D(AX)
Dā(a) (f(2) 2
d
05(0)
g=1 ==1
G
N₂
2d
Σ>>>>
(i(0) (1))
9=1 1=1
G
N
2d
Σ
(N (Nga((9)
r
(n) 0
i=
Consequently, obtain
E (0) x
((9)
(19)
Ng
which is the average of each group. In other words, we can minimize the boundary function by approximating the elements in each group their group average.
B RESPONSES GENERATED BY LLAMA2-7B FINE-TUNED USING ALAM
Prompt: Translate the sentence 'Deux ou trois choses que je sais d' d'elle' into English.
Baseline I know two three things about her.
ALAM (1-bit): Two or three things that know about her.
17

---

Published as a conference paper at ICLR 2024
Prompt: Recommend movie for me to watch during the weekend and explain the reason.
Baseline: Movie: Spiderman: No way home
Reason: If you are fan of Marvel universe, Spiderman: No way home is the movie that you should watch during the weekend. Spiderman: No way home tells the story of Peter's friend from high school and his classmate that get evicted from their homes, they steal something from Doctor Strange's transport vehicle, called the Sanctum Sanctorum. Doctor Strange the Pretorians to capture the thieves and return the stolen item. But the Pretorians overpower Doctor Strange and take him to the Void, where they torture him. The thieves are brought the Sanctum Sanctorum and Doetor Strange grants them their wish. For Peter, he puts Doctor Strange's spell on him so that he can retrieve everyone who is stuck in different dimensions.
ALAM (1-bit) There are a few suggestions that would recommend you to watch this weekend. I think you would enjoy both of them; "Spiderman" and "Avatar". Spiderman is considered to be one of the best Marvel films. It is about a teenager who is bitten by a radioactive spider and gains a few superhuman powers that have helped him save the world from a few villains. Avatar is one of the most successful films of the 21 century. The film is centred around human who infiltrates the world of the Na vi, the indigenous inhabitants of planet known as Pandora in the twenty-firs first century on Earth.
C TRAINING TIME OF ACTIVATION-COMPRESSER TRAINING
In this section, we analyze the time overhead of activation compressed training including ALAM. 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 amalysis. Our ALAM utilizes average quantization, which involves average calculation and quantization including searching for min and max values, multiplication with scaling factor, and stochastic rounding. These operations have time complexity of (n). where n represents the number of activations. Korthikanti et al. (0222) calculated the FLOPs and the number of activations of a transformer block with flash-attention Dao et al. (022)) as 24bsh² + 4hs and 34bsh, respectively, where 6, 8, and h denote the micro-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
6 below.
Table Time complexity of training and activation compression required by ALAM.
Training Obbhh*)
Activation compression
Time complexity
O(bsh)
Therefore, as the model size h increases, the training time overhead due ALAM becomes rel- atively smaller. To verify this, we trained models larger than 7B on the LIMA (IK) dataset for epochs, and the results are displayed Table E.2 below. Table shows that while there 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.
Table Time complexity of training and activation compression required by ALAM.
Model
Baseline
GACT (4-bit)
ALAM (1-bit)
LLaMA2-7B
30 min
36 min
37 min
LLaMA2-13B
51 min
58 min
60 min
LLaMA-300
115 min
127 min
128 min
18

---

Published as a conference paper at ICLR 2024
D EXPERIMENTAL DETAILS
1 We trained an MLP (Haykin 1994) 512-512-10 model on the MNIST dataset (Deng 2012) using various approximate activation methods such averaging and median. In this experiment, training was executed with a fixed learning rate for 30 epochs using the SGD optimizer (Ruder 2016). The batch size and learning rate were set 128 and 0.01, respectively. A single RTX- 3090 Ti GPU was used.
Training MLP to confirm Theorem
Training VGG-11 to confirm quality of activations We conducted experiments to verify the quality of activations compressed by Average Quantization in VGG-11 (Simonyan & Zisserman 2013) when trained on the CIFAR-100 dataset (Krizhevsky et al. 2009). The learning rate and batch size were set at 0.01 and 128, respectively, and the SGD optimizer (Ruder (016) lized. The training lasted for 100 epochs, and we scheduled the learning rate using cosine annealing Loshchilov & Hutter 2016). This experiment conducted using single RTX-3090 Ti GPU.
uti-
Fine-tuning BERT models on the datasets in the GLUE benchmark We evaluated the perfor- mance of ALAM by fine-tuning various BERT models on the datasets in the GLUE benchmark (Wang et al., 2019). We trained DistilBERT-Bass (Sanh et al., 2019) on the STS-B and WNLI datasets, using the distilbert-base-uncaan
as the pretrained model from HuggingFace (Wolf et al. 2020 We set the batch size to 16 and trained for 5 epochs with different learning rates (5 x 10-5 x 10 1). and we reported the highest accuracy. Furthermore, we trained the BERT-Base model Devlin al. 2019) on the CoLA, RTE, and MRPC datasets, utilizing the bert-base-cased from HuggingFace (Woff et al. 2020) as our pretrained model. With a batch size of 16, we trained for epochs varying learning rates (2 x 10-5, 3 x 10-5 4 x 10-5, and 5 10 10~5) and reported the highest accuracy. The BERT-Large model (Devlin et al. 2019) was trained on the SST-2, QNLI, MNLI, and QQP datasets using the bert-base-cased the pretrained model from HuggingFace (Wolf al.) 2020). We trained with a batch size of 16 for 10 epochs at a learning rate of 1 x 10 a setting proposed by Mosbach et al. (202I) for robust training. All experiments were conducted on a single RTX-3090 Ti GPU, and we reported the results as the average across 5 seeds.
Fine-tuning large language models on the Alpaca dataset To demonstrate the effectiveness of ALAM in training large language models, conducted the following experiments: We fine-tuned
LLaMA2-7B model (Tovrron et al. 2023b) on the Alpaca dataset (Taori et al., (023) for epoch
with a learning rate of 1 x 10 Instead of the full fine-tuning, we used LoRA (Hu et al. 2022),
which is a parameter-efficient fine-tuning (PEFT) technique. Similarly, we fine-tuned the I.LaMA2-
13B model (Touvron et al. 2023b) and I.LMMA33BB model (Touvron et al. (023a] the LIMA
dataset (Zhou et al. (023) for 5 epochs with a learning rate of 5 x 10 and 2 x 10 respectively.
In this experiment, we applied the QLoRA (Dettmers et al. 2023 as PEFT.
All experiments were carried out with a max sequence length of 512, utilizing a single NVIDIA A6000 GPU. One exception is the experiments in Section 6 where we specifically investigated how much we could increase the micro--btch size using a single NVIDIA RTX-3090 Ti GPU. The fine- tuned models were then evaluated obtain 5-shot accuracy for MMLJU (Hendrycks et al.) 2021) and 0-shot accuracy for common sense reasoning (Bhakthavatsalam et al., 2021 Bisk et al., 2020 Zellers et al. 2019 Win 2019), reading comprehension (Clark et al. 2019). and truthfulness tasks (in et al. 2021 using the Im-harness library"
19


Reviews:
[{"review_id": "uEFqy0Mfnv", "reviewer": "Reviewer_HytY", "summary": "This work focuses on the challenge of high GPU memory consumption for storing activation for deep neural network training, especially for large language models (LLMs). Existing activation-compressed training (ACT) approaches introduce significant performance drops when applied to LLMs. This paper proposes a new ACT framework, ALAM, that applies average quantization and a lightweight sensitivity calculation scheme to enable substantial memory savings for LLM training while maintaining training performance. The experiments show that the ALAM framework can support up to a 12.5x compression rate with 1-bit quantization without compromising accuracy.", "questions": "Please answer the questions in weaknesses section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The paper is well-written and easy to follow.\n+ The proposed average quantization in ACT is simple but effective.\n+ The proposed GradNormVar algorithm significantly reduces the memory requirement for storing the gradients during the sensitivity evaluation.\n+ The evaluation results cover both the classical model BERT and the latest model LLaMa in terms of large language model training/finetuning. The results of fine-tuning LLaMa-2 with 2-bit and even 1-bit ALAM are promising.", "weaknesses": "- It would be better to add QLoRA as a baseline for evaluation.\n- It is unclear how the proposed ALAM performs on larger models, such as LLaMa-13B, LLaMa-30B, etc. It would be better to show the results on these larger models.", "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": "", "extra_scores": {}}, {"review_id": "3a8XwFq4b8", "reviewer": "Reviewer_Eyz9", "summary": "This paper introduces ALAM, an ACT framework that applies average quantization and a novel sensitivity calculation method to reduce memory usage while preserving the performance of LLMs. By compressing activations to their group averages, ALAM minimizes the impact on gradient variance, enabling deep compression with an effective precision of less than 1 bit. ALAM’s sensitivity calculation is based on the L2 norm of parameter gradients, which is a memory-efficient approach. The paper's experiments show that ALAM can achieve up to 12.5× activation memory compression in LLMs without sacrificing accuracy. This method represents a significant advancement in making the training of large neural networks more memory-efficient.", "questions": "1. It would be better to provide the evaluation result with other compression techniques.\n2. In Table 1, there is an accuracy drop from 98.72 to 96.09 in a small model. However, in the evaluation result, the compressed model retains the accuracy performance. Can you identify the reason for this? This is because the large model has more redundancy or other reasons?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The quantization method based on group average value is very easy to understand and very practical, given the experimental demonstration and equation proof in the paper.\n2. The training overhead is acceptable, since the efficient sensitivity calculation method is proposed, and the training time is comparable with the baseline.\n3. The evaluation result is very promising and the improvement compared with GACT (SOTA work based on activation compression) is very significant.", "weaknesses": "1. 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. \n2. How to choose the parameter group is unclear. Is this parameter related to the sensitivity? What if the adjacent activations have little similarity?", "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": "", "extra_scores": {}}, {"review_id": "Jle5RDYQFg", "reviewer": "Reviewer_UZqy", "summary": "This paper proposes a method for reducing the memory usage of activations during the training of large-scale models. The techniques extend well-established methods like GACT, further improving the activation compression rate by introducing group-wise quantization. The methods are evaluated on popular transformer models such as BERT and LLaMa, demonstrating enhanced compression rate while preserving baseline accuracy.", "questions": "Could the authors explain the overhead involved in doing the group-wise quantization?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1). The paper is well-written. \n\n2). The compression rates have been improved compared with previous methods.", "weaknesses": "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 quation 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.", "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": "", "extra_scores": {}}, {"review_id": "Yv1Fiu3iKq", "reviewer": "Reviewer_7RJR", "summary": "This paper proposes a novel ACT framework--ALAM to quantize activations for memory-efficient training of transformer models. ALAM utilizes average quantization and a lightweight sensitivity calculation scheme. The experiments show that ALAM significantly reduces activation memory.", "questions": "Please refer to the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe proposed ALAM provides high-quality deeply compressed activations with precision of less than 1 bit, which enables further compression without sacrificing accuracy.\n2.\tThe proposed ALAM is demonstrated to successfully compress activations in various transformer models.", "weaknesses": "The complexity of ALAM is not analyzed. And the paper only compares training time for LLaMA-3B.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)