Title: Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model
Venue: NeurIPS 2023
Accepted: yes
Keywords: memory-efficient tuning

language model

transformers
Abstract: As the model size grows rapidly, fine-tuning the large pre-trained language model has become increasingly difficult due to its extensive memory usage. 
Previous works usually focus on reducing the number of trainable parameters in the network. 
While the model parameters do contribute to memory usage, the primary memory bottleneck during training arises from storing feature maps, also known as activations, as they are crucial for gradient calculation. 
Notably, machine learning models are typically trained using stochastic gradient descent.
We argue that in stochastic optimization, models can handle noisy gradients as long as the gradient estimator is unbiased with reasonable variance.
Following this motivation, we propose a new family of unbiased estimators called \sas, for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient.
Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
By replacing the linear operation with our approximated one in transformers, we can achieve up to 2.7X peak memory reduction with almost no accuracy drop and enables up to $6.4\times$ larger batch size.
Under the same hardware, \sas enables better down-streaming task performance by applying larger models and/or faster training speed with larger batch sizes.
The code is available at https://anonymous.4open.science/r/WTACRS-A5C5/.

Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model

Abstract

As the model size grows rapidly, fine-tuning the large pre-trained language model has become increasingly difficult due to its extensive memory usage. Previous works usually focus on reducing the number of trainable parameters in the network. While the model parameters do contribute to memory usage, the primary memory bottleneck during training arises from storing feature maps, also known as activations, as they are crucial for gradient calculation. Notably, machine learning models are typically trained using stochastic gradient descent. We argue that in stochastic optimization, models can handle noisy gradients as long as the gradient estimator is unbiased with reasonable variance. Following this motivation, we propose a new family of unbiased estimators called WTA-CRS , for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones. By replacing the linear operation with our approximated one in transformers, we can achieve up to 2.7 × peak memory reduction with almost no accuracy drop and enables up to 6 . 4 × larger batch size. Under the same hardware, WTA-CRS enables better down-streaming task performance by applying larger models and/or faster training speed with larger batch sizes. The code is available at https://github.com/zirui-ray-liu/WTACRS/ .

1 Introduction

Pre-trained language models (LMs) with transformer architecture have achieved remarkable success in numerous natural language processing (NLP) tasks [1, 2, 3, 4, 5]. Specifically, these models are trained on vast text corpora to acquire general-purpose representations, which are then adapted to a specific task by fine-tuning on task-specific data. In recent studies, it has been convincingly demonstrated that significantly increasing the number of parameters in pre-trained LMs leads to remarkable improvements in performance [6, 7]. As a result, there is now an urgent necessity to effectively adapt these models, equipped with billion-scale parameters, to a wide range of tasks.

However, a significant disparity exists between the memory requirements of pre-trained LMs and the capacity of current hardware, particularly GPUs. For example, even a GPU with 24GB memory cannot accommodate the fine-tuning process of the T5-3B model [3] with batch size one, which boasts three billion parameters. Without additional techniques, attempting to fine-tune billion-scale LMs on a single GPU is impossible. Although model-parallel fine-tuning is feasible, the majority of the time, we cannot bear the expense of acquiring multiple GPUs or the communication overhead

involved. To ensure the smooth deployment of language models during the fine-tuning process, it is crucial to adapt them for operation on a single GPU.

To address this issue, several parameter-efficient tuning methods are proposed [8, 9, 10, 11, 12, 13, 14]. Specifically, adapters [14, 13] insert a small module into the transformer blocks and only update it while keeping other parameters fixed. Similarly, prompt tuning [8] introduces a small vector that is concatenated with the input embeddings and updated during the tuning process. LoRA [12] injects trainable rank decomposition matrices into the transformer block, updating them while freezing the others. Parameter-efficient tuning methods mainly reduce the memory taken by the optimizer states [15, 12]. Although the optimizer states contribute to the memory footprint, storing activations (or feature maps) is the main memory bottleneck during training (often > 70% ) [16, 17, 18, 19]. Thus, parameter-efficient methods often do not reduce memory usage by much [9, 19].

In parallel, we can reduce the main memory bottleneck by reducing the activation storage in finetuning. Since transformer-based models are mainly built based on the linear layer, a less-explored direction is to replace the expensive matrix multiplication operation with its memory-efficient estimations using column-row sampling (CRS) [21, 22]. The key idea of CRS is to sub-sample tensors onto low-dimensional spaces and perform the original operations here. Specifically, for the linear operation between two matrices A ∈ R n × m and B ∈ R m × q (in the context of machine learning, A is often activations), we first sample k ( k < m ) column-row pairs according to a pre-defined distribution. Then we obtain A ′ ∈ R n × k and B ′ ∈ R k × q ( k < m ) by picking k columns of A and the corresponding rows of B according to the sampled column-row pairs [22]. Finally, we estimate AB ≈ A ′ B ′ . In this way, we only need to store the sub-matrix A ′ and B ′ in GPU memory to perform the computation. Moreover, transformer-based models training/tuning are performed with the first-order stochastic optimizer, e.g., Adam [15]. In stochastic optimization, models can work with noisy gradients, as long as the gradient estimator is unbiased and has a reasonable variance. In view of such, we ask: why spend resources on obtaining exact gradients when we are using stochastic optimization? Motivated by this, we focus on obtaining unbiased gradients cheaply with approximated matrix multiplication.

The approximation method reduces the memory usage at the cost of giving outputs with variance. Thus there naturally exists an accuracy-memory trade-off. The main challenge is how to integrate the approximated matrix multiplication into transformer with minimal gradient variance. In this paper, we propose a new family of unbiased estimator for matrix multiplication with reduced variance, dubbed Winner-TakeAll Column-Row Sampling ( WTA-CRS ). Compared to CRS, WTA-CRS reduces the variance of an estimator by focusing more on high-probability regions of the sampling distribution. Moreover, WTA-CRS can serve as a drop-in replacement for the linear operation in transformers, providing an unbiased weight gradient with reduced memory usage. As shown in Figure 1, our method achieves better accuracy-memory trade-off than

20

30

40

Training Peak Memory (GB)

75

80

85

90

Average Validation Result (%)

Full (T5-Large)

LoRA (T5-Large)

LST (T5-Large)

LoRA+WTA-CRS@0.3 (T5-Large)

LoRA+WTA-CRS@0.3 (T5-3B)

LoRA+WTA-CRS@0.1 (T5-Large)

LoRA+WTA-CRS@0.1 (T5-3B)

state-of-the-art memory-efficient tuning methods, e.g., LST [9] and LoRA [12]. Moreover, since WTA-CRS executed at the operation level, it is orthogonal to most of the existing parameter-efficient tuning methods. Our contributions are highlighted as follows:

We design a new family of unbiased estimator for matrix multiplication with reduced variance. We theoretically and experimentally verify that it has smaller variance than the established one under the context of tuning transformer.

By replacing the linear operation with WTA-CRS in transformers, we can achieve up to 2.7 × peak memory reduction with almost no accuracy drop, and enables up to 6 . 4 × larger batch size. For example, to fully tune T5-3B, it requires 66.4GB of memory for in a mini-batch size of 32 , which relies on A100 GPUs with a capacity of 80GB. With WTA-CRS , it only requires 21.6GB of memory for fine-tuning, which can run on a GPU with 24GB memory, e.g. RTX3090Ti.

We implement WTA-CRS as a ready-to-use extension for Pytorch with an easy-to-use API that can also be combined with other memory-saving techniques.

2 Background and Preliminary

In this section, we first analyze the memory usage of transformers. Then we introduce the background on the approximated matrix multiplication.

2.1 The Memory Usage of Transformers

In each training step of backpropagation, it has exactly two phases, i.e., one forward phase and one backward phase. Transformer-based models are mainly built based on the linear operation, which can be written as:

where

H

(S=128)

(S=256)

T5-Large

(S=128)

T5-Large

(S=256)

0

10

20

30

40

50

60

Memory (GB)

77.0%

88.3%

73.4%

86.1%

Memory Usage Breakdown

Model

Optimizer

Activation

Other

GEMM

and

Z

(

·

,

·

)

is the General Matrix Multiplication operation, are the activation (or input feature maps) and output

feature maps, respectively. W is the weight of the linear layer. ∇ H , ∇ W , and ∇ Z are the gradient of H , W , and Z , respectively. From Equation (1c), activations H are used in the backward phase. In commonly used deep learning framework [24, 25], it requires storing H in GPU memory during the forward pass, for calculating the weight gradient ∇ W in the backward pass.

Previous works show that although the model parameters contribute to the memory footprint, activations (e.g., storing H ) are the main memory bottleneck during training [16, 17, 18, 19]. To get a sense of the scale, we show in Figure 2 that for popular transformer models like T5, activations may take roughly 73 ∼ 88% of the total memory, depending on the batch size B and sequential length S .

2.2 Approximated GEMM With Sampling

Let X ∈ R n × m , Y ∈ R m × q be two matrices. The goal is to efficiently estimate the matrix production XY . Singular Value Decomposition (SVD) outputs provably optimal low-rank estimation of XY [21]. However, SVD is almost as expensive as matrix production itself. Instead, the sampling algorithm is proposed to approximate the matrix product XY by sampling k columns of X and corresponding rows of Y to form smaller matrices, which are then multiplied as usual [22, 26]:

where X : ,i ∈ R n × 1 and Y i, : ∈ R 1 × q are the i th column and row of X and Y , respectively. In this paper, we call ( X : ,i , Y i, : ) the i th column-row pair. k is the number of sampled pairs ( 1 ≤ k ≤ m ). P = { p i } i m =1 is a probability distribution over the column-row pairs. i t ∈ { 1 , · · · m } is the index of the sampled column-row pair at the t th trial. s t is the scale factor. X ′ ∈ R n × k and Y ′ ∈ R k × q are the normalized sub-matrices sliced according to the sampled column-row pairs.

Existing work [22] shows X ′ Y ′ is an unbiased estimation of XY , i.e., E [ X ′ Y ′ ] = XY . Furthermore, the approximation error E [ || XY -X ′ Y ′ || F ] is minimized when the probabilities { p i } i m =1 are proportional to the product of the column-row Euclidean norms [22] (Proof in Appendix C):

As we analyzed in Section 2.1, storing the activation H is the major memory bottleneck. If we can replace GEMM ( H ⊤ , ∇ Z ) in Equation (1c) with H ′⊤ ∇ Z ′ following the paradigm of Equation (2), then we only need H ′ instead of H in GPU memory to compute the gradient, which significantly decreases the memory usage of activations. This estimation linearly reduces the memory complexity from O ( nm ) to O ( nk ) . Also, the total number of floating point operations (FLOPs) is reduced as

well since the computation is executed on two smaller matrices. For the ease of illustration, in this paper we call the distribution in Equation (3) the column-row index distribution . In the next section, we explore how to reduce memory usage via sampling-based matrix multiplication.

3 Methodology

In recent years, we have observed that deep neural network training can be performed almost entirely with first-order stochastic optimization [15]. Thus intuitively, in stochastic optimization we can reduce the resources spent on obtaining gradients, as long as the estimated gradient is unbiased with reasonable variance [27, 28, 29]. Following this motivation, we first design a new unbiased estimator for matrix multiplication with reduced variance compared to the one in Equation (2) (Section 3.1 ). Then we introduce how to replace the GEMM in Transformer with its approximated version to reduce the memory usage (Section 3.2).

3.1 Winner-Take-All Column-Row Sampling: A New Unbiased Estimator for GEMM

In this section, we mathematically design a new unbiased estimator for GEMM with reduced variance called WTA-CRS (Winner-Take-All Column-Row Sampling). Following the notation in Section 2.2, let X ∈ R n × m , Y ∈ R m × q be two matrices. P = { p i } i m =1 is the column-row index distribution in Equation (3) 2 . We first define the variable f ( i ) as f ( i ) = X : i Y i : p i .

f ( i ) is an unbiased estimation for the matrix production between X and Y . To see this,

We note that the prior approximated matrix multiplication in Equation (2) is the direct extension of f ( i ) by taking the average of { f ( i t ) } k t =1 among k independent random trials to reduce the variance. Here we explore an alternative approach to reduce the variance of f ( i ) beyond simple averaging. Our core idea is to partition the column-row index distribution P = { p i } i m =1 into two complementary regions based on the probability mass: a high-probability region P C and a low-probability region P D\C , where D = { 1 , · · · , m } is the whole set and C is the set of the column-row index with the largest probability. Let C be the set of column-row pair indices associated with |C| largest p i . We define WTA-CRS estimator for XY as follows:

We note that the random variable in Equation (4) is the column-row pair index j , and is only sampled from D\C . The estimator defined in Equation (4) contains two parts. The first part ∑ c ∈C f ( c ) p c has no relationship with the random variable j and is summed deterministically. The second part f ( j ) is sampled stocastically, but scaled by the factor (1 -∑ c ∈C p c ) . When P = { p i } i m =1 is concentrated on a small number of atoms, the scaling factor (1 -∑ c ∈C p c ) for the stochastic term should be small. Therefore, we intuitively expect the estimator to have a small variance in this case due to a small scaling factor. In this way, we reduce the variance of an estimator by focusing more on high-probability regions of the distribution (winner-take-all). Below we formalize this intuition by showing the statistical property of our estimator regarding the bias and variance, respectively.

Theorem 1 (Proof in Appendix C.2) . The estimator defined in Equation (4) is an unbiased estimator for matrix production XY , i.e, E j ∼P D\C [ ∑ c ∈C f ( c ) p c +(1 -∑ c ∈C p c ) f ( j )] = XY .

Theorem 1 states that our proposed estimator in Equation (4) is unbiased. Below we compare our proposed estimator to the CRS estimator in Equation (2) in terms of the variance. Suppose we have the budget of only utilizing k column-row pairs for approximating the matrix production. From the implementation perspective, the estimator defined in Equation (2) estimates GEMM ( X , Y ) as:

Our estimator defined in Equation (4) splits the budget k into two parts. Namely, the first part explicitly sums the expectation terms for the largest probability group C ( |C| < k ), while stochastically average k -|C| samples drawn from D\C to estimate the remaining terms, up to scale:

Theorem 2 (Proof in Appendix C.3) . Suppose the total budget of column-row pairs is k . If C satisfies

then we have Var[ˆ g ( X , Y )] < Var[ g ( X , Y )] . Moreover, Var[ˆ g ( X , Y )] is minimized when |C| = min |C|∈{ 0 , ··· ,k } 1 -∑ c ∈C p c k -|C| .

Both the left- and right-hand sides of Equation (7) depend on the size of the highest probability group |C| , which controls the number of high probability column-row pairs that are directly added without sampling. Below we experimentally investigate whether Equation (7) holds under the context of fine-tuning the transformer-based model with varying |C| .

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

0.8

Probability Mass

c

pc

|

|

)

2-th layer

4-th layer

6-th layer

|

|

k

Ratio

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

|

|

)

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

0.8

|

|

)

Experimental analysis. As shown in Figure 3, we visualize the two terms in Equation (3) for the column-row index distribution of query, key, and value projection in the self-attention module, respectively [1]. Specifically, we fix the total column-row pair budget k = 0 . 3 |D| and change the size of the highest probability group |C| from 0 to k . We conclude that Equation (7) holds for most of the layers when fine-tuning transformers. Thus, we expect our WTA-CRS has better performance than CRS for adapting transformer-based models, which is later experimentally verified in Section 5.

3.2 Compress GEMM in Transformers with WTA-CRS

Previous work has shown that unbiasedness of the estimated gradient is crucial for the proper convergence of stochastic gradient descent [27, 28, 16, 30]. As shown in Section 2.1, we have three GEMM in the linear layer. Below we investigate how to replace GEMM with its approximated version in a way that the estimated gradient is unbiased.

̸

GEMM

Approx-

GEMM

GEMM

Lived in GPU

memory

Equation (5)

Unbiasedness. Previous work has shown that to ensure the unbiasedness of the gradient, the approximation can only be applied during the backward pass [16, 31, 21]. The rationale behind this conclusion is that we have E [ f ( x )] = f ( E [ x ]) for any non-linear function f ( · ) , e.g.,

̸

E [ x 2 ] = E 2 [ x ] . Thus if we replace the forward GEMM in Equation (1a), even when the approximation method gives an unbiased estimation, i.e., E [ˆ g ( H , W )] = HW = Z , the output activations (e.g., GeLU ( Z ) ) are still biased since the activation function is non-linear, namely,

̸

Linear-Q

Linear-K

Linear-V

TensorMul-1

TensorMul-2

GELU

Dropout

Dropout

Linear-O

LayerNorm

Linear-U

Linear-D

Dropout

LayerNorm

Output of

last block

Q

K

V

Q, -K, and -V

Reshape

[

B

×

S

×

D

model

]

[

B

×

N

head

×

S

×

D

head

]

[

B

×

N

head

×

D

head

×

S

]

[

B

×

S

×

D

model

]

[

B

×

N

head

×

S

×

S

]

[

B

×

N

head

×

S

×

S

]

[

B

×

S

×

D

model

]

[

B

×

S

×

D

model

]

[

B

×

S

×

4

D

head

]

[

B

×

S

×

4

D

model

]

[

B

×

S

×

D

model

]

[

B

×

N

head

×

S

×

D

head

]

[

B

×

N

head

×

S

×

D

head

]

To ensure the unbiasness of gradient and reduce the memory usage of storing H , as shown in the example of Figure 4, we only replace GEMM in the backward pass with its approximation (e.g., Equation (1c)), while leaving the forward one unchanged (e.g., Equation (1a)). We show in Appendix B that the estimated weight gradient is unbiased in this case.

System Implementation. Here we present how we implement WTA-CRS in Equation (6) in practice. For the linear layer, as we analyzed, we only replace GEMM in Equation (1c) with its approximated version. In this case, X and Y in Equation (6) are activation H ⊤ and output gradient ∇ Z , respectively. Given the total column-row pair budget k , the first step is to build the deterministic index set C , where each element is summed explicitly without sampling. Note that C is a set of indices with the highest probabilities in Equation (3). Thus, to build C , we only need to determine its size, denoted as |C| , which minimizes the variance of the estimator. As Theorem 2 suggested, we set |C| = min |C|∈{ 0 , ··· ,k } 1 -∑ c ∈C p c k -|C| . The second step is to sample k -|C| column-row indices from the remaining distribution P D\C to obtain the set C stoc, where |C stoc | = k -|C| . The third step is to build sub-sampled H ′ with only rows from C ∪ C stoc. Note that for rows in H ′ from C stoc, we need to normalize it by 1 -∑ c ∈C p c k -|C| according to Equation (6). We illustrate the above process in Figure 4.

Scope. Here we show which operation can be replaced with its approximation version. As shown in Figure 5, the transformer is mainly consisted of linear layer, TensorMul, and other operations (e.g., GeLU, Dropout, LayerNorm). TensorMul in Figure 5 refers to the multiplication between two four-dimensional tensors. Our WTA-CRS can be applied to Linear-Q, -K, -V, -O, -U, -D, TensorMul-1, and TensorMul-2 (in green). The activations of Dropout and GELU operations (in blue) can be losslessly compressed. The Softmax and LayerNorm operators (in gray) remain unchanged.

4 Related Work and Discussion

Due to the page limit, we discuss the related work on approximated matrix multiplication and aproximation in LLM inference. Other related topics, e.g., parameter-efficient fine-tuning, activation quantization, and gradient checkpointing, can be found in Appendix A. We also discuss the limitation and potential negative social impact in Appendix A.

Approximated Matrix Multiplication. In the context of neural networks, approximated matrix multiplication methods can be broadly categorized into two main groups: (1) Butterfly-based methods [33, 34] replace dense weight matrices with butterfly matrices. We note that they focus on the weight matrix and are orthogonal to our research, as we concentrate on sub-sampling the activation matrix. (2) Column-row sampling (CRS) methods[22, 21, 31] select important rows and columns from the input matrices and perform the multiplication on the sampled matrix. Our work is closely aligned

with this second research line. [21, 31] share similarities with our research in terms of utilizing CRS for approximating matrix multiplication within neural networks. The main distinction lies in how to select the column-row pairs. Specifically, [21] deterministically selects column-row pairs without scaling, whereas our estimator divides the column-row pairs into a deterministic component and a stochastic component. As we analyzed, selecting column-row pairs deterministically is biased. Later we show that this approach may cause a significant accuracy drop ( 'Deterministic' in Figure 12 ).

Approximation in Inference There are two approximation techniques to reduce inference latency. (1) Sparse Modeling, which only involve a subset of weights during computation to reduce both computational and memory I/O demands [35, 36, 37]. (2) Quantization, which compresses the trained weight into lower numerical precision [38, 39, 40]. All these approximation techniques trade-off model quality in return for improved efficiency. Besides approximation technique itself, researchers also propose methods to recover the accuracy drop of compressed models [41, 42], e.g., by prompting compressed models. These methods can greatly improve the accuracy-efficiency trade-off of LLMs.

5 Experiments

In this section, we design experiments to answer the following research questions: RQ1: How effective is WTA-CRS in terms of accuracy with reduced memory usage? RQ2: How sensitive is WTA-CRS affected by its key hyper-parameters? RQ3: WTA-CRS contains two parts, i.e., the deterministic summation part and the statistical sampling part. Are they both necessary? RQ4: How is the fine-tuning speed affected by WTA-CRS ?

5.1 Experiment Setup

Datasets and Evaluation Protocol. Following most of the previous work, we adopt GLUE benchmark [43] to evaluate the effectiveness of different methods, including the CoLA, SST-2, MRPC, QQP, MNLI, QNLI, RTE, and STS-B datasets. To evaluate the memory usage, we report the peak GPU memory usage and compression rate during the fine-tuning process with Huggingface API [44].

Compared Methods and Adopted Models. We consider three methods to compare in this paper: Full fine-tuning (Full), LoRA [12], and Ladder Side-tuning (LST) [9]. Specifically, Full tunes all of the parameters in the model to provide an upper bound of accuracy; LoRA inserts trainable low-rank matrices into the model to parameterize the weights' changes; LST injects a trainable ladder side structure. Since WTA-CRS essentially replace the linear operation with approximated one, we emphasize that our WTA-CRS is compatible with all these three compared methods, i.e., they can be combined together towards smaller memory usage. For the backbone model, we follow the previous work [9, 14, 12] to adopt the OPT [45], Bert-Base [2], Bert-Large, T5-Base, T5-Large, and T5-3B [46] for evaluating the effectiveness of different methods.

Hyperparameter Settings. For WTA-CRS , it only has one hyperparameter k , which controls the column-row pair budget. We assign the same k to all replaceable linear operations in the model. We consider the normalized column-row pair budget k/ |D| ∈ { 0 . 3 , 0 . 1 } , which are denoted as WTA-CRS @0.3 and WTA-CRS @0.1, respectively. We also consider the combination of WTA-CRS and LoRA to further reduce the memory cost of both optimizer and activations. The detailed hyperparameters are given in Appendix F. All reported results are averaged over three random trials.

5.2 Accuracy versus Memory Usage (RQ1)

To answer RQ1 , we first analyze the trade-off between the model performance and memory saving. The evaluation results and peak memory usage are given in Tables 1 and 2, respectively. We observe:

❶ WTA-CRS achieves a superior trade-off between accuracy and memory usage compared to baselines. Specifically, WTA-CRS has negligible accuracy drop, while the peak memory usage is reduced by 2 . 1 × ∼ 2 . 7 × (when combined with LoRA).

As we analyzed, LoRA mainly reduces the memory of optimizer states. Thus, although it has negligible accuracy drop, it can only achieve ∼ 1 . 3 × peak memory saving. LST can reduce memory usage up to 3 × , but its accuracy drop is much larger than LoRA and WTA-CRS . Since WTA-CRS executes at the operation level and focuses on the activations, we further combine LoRA with WTA-CRS to reduce memory usage more aggressively. When combing with LoRA, WTA-CRS achieves 2 . 7 × memory usage saving with almost no accuracy drop. To fully tune T5-3B, it requires 66.4GB of memory for in a mini-batch size of 32 , which relies on A100 GPUs with a capacity of 80GB. On the other hand, LoRA+ WTA-CRS only requires 21.6GB of memory for finetuning with a mini-batch size of 32 , which can run on a GPU with 24GB memory, e.g. RTX3090Ti or A5000. We have experimentally confirmed this

100

200

Batch-size

40

60

80

Full

LoRA

LoRA+WTA-CRS@0.3

LoRA+WTA-CRS@0.1

conclusion. Under the same hardware, WTA-CRS enables the tuning of larger models, resulting in improved down-streaming task performance. Thus as shown in Figure 1, ❷ under the similar memory budget, WTA-CRS outperforms the other methods in terms of the accuracy.

Also, according to Figure 6, for T5-3B, LoRA itself can enable 1 . 9 × larger batch size. ❸ When combined with LoRA, WTA-CRS enables 4 . 8 × ( k =0 . 3 |D| ) to 6 . 4 × ( k =0 . 1 |D| ) larger batch-size.

Influence of Row-column Pairs Budget (RQ2). As we analyzed in Section 3.2, WTA-CRS only have one hyperparameter, i.e., the total column-row pair budgets k . We conduct the ablation study with different budget k in Figure 7. We observe that ❹ It has almost no accuracy drop when k = 0 . 3 |D| . And the accuracy drop is about 1% when k = 0 . 1 |D| . Notably, for T5-3B, the accuracy drop is only 0 . 4% when k = 0 . 1 |D| , which is much smaller than T5-Base and T5-Large. This suggests that larger models are more compressible because they have more redundant activations, which is consistent with previous observations [47].

5.3 Ablation Study (RQ3 and RQ4)

To answer RQ3 , WTA-CRS is compared with two compositional methods to demonstrate its superiority. Namely, (1) the Deterministic method selects row-column pairs with top k probability of Equation (3). We note that this is the estimator proposed in [21]. (2) The CRS method follows Equation (3) to

w/ LoRA

80.0

82.5

85.0

Average Validation Result (%)

85.7

85.4

84.5

85.7

85.4

84.7

(a) T5-Base

w/ LoRA

84

86

88

Average Validation Result (%)

87.7

87.4

86.5

87.7

87.4

87.1

(b) T5-Large

w/ LoRA

86

88

90

Average Validation Result (%)

89.5

89.8

89.1

Full

k

/|

| = 0.3

k

/|

| = 0.1

(c) T5-3B

sample the row-column pairs. All methods are deployed to GEMM in the backward pass, while leaving the forward one unchanged. The experiments are conducted on the training of T5-base language model on the SST2, MNLI, and QQP datasets; The column-row pair budget takes k/ |D| = 0 . 1 for all methods. The validation accuracy versus training epoch is given in Figure 12. We observe:

❺ WTA-CRS outperforms all compared methods, especially as the training epoch grows. The deterministic selection of top k column-row pairs suffers from accumulation of bias error that ultimately results in a failure of convergence. For CRS, it also enables the unbiased weight gradient. However, as we theoretically and experimentally analyzed in Theorem 7 and Figure 3, it is worse than WTA-CRS due to larger variance. In summary, both the deterministic and stochastic parts contribute to the effectiveness of WTA-CRS , which is consistent with our theoretical analysis.

The speed of WTA-CRS (RQ4). The configuration of computational infrastructure is given in Appendix F.1. We note that WTA-CRS does not add any extra parameters to the model. Thus, WTA-CRS only affect the fine-tuning speed, without affecting the inference speed. Below we analyze how the fine-tuning speed affected by WTA-CRS . As we analyzed in Appendix A limitation, the current implementation is not heavily optimized and thus the execution time of WTA-CRS is still slower than the original linear operation (details are shown in Appendix E.2). However, under the same

200

400

600

Batch-size

45

50

55

Training Throughput

Full

WTA-CRS@0.3

WTA-CRS@0.1

(a) T5-Large

50

100

150

200

Batch-size

12

14

16

Training Throughput

Full

WTA-CRS@0.3

WTA-CRS@0.1

(b) T5-3B

hardware, a reduction in activation memory enables the use of larger batch sizes, thereby improving training speed due to increased GPU utilization [48, 9]. As we analyzed in Figure 6, WTA-CRS can enlarge the available batch size by up to 4 . 8 × larger. This enhancement is expected to result in a acceleration of the training speed. To illustrate, Figure 8 presents a visualization of batch size against training throughput (sentences per second) for both T5-Large and T5-3B models. We observe that ❻ WTA-CRS enables faster training speed under the same hardware. Specifically, on the T5-Large model, WTA-CRS @0.1 shows 1 . 08 × higher training throughput; and on the T5-3B model, WTA-CRS @0.3 and WTA-CRS @0.1 achieve 1 . 14 × and 1 . 21 × higher training throughput, respectively.

7 Conclusion

In this paper, we propose WTA-CRS , a new unbiased estimator for matrix production with reduced variance. We theoretically and experimentally show when and why the estimator is better than the traditional unbiased estimator in terms of the variance. In the context of adapting transformers, it almost has no accuracy drop while reducing the peak memory usage by up to 2 . 7 × , and it enables a 6 . 4 × larger batch size, which in return resulting in 1 . 2 × higher training throughput.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Jingfeng Yang, Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Bing Yin, and Xia Hu. Harnessing the power of llms in practice: A survey on chatgpt and beyond. arXiv preprint arXiv:2304.13712 , 2023.

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. arXiv preprint arXiv:2001.08361 , 2020.

Ruixiang Tang, Yu-Neng Chuang, and Xia Hu. The science of detecting llm-generated texts. arXiv preprint arXiv:2303.07205 , 2023.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 , 2021.

Yi-Lin Sung, Jaemin Cho, and Mohit Bansal. Lst: Ladder side-tuning for parameter and memory efficient transfer learning. arXiv preprint arXiv:2206.06522 , 2022.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 , 2021.

Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. arXiv preprint arXiv:2106.10199 , 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. arXiv preprint arXiv:2106.09685 , 2021.

Rabeeh Karimi Mahabadi, James Henderson, and Sebastian Ruder. Compacter: Efficient low-rank hypercomplex adapter layers. Advances in Neural Information Processing Systems , 34:1022-1035, 2021.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning , pages 2790-2799. PMLR, 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Actnn: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning , pages 1803-1813. PMLR, 2021.

Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems , 2:497-511, 2020.

Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. arXiv preprint arXiv:2006.09616 , 2020.

Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko. Tempo: Accelerating transformer-based model training through memory footprint reduction. arXiv preprint arXiv:2210.10246 , 2022.

Yu-Neng Chuang, Guanchu Wang, Chia-Yuan Chang, Kwei-Herng Lai, Daochen Zha, Ruixiang Tang, Fan Yang, Alfredo Costilla Reyes, Kaixiong Zhou, Xiaoqian Jiang, et al. Discoverpath: A knowledge refinement and retrieval system for interdisciplinarity on biomedical research. arXiv preprint arXiv:2309.01808 , 2023.

Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations. Advances in Neural Information Processing Systems , 34:27877-27889, 2021.

Petros Drineas, Ravi Kannan, and Michael W Mahoney. Fast monte carlo algorithms for matrices i: Approximating matrix multiplication. SIAM Journal on Computing , 36(1):132-157, 2006.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018.

Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: a system for large-scale machine learning. In Osdi , volume 16, pages 265-283. Savannah, GA, USA, 2016.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Petros Drineas and Ravi Kannan. Fast monte-carlo algorithms for approximate matrix multiplication. In Proceedings 42nd IEEE Symposium on Foundations of Computer Science , pages 452-459. IEEE, 2001.

Brian Chmiel, Itay Hubara, Ron Banner, and Daniel Soudry. Minimum variance unbiased n: Msparsity for the neural gradients. In The Eleventh International Conference on Learning Representations , 2023.

Brian Chmiel, Ron Banner, Elad Hoffer, Hilla Ben Yaacov, and Daniel Soudry. Logarithmic unbiased quantization: Simple 4-bit training in deep learning. arXiv preprint arXiv:2112.10769 , 2021.

Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams. Randomized automatic differentiation. arXiv preprint arXiv:2007.10412 , 2020.

Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for generic network architectures. In International Conference on Machine Learning , pages 14139-14152. PMLR, 2022.

Zirui Liu, Shengyuan Chen, Kaixiong Zhou, Daochen Zha, Xiao Huang, and Xia Hu. Rsc: Accelerating graph neural networks training via randomized sparse computations. arXiv preprint arXiv:2210.10737 , 2022.

Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko. Tempo: Accelerating transformer-based model training through memory footprint reduction. arXiv preprint arXiv:2210.10246 , 2022.

Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint arXiv:2112.00029 , 2021.

Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning , pages 4690-4721. PMLR, 2022.

Shaochen Zhong, Guanqun Zhang, Ningjia Huang, and Shuai Xu. Revisit kernel pruning with lottery regulated grouped convolutions. In International Conference on Learning Representations , 2022.

Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. 2023.

Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning , pages 22137-22176. PMLR, 2023.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.

Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978 , 2023.

Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438 , 2022.

Zhaozhuo Xu, Zirui Liu, Beidi Chen, Yuxin Tang, Jue Wang, Kaixiong Zhou, Xia Hu, and Anshumali Shrivastava. Compress, then prompt: Improving accuracy-efficiency trade-off of llm inference with transferable prompt. arXiv preprint arXiv:2305.11186 , 2023.

Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.11695 , 2023.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: Stateof-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations , pages 38-45, 2020.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. Train big, then compress: Rethinking model size for efficient training and inference of transformers. In International Conference on machine learning , pages 5958-5968. PMLR, 2020.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. MONGOOSE: A learnable LSH framework for efficient neural network training. In International Conference on Learning Representations (ICLR) , 2021.

Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. EXACT: Scalable graph neural networks training via extreme activation compression. In International Conference on Learning Representations , 2022.

Guanchu Wang, Zirui Liu, Zhimeng Jiang, Ninghao Liu, Na Zou, and Xia Hu. Towards memory efficient training via dual activation precision. arXiv preprint arXiv:2208.04187 , 2022.

Guanchu Wang, Zaid Pervaiz Bhat, Zhimeng Jiang, Yi-Wei Chen, Daochen Zha, Alfredo Costilla Reyes, Afshin Niktash, Gorkem Ulkar, Erman Okman, Xuanting Cai, et al. Bed: A real-time object detection system for edge devices. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management , pages 4994-4998, 2022.

Appendix

A Extended Related Work and Discussion

Parameter-Efficient Fine-tuning. Parameter-efficient tuning methods select a small subset of parameters or insert a few parameters to a pre-trained network. Then they only update the small subset of parameters, while keeping others fixed [8, 9, 10, 11, 12, 13, 14, 49]. For example, Adapters [14, 13] insert a small module into the transformer blocks and only update it. Similarly, prompt tuning [8] introduces a small vector that is concatenated with the input embeddings. BitFit [11] only tunes the bias term of the model. LoRA [12] injects trainable rank decomposition matrices into the transformer block. Although these methods are 'parameter-efficient', they actually cannot reduce the memory usage of the model itself. This is because we still needs to build the computation graph for the whole model. Instead, the memory usage of optimizer states will be significantly reduced, which is in proportional to the number of trainable parameters [15].

Activation Quantization. The activation quantization methods focus on quantizing the activation into low numerical precision numbers, e.g., 8-bit integers [16, 30, 50, 51, 52]. Our work also compresses the activation, but in a different way. We emphasize that our work is orthogonal to activation quantization in the sense that our work essentially reduces the dimension of activation. This distinction allows our method to be readily combined with activation quantization techniques, offering the potential for even more aggressive compression.

Gradient Checkpointing. Gradient checkpointing helps decrease activation memory usage by saving only a selection of activations. However, it demands additional computation during the backward pass, as discarded activations must be recalculated [18, 17]. According to the report of Checkmate 3 [17], it achieves 'a 2.3x memory reduction when training a BERT model with Checkmate optimizations (at 1x extra overhead for rematerialization)'.

Limitations Although WTA-CRS significantly reduces the computation of the backward pass in a hardware-friendly way i.e., dropping entire rows/columns in the tensor, the current implementation still hampers the execution time of linear operations. This is because the extra sampling process and data movement counteract the acceleration. However, we note that (1) the overhead can be greatly reduced with better implementation, e.g., using prefetch and operation-fusion technique [30]; (2) the existing implementation can still yield a large speedup when employing larger batch sizes (Figure 8).

Potential Negative Societal Impacts. Our research primarily focuses on reducing the memory requirement of fine-tuning Language Models (LMs). The carbon emissions produced by LM finetuning may pose environmental issues. Our next step is to further improve the efficiency of LM fine-tuning, particularly on hardware with lower energy consumption.

B Unbiasedness of Weight Gradient

This part we directly follow the proof of Theorem 1 in ActNN [16]. For completeness, we provide the proof sketch here that is short and easy to follow. Specifically, here we use ReLU as the activation function for illustration convenience. We note that the conclusion in this section holds for any non-linear activation function. Specifically, the forward pass of ReLU-Linear at the l th layer is

and the backward pass of ReLU is:

where ⊙ is the element-wise product and ✶ is the indicator function. The element-wise product is linear operation and ✶ Z ( l +1) > 0 is only related to the pre-activation Z ( l +1) in the forward pass. We only apply the approximation during the backward pass so ✶ Z ( l +1) > 0 can be extracted from the expectation. We know that for the last layer L , we have E [ ∇ H ( L ) ] = H ( L ) since we do not apply activation at the output layer. We then can prove by induction that E [ ∇ H ( l +1) ] = H ( l +1) and E [ ∇ W ( l ) ] = W ( l ) for any layer l .

C Proof

C.1 Derivation of Equation (3)

Let X ∈ R n × m , Y ∈ R m × q be two matrices. The matrix multiplication XY can be estimated as

Equation (3) shows the approximation error E [ || XY -X ′ Y ′ || F ] is minimized when the probabilities

Proof. Let f ( i ) = X : i Y i : p i ∈ R n × q . We note that f ( i ) is an unbiased estimation of XY . Namely,

where i 1 , · · · , i t are the index of the sampled column-row pairs at t th random trials. For each i t , its variance is

where the first step follows from the fact that Var[ x ] = E [ x 2 ] -E 2 [ x ] . Then we have,

Then we have

By combining Equation (8) and Equation (9) into the above equation, we have

To minimize E [ || XY -X ′ Y ′ || F ] , the optimal probability distribution can be obtained via solving the following optimization problem:

The solution to the above convex problem is the distribution defined in Equation (3). Namely,

C.2 Unbiasedness of Our Proposed Estimator

Theorem 1 (Proof in Appendix C.2) . The estimator defined in Equation (4) is an unbiased estimator for matrix production XY , i.e, E j ∼P D\C [ ∑ c ∈C f ( c ) p c +(1 -∑ c ∈C p c ) f ( j )] = XY .

Proof.

C.3 Variance of Our Proposed Estimator

Theorem 2 (Proof in Appendix C.3) . Suppose the total budget of column-row pairs is k . If C satisfies

then we have Var[ˆ g ( X , Y )] < Var[ g ( X , Y )] . Moreover, Var[ˆ g ( X , Y )] is minimized when |C| = min |C|∈{ 0 , ··· ,k } 1 -∑ c ∈C p c k -|C| .

Proof. Recall that the original estimator for matrix production XY is defined as

and our proposed family of estimator is defined as:

We first define three independent random variables as belows:

According to the Law of total variance, we have

where the first step follows from the fact that for any random variance x , y , we have Var[ y ] = E [Var[ y | x ]] + Var[ E [ y | x ]] . Also, by Equation (11), we have

By combining the above two inequality, we have

Equation (17) quantitatively shows the variance reduction of h ( j ) over f ( i ) . Then we compare our estimator ˆ g ( X , Y ) and g ( X , Y ) in terms of variance.

First, because g ( X , Y ) = 1 k ∑ k t =1 f ( i t ) , i 1 , · · · i k i.i.d ∼ P . Thus we have

By combining Equation (17) into the above two equations, we have

Similarly, we have

where the first step follows from Equation (17). By setting 1 -∑ c ∈C p c k -|C| k ≤ 1 , we arrive the conclusion that when ∑ c ∈C p c > |C| k , we have Var[ˆ g ( X , Y )] ≤ Var[ g ( X , Y )] .

D Implementation Details

The pseudocode for approximated linear layer with WTA-CRS and standard line layer is given in Algorithm 1 and Algorithm 3, respectively. The column-row pair sampling procedure is given in Algorithm 2. For the ease of illustration, we ignore the sequential length. As we mentioned in the main text, we only replace the GEMM in the backward pass with WTA-CRS . According to Equation (1c), we need the activation gradient ∇ Z to perform the column-row pair sampling during the forward pass. Thus we initialize a cache in CPU memory to store the gradient norm of activations from the last step. When performing column-row pair selection, we need to swap the gradient norm of activations between CPU and GPU, which will cause extra time overhead due to the data movement. Fortunately, we note that the number of elements in the gradient norm of activations is significantly less than the one in activations, which does not cause a significant time overhead.

Algorithm 2: SUBSAMPLE

Algorithm 3: Forward & Backward pass of the standard Linear layer

E More Experimental Results

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

Probability Mass

c

pc

|

|

)

2-th layer

4-th layer

6-th layer

|

|

k

Ratio

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

|

|

)

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.2

0.4

0.6

|

|

)

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.25

0.50

0.75

Probability Mass

c

pc

|

|

)

2-th layer

4-th layer

6-th layer

|

|

k

Ratio

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.25

0.50

0.75

|

|

)

0.1

0.2

0.3

0.4

0.5

0.6

0.7

|

|

k

Ratio

0.25

0.50

0.75

|

|

)

0

20

40

60

80

100

# Iteration

0.6

0.7

Probability Mass

c

pc

|

|

|

|

=0.1

)

2-th layer

4-th layer

6-th layer

0

20

40

60

80

100

# Iteration

0.6

0.7

|

|

|

|

=0.1

)

0

20

40

60

80

100

# Iteration

0.6

0.7

|

|

|

|

=0.1

)

2

4

6

8

10

Epoch

92

94

Accuracy (%)

(a) SST2

2

4

6

8

10

Epoch

70

80

Accuracy (%)

(b) MNLI

2

4

6

8

10

Epoch

60

80

F1 score (%)

Deterministic

CRS

WTA-CRS

(c) QQP

E.1 More Experimental Analysis on Theorem 2

To evaluate Theorem 2 more comprehensively, below we also plot the ∑ c ∈C p c versus |C| k at k = 0 . 1 |D| and k = 0 . 5 |D| in Figure 9 and 10, respectively. We also plot ∑ c ∈C p c versus iterations in Figure 11. We summarize that the the column-row index distribution is concentrated on a few column-row pairs. Thus, the assumption in Theorem 2 holds under the context of fine-tuning transformers.

E.2 More Experimental Speed Analysis

Increasing the batch size can often result in faster model convergence and/or enhance the final performance. Ideally, we should adjust the batch size according to the requirements of our model

rather than being constrained by the GPU's memory capacity. To illustrate this, we have represented the correlation between peak memory usage and maximum mini-batch size for T5-Base, T5-Large, and T5-3B in Figure 13. Our observations highlight that WTA-CRS effectively increases the maximum available batch size.

We also provide the apple-to-apple speed comparison for linear operation with and without WTA-CRS in Table 3. In Table 3, 'Fwd', 'Bwd', and 'F-B' are the time of forward pass, the time of backward pass, and the total time for both the forward and backward pass, respectively. We summarize that under the same workload, the current implementation of WTA-CRS may roughly slow down the linear operation about 20% . This is because the extra sampling process and data movement counteract the acceleration (see Algorithm 1). However, we note that (1) the overhead can be greatly reduced with better implementation, e.g., using prefetch and operation-fusion technique [30]; (2) the existing implementation can still yield a large speedup when employing larger batch sizes (Figure 8).

500

1000

1500

Batch-size

20

40

60

80

Full

LoRA

LoRA+WTA-CRS@0.3

LoRA+WTA-CRS@0.1

(a) T5-Base

200

400

600

Batch-size

20

40

60

80

Full

LoRA

LoRA+WTA-CRS@0.3

LoRA+WTA-CRS@0.1

(b) T5-Large

100

200

Batch-size

40

60

80

Full

LoRA

LoRA+WTA-CRS@0.3

LoRA+WTA-CRS@0.1

(c) T5-3B

E.3 Performance of WTA-CRS on OPT models

We conduct experiments of WTA-CRS on OPT models. It is observed in Table 4 that WTA-CRS exhibits almost no drop in accuracy compared to Full training and LoRA on all three architectures. This observation indicates the effectiveness of WTA-CRS, especially considering its consistent performance across the encoder-only, decoder-only and encoder-decoder transformer architectures.

E.4 Performance of Activation Quantization+ WTA-CRS

To demonstrate the orthogonality of WTA-CRS with activation quantization, we conducted additional experiments of combining WTA-CRS@0.3, along with activation quantization@8bit, on the T5-

base model. The experiment results are presented in Table 5. It is notabe that the combination of LoRA+WTA-CRS@0.3+Quant@8bit exhibits almost no drop in accuracy compared with Full training, LoRA, WTA-CRS@0.3, and LoRA+WTA-CRS@0.3. This observation clearly demonstrates the orthogonality of WTA-CRS with activation quantization, indicating that they can be effectively applied together without compromising performance.

E.5 Comparison of WTA-CRS with QLoRA

We conduct experiments to compare WTA-CRS with QLoRA on the T5-Base model. Specifically, we first quantize the model into NF4 data format using 'bitandbytes' and keep it fixed. Then we apply LoRA+ WTA-CRS @0.3 over the quantized model. Below in Table 6, we demonstrate that applying QLoRA and WTA-CRS over a base model quantized in NF4 results in almost no performance loss against the QLoRA baseline. However, such joint applications may enjoy the training-time memory efficiency offered by WTA-CRS, as illustrated in Table 7.

F Experimental Settings

We give the detailed hyper-parameter setting in this section. Specifically, for both T5 and BERT models, the parameters are updated with the AdamW optimizer β 1 = 0 . 9 β 2 = 0 . 999 ϵ = 10 -8 and weight decay = 0 . The the learning rate is adjusted with a linear LR Scheduler, which maintains a constant learning rate for the initial 500 steps, and adjusts it gradually thereafter. The input sequences are padded to the maximum length 128. WTA-CRS has a LoRA dimension 32 if it is combined with LoRA. To achieve the optimal solution, the T5-Base, Large, 3B and BERT-Base and Large models have different learning rate, training epoch number, and mini-batch size on different datasets, which are given in Tables 9, 10, 11, respectively.

F.1 Computational Infrastructure

The computational infrastructure information is given in Table 8.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n>**Rebuttal:** The provided details satisfy my concerns. I think this paper should be accepted after applying the agreed changes.\n\n>**TL;DR:** **Good paper.** The proposed WTA-CRS algorithm is based on the existing CRS algorithm and is used to reduce activation memory during training. WTA-CRS achieves up to 2.7× peak memory reduction with almost no accuracy drop and enables up to 6.4× larger batch size. However, WTA-CRS comes with computational overhead, which is discussed and explore. Addressing my concerns and questions would improve my score.\n\nThe paper proposes the WTA-CRS algorithm to reduce the neural networks training activation memory, where the paper claims that activation memory is primary memory bottleneck during training. The WTA-CRS algorithm is an unbiased estimators for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient. WTA-CRS achieves up to 2.7× peak memory reduction with almost no accuracy drop and enables up to 6.4× larger batch size.\n\nThe WTA-CRS algorithm works by sampling columns and rows to create an unbiased estimation of the original GEMM for the backpropagation. The WTA-CRS algorithm does not alter the neural architecture, and therefore the inference speed is left in tact. The experimental section shows that WTA-CRS outperforms existing prior work and is compatible with existing PEFT techniques. WTA-CRS adds a computational overhead due to sampling, however, WTA-CRS enables training on much larger batch sizes, which results in a 1.2× higher training throughput.\n\nSTRENGTHS:\n* **S.1.** The proposed WTA-CRS algorithm tackles an important problem in existing PEFT techniques, which makes LLM PEFT training more accessible to researchers with low resources.\n* **S.2.** The paper provides a theoretical analysis on WTA-CRS.\n* **S.3.** The proposed WTA-CRS algorithm outperform existing algorithms.\n* **S.4.** An anonymized code repository is provided as part of the submission for reproduction .\n\nWEAKNESSES:\n* **W.1.** Popular existing memory efficient training techniques such as tensor rematerialization (gradient checkpointing) [2][3] and ZeRO [1] are not compared to, although some are partially discussed in Appendix A.\n* **W.2.** The experiments are conducted on single neural network architecture (T5), although the proposed technique does not seem to be confined solely to that setting.\n* **W.3.** It is common practice today to train neural networks at a lower precision (quantization), however, it is not clear whether quantization (16bit) was used. Therefore, there is insufficient proof that the combined noise of WTA-CRS and quantization would be compatible.\n\n\n**Typos.**\n* Line #62: \"Thus\" → \"Thus,\"\n* Line #240: \"mAccording\" → \"According\"\n* Line #297: \"Thus\" → \"Thus,\"\n\n[1] Ren, J., Rajbhandari, S., Aminabadi, R.Y., Ruwase, O., Yang, S., Zhang, M., Li, D. and He, Y., 2021, July. ZeRO-Offload: Democratizing Billion-Scale Model Training. In USENIX Annual Technical Conference (pp. 551-564).\n\n[2] Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K. and Stoica, I., 2020. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2, pp.497-511.\n\n[3] Beaumont, O., Eyraud-Dubois, L. and Shilova, A., 2021. Efficient combination of rematerialization and offloading for training dnns. Advances in Neural Information Processing Systems, 34, pp.23844-23857.\n\nQUESTIONS:\n* **Q.1.** In line #43 and Figure 2 it is noted that \"storing activations (or feature maps) is the main memory bottleneck during training\". Does this hold true for all model architectures? What about LLM training where the fine-tuning batch size is usually very small?\n* **Q.2.** Why was the WTA-CRS algorithm compared to the Deterministic top-k from [1] but not to the Bernoulli-CRS from [1]? What are the key differences between WTA-CRS and Bernoulli-CRS?\n* **Q.3.** The paper proposes WTA-CRS which sacrifices computation speed at the cost of lower peak memory. There are several existing common approaches (such as gradient checkpointing and DeepSpeed) for general memory efficient training which are compatible with PEFT techniques. Why are these comparisons not explored or detailed in the main paper?\n\n[1] Adelman, Menachem, Kfir Levy, Ido Hakimi, and Mark Silberstein. \"Faster neural network training with approximate tensor operations.\" Advances in Neural Information Processing Systems 34 (2021): 27877-27889.\n\nLIMITATIONS:\nThe limitations are discussed in Appendix A.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, we propose a new method called WTA-CRS (Winner-Take-All Column Row Sampling) to address the main memory bottleneck issue during training, which arises from storing feature maps. To reduce memory usage during training, we sample the most likely column indices during backpropagation.\n\nFurthermore, \bthey proposed method demonstrates the ability to significantly reduce peak memory usage, by approximately up to 2.7 times, when fine-tuning downstream tasks. It also showcases the potential for higher throughput, enabling more efficient training.\n\nSTRENGTHS:\n1. The work clearly states its motivation and its solution and is easy to follow.\n2. The authors show that their method reaches comparable performance with backpropagation using the full activation when combined with LoRA.\n3. They also empirically measure throughput gains obtained by increasing batch size, which demonstrates the practical applicability of their method.\n\nWEAKNESSES:\n1. The paper needs a comparative analysis of other researchs, such as gradient checkpoint/recalculation and CRS, aimed at reducing activation memory during the training phase, as shown in Fig. 6 and Fig. 9.\n2. The paper should include an analysis of the overhead associated with the proposed WTS-CRS method, which involves sampling rows and columns. It is crucial to consider factors such as the computational cost of Equation 3 and any potential effects of lowering on the overall performance. Providing this analysis would enhance the clarity and completeness of the research.\n3. There is a need of analysis on the effectiveness of the proposed approach, WTS-CRS, in distributed training environments such as tensor parallelism or pipeline parallelism.\n4. It seems necessary to conduct performance evaluations on various LLMs of the GPT family, such as LLaMA and OPT.\n\nQUESTIONS:\n* In Figure 9, it can be observed that the throughput of WTS-CRS is lower than that of full when the batch size is small. Is this due to the overhead caused by lowering?\n* When comparing the training throughput, how does CRS differ from full in terms of throughput?\n* Could the authors include statistics for GPU utilization in their experiments? It would be helpful to analyze the causes of improved performance more thoroughly.\n* Considering that most large models are trained using multiple levels of parallelism, would it be possible to verify results for pipeline parallel, tensor parallel, etc.? Also, it is unclear from the paper whether the data parallelism used was distributed data parallelism or naïve data parallelism.\n\nLIMITATIONS:\n* As previously mentioned, it would be valuable to include additional experimental results for models that are more challenging to quantify, such as GPT-series (OPT, LLaMA). This would enhance the validity and applicability of the proposed method across a broader range of models.\n* Considering that most large-scale models are trained using multiple levels of parallelism, it is important to assess how much the proposed methods, such as pipeline parallelism and tensor parallelism, can increase throughput while taking into account overhead (such as GPU-to-GPU or node-to-node communication), memory reduction, and computational cost. Furthermore, it is not clear from the paper whether the data parallel processing used is distributed data parallelism or naive data parallelism.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors studied fine-tuning LLMs with limited memory. As the increased scale of current LLMs, the memory cost during fine-tuning is of great importance when adapting the pretrained LLMs to down-streaming tasks. In contrast to the existing work that mainly focus on the number of updated weights, this paper proposed to reduce the number of stored activations, also the inputs to each layer. Given the widely used stochastic gradient descent optimization pipeline, the authors proposed to store a subset of activations that can generate an unbiased gradient estimation. This way, the training memory and the training time decreased significantly. The authors provide both theoretical and experimental analysis on their CRS methods.\n\nSTRENGTHS:\n- This paper studied an important problem in LLM fine-tuning, i.e., how to fine-tuning LLMs with less memory consumption without increasing the computation cost. The authors provided solid quantitative results to show that the main memory consumption is from storing the intermediate activations. \n- The authors provided a general solution for fine-tuning LLMs under memory constraints. The solution can be applied in most transformer-based network architectures.  \n- The authors provided solid mathematical proof on the unbiased gradient estimation, which is especially encouraged. \n- The extensive experiments on different network architectures showed the efficacy of the methods.\n- The released code can benefit the following researchers studying efficient LLM fine-tuning.\n\nWEAKNESSES:\n- I am not fully convinced by the comment made in Line241-244, i.e., the methods in the paper is orthogonal to the activation quantization. When activation is quantized into a lower bit width, it is very possible that the number of less important activations will decrease. This way, the selection on the top-k columns in activation matrices with the proposed methods may hurt the training accuracy or convergence. It would be great if the authors can provide some theoretical analysis or experimental results on this combination. Otherwise, it would be necessary to provide some comparison results w.r.t. the activation quantization.\n- It would be great if the authors can discuss the main difference of their paper w.r.t. [Randomized Automatic Differentiation, ICLR2021].\n\nQUESTIONS:\nOverall, I think this paper has a relatively high quality in both writing and scientific contribution.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper's contribution is in proposing a practical, intuitive yet not trivial unbiased approximation to gradient training of matrix multiplication. It shows that even though totally deterministic sampling is biased, somewhat deterministic sampling is unbiased, and a judicious allocation of sampling to those pairs favored by deterministic thinking can lead to the use of a larger batch size with empirically negligible performance loss. This reviewer must declare that he does not check the derivation very carefully.

STRENGTHS:
The proposed idea is practical and can be readily combined with virtually all first-order gradient-based training methods.
The paper also derived why deterministic sampling is a biased estimator and empirically shown the associated bad performance, thus proving that the additional complexity of stochastic sampling over deterministic sampling is not only sufficiently better but also necessary.

WEAKNESSES:
It's just a few empirical comparisons, but the performance gap between CRS and WTA-CRS seems modest.

QUESTIONS:
This reviewer does not have a question.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06