Title: CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor Optimization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Compressed Training

Memory-efficient

Tensor
Abstract: Training large AI models such as LLMs and DLRMs costs massive GPUs and computing time. The high training cost has become only affordable to big tech companies, meanwhile also causing increasing concerns about the environmental impact. This paper presents CoMERA, a **Co**mputing- and **M**emory-**E**fficient training method via **R**ank-**A**daptive tensor optimization. CoMERA achieves end-to-end rank-adaptive tensor-compressed training via a multi-objective optimization formulation, and improves the training to provide both a high compression ratio and excellent accuracy in the training process. Our optimized numerical computation (e.g., optimized tensorized embedding and tensor-vector contractions) and GPU implementation eliminate part of the run-time overhead in the tensorized training on GPU. This leads to, for the first time, $2-3\times$ speedup per training epoch compared with standard training. CoMERA also outperforms the recent GaLore in terms of both memory and computing efficiency. Specifically, CoMERA is $2\times$ faster per training epoch and $9\times$ more memory-efficient than GaLore on a tested six-encoder transformer with single-batch training. Our method also shows $\sim 2\times$ speedup than standard pre-training on a BERT-like code-generation LLM while achieving $4.23\times$  compression ratio in pre-training.
With further HPC optimization, CoMERA may reduce the pre-training cost of many other LLMs.  An implementation of CoMERA is available at <https://github.com/ziyangjoy/CoMERA>.

CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor Optimization

Abstract

Training large AI models such as LLMs and DLRMs costs massive GPUs and computing time. The high training cost has become only affordable to big tech companies, meanwhile also causing increasing concerns about the environmental impact. This paper presents CoMERA, a Co mputing- and M emoryE fficient training method via R ankA daptive tensor optimization. CoMERA achieves rankadaptive tensor-compressed (pre)-training via a multi-objective optimization formulation and improves the training to provide both a high compression ratio and excellent accuracy in the training process. Our optimized numerical computation (e.g., optimized tensorized embedding and tensor-network contractions) and GPU implementation eliminate part of the run-time overhead in the tensorized training on GPU. This leads to, for the first time, 2 -3 × speedup per training epoch compared with standard training. CoMERA also outperforms the recent GaLore in terms of both memory and computing efficiency. Specifically, CoMERA is 2 × faster per training epoch and 9 × more memory-efficient than GaLore on a tested six-encoder transformer with single-batch training. Our method also shows ∼ 2 × speedup than standard pre-training on a BERT-like code-generation LLM while achieving 4 . 23 × compression ratio in pre-training. With further HPC optimization, CoMERA may reduce the pre-training cost of many other LLMs. An implementation of CoMERA is available at https://github.com/ziyangjoy/CoMERA .

1 Introduction

Deep neural networks have gained success in solving numerous engineering problems. These approaches usually use a huge number of variables to parametrize a network, and require massive hardware resources to train the model. For instance, the Deep Learning Recommendation Model (DLRM) released by Meta (which is smaller than the practical product) [32] has 4.2 billion parameters; GPT-3 [4] has 175 billion parameters. OpenAI shows that the computing power required for key AI tasks has doubled every 3.4 months [1] since 2012. Training a large language model like ChatGPT and LLaMA from scratch often takes several weeks or months on thousands of GPUs [3, 2].

(a) Training time per epoch.

(b) Peak memory consumption.

100

CoMERA

GaLore

LTE

CoMERA

GaLore

LTE

80

78.3

B

10000

10922

m

n

60

45.3

memory

>

6400

6998

7710

time

40

5000

3358

4320

3800

4532

20

14.8

19.8

13.8 16.5

13.2 15.5

1674

2174

N/A

7.2

6.4

5.5

184

N/A

32

64

128

1

32

64

128

batch size

batch size

Large AI models often have much redundancy. Therefore, numerous methods have been developed to reduce the cost of AI inference [5, 25, 7, 13, 16, 31, 30, 17]). However, training large AI models (especially from scratch) remains an extremely challenging task. Low-precision training [19, 27, 12, 36] has been popular in on-device setting, but its memory reduction is quite limited. Furthermore, it is hard to utilize ultra low-precision training on GPU since current GPUs only support limited precision formats for truly quantized training. Chen [6] employed the idea of robust matrix factorization to reduce the training cost. Similar low-rank matrix approximation techniques have been applied to train large AI models including large language models [41, 20]. Among them, GaLore [41] can train the 7B LLaMA model on an RTX 4090 GPU using a single-batch and layer-wise setting. However, this setting can lead to extremely long training time, which is infeasible in practical settings.

Compared with matrix compression, low-rank tensor compression has achieved much higher compression ratios on various neural networks [11, 22, 25, 29, 33, 37, 40, 43]. This idea has been studied in structure search for compact representations [43], post-training compression [11], fixed-rank training [5, 33, 24], zeroth-order training [42] and parameter-efficient fine tuning [39]. The recent work [14, 15] provides a rank-adaptive tensor-compressed training from a Bayesian perspective. However, to achieve a reduction in both memory and training time (especially on transformers), two open problems need to be addressed. Firstly, a more robust rank-adaptive tensor-compressed training model is desired, since the method in [14, 15] relies on a heuristic fixed-rank warm-up training. Secondly, while modern GPUs are well-optimized for large-size matrix computations, they are unfriendly for low-rank tensor-compressed training. Specifically, most operations in tensor-compressed training are small-size tensor contractions, which can cause significant runtime overhead on GPUs even though the theoretical computing FLOPS is very low. As a result, as far as we know no papers have reported real training speedup on GPU. This issue was also observed in [18]. SVDinsTN [43] controls tensor ranks to search for a compact tensor structure of a given tensor. HEAT [11] uses tensor decompositions for post-training model compression of trained models. Detailed comparisons with these works are shown in Appendix A.1.

Paper Contributions. In this work, we propose CoMERA, a tensor-compressed training method that can achieve, for the first time, simultaneous reduction of both memory and runtime on GPU . Our specific contributions are summarized as follows.

Multi-Objective Optimization for Rank-Adaptive Tensor-Compressed Training . We propose a multi-objective optimization formulation to balance the compression ratio and model accuracy and to customize the model for a specific resource requirement. One by-product of this method is the partial capability of automatic architecture search: some layers are identified as unnecessary and can be completely removed by rank-adaptive training.

Performance Optimization of Tensor-Compressed Training. While tensor-compressed training greatly reduces the memory cost and computing FLOPS, it often slows down the practical training on GPU. We propose three approaches to achieve real training speedup: 1 ⃝ optimizing the lookup process of tensorized embedding tables, 2 ⃝ optimizing the tensor-network contractions in both forward and backward propagation, 3 ⃝ eliminating the GPU backend overhead via CUDA Graph.

Experimental Results. Weevaluate our method on the end-to-end training of a transformer with six encoders and the deep learning recommendation system model (DLRM). On these two benchmarks, our method achieves 80 × and 99 × compression ratios respectively, while maintaining the testing accuracy of standard uncompressed training. CoMERA also achieves 2 -3 × speedup per training epoch compared with standard training methods on the transformer model. In a preliminary study of LLM pre-training, CoMERA shows 1 . 9 × to 2 . 3 × speedup in different pre-training stages on CodeBERT [21], while achieving 4 . 23 × overall model reduction in the pre-training process.

matrix-vector

matrix-matrix

multiplication

multiplication

Order-1

Order-2

Order-3

Order-4

tensor-tensor

multi-tensor

contraction

contraction

(a)

(b)

Figure 1 compares our CoMERA with GaLore [41] and the recent LoRA-based training method LTE [20] on the six-encoder transformer. When data and back-end memory cost are considered, CoMERA's memory consumption is 9 × less than GaLore in the single-batch training as adopted in [41], and it uses the least memory under all batch sizes. Our method is 2 -3 × faster than GaLore and LTE in each training epoch, although CoMERA has not yet been fully optimized on GPU.

While this work focuses on reducing the memory and computing cost of training, it can also reduce the communication cost by orders of magnitude: only low-rank tensorized model parameters and gradients need to be communicated in a distributed setting. The CoMERA framework can also be implemented on resource-constraint edge devices to achieve energy-efficient on-device learning.

2 Background

The tensor [26, 28] A ∈ R n 1 × n 2 ×···× n d is indexed as A = ( a i 1 ··· i d ) 1 ≤ i j ≤ n j and is said to have order d and dimension n 1 , . . . , n d . The Frobenius norm of tensor A is defined as ∥ A ∥ := √ ∑ n 1 ,...,n d i 1 ,...,i d a 2 i 1 ··· i d . In tensor networks, the orderd tensor A is represented as a node with d edges. Some tensor network representations are illustrated in Fig. 2 (a).

̸

Tensor Contraction. Let A ∈ R n 1 × n 2 ×···× n d and B ∈ R l 1 × l 2 ×···× l m be two tensors with n s = l t . The tensor contraction product C = A × s,t B has dimension Π i = s n i × Π j = t l j and the entries are

̸

̸

̸

This definition can be naturally generalized to multiple pairs. Figure 2(b) illustrates some tensor contractions. For general operations among multiple tensors, we use PyTorch einsum in the following

where each S i is a string of characters that specifies the dimension of A i . The output tensor B is obtained by summing over all other dimensions that are not in T . In the following, we show a few commonly used einsum operations. The Tensor-Train decomposition as in (3) is

For the batched matrices A ∈ R b × m × k , B ∈ R b × k × n , the batched matrix multiplication is

Tensor Decomposition. In this paper, we will mainly use tensor-train (TT) [35] and tensor-train matrix (TTM) [34] decomposition for compressed neural network training. TT [35] represents the tensor A ∈ R n 1 × n 2 ×···× n d as a set of small-size cores G 1 , . . . , G d such that G i ∈ R r i -1 × n i × r i and

The tuple ( r 0 , r 1 , . . . , r d ) is the TT rank of the TT decomposition (3) and must satisfy r 0 = r d = 1 . TTM considers an order2 d tensor B of dimension m 1 × n 2 ×··· × m d × n d , and represents B as

where F i ∈ R r i -1 × m i × n i × r i for i = 1 , . . . , d and r 0 = r d = 1 . Figure 3 shows the tensor-network representations of TT and TTM decomposition.

In tensor-compressed neural networks, large weight matrices are reshaped to high-order tensors and compressed into small tensor cores in TT or TTM format. The weights of linear layers are often compressed into the TT format due to its efficiency in tensor-vector multiplications. The TTM format is more suitable for embedding tables whose dimension is highly unbalanced.

n1

n2

nd-1

nd

n1

2

nd-1

Pu

r1

G2

r2

rd-2

Gd-1

rd-1

Gd

r1

92

r2

rd-2

Gd-1

rd-1

Gd

m1

m2

md-1

md

a)

(b)

3 The CoMERA Training Framework

The size of the tensor-compressed neural networks can be adjusted by modifying the tensor ranks. However, it also brings in an important problem: how can we determine the tensor ranks automatically for a given resource limit? We propose a multi-objective optimization to address this issue.

3.1 Multi-Objective Training Model

A Modified TT Representation. We consider the tensor-compressed training for a generic neural network. Suppose that the neural network is parameterized as f ( x |{ G i 1 , . . . , G i d i } P i =1 ) , where { G i j ∈ R r i j -1 × n i j × r i j } d i j =1 compress the original weight W i . Let { x k , y k } N k =1 be training data and L be the loss function. The training is to minimize the following objective function

We modify the TT compression and control the ranks of G i 1 , . . . , G i d i by a set of diagonal matrices { D i j ∈ R r i j × r i j } d -1 j =1 . Specifically, let W i be the reshape of W i , and the compression of W i is

Now the tensor cores for W i have S i = n i 1 ∥ D i 1 ∥ 0 + n i d i ∥ D i d i -1 ∥ 0 + ∑ d i -1 j =2 n i j ∥ D i j -1 ∥ 0 ∥ D i j ∥ 0 variables. For simplicity, we denote G := { G i 1 , . . . , G i d i } P i =1 and D := { D i 1 , . . . , D i d i -1 } P i =1 .

Multi-Objective Optimization. We intend to minimize both the loss and compressed network size, which can be formulated as a multi-objective optimization min G , D { L ( G , D ) , S ( D ) } , where S ( D ) := ∑ P i =1 S i ( D ) . In most cases, we cannot find a point that minimizes the loss and model size simultaneously. Therefore, we look for a Pareto point ( G ∗ , D ∗ ) , meaning that there exist no G and D such that L ( G , D ) ≤ L ( G ∗ , D ∗ ) , S ( D ) ≤ S ( D ∗ ) , and at least one of inequalities is strict.

3.2 Training Methods

We convert a multi-objective optimization to a single-objective one via scalarization. We use different scalarization methods at the early and late stage of training. The late stage is optional, and it can further compress the model to enable efficient deployment on resource-constraint platforms.

Early Stage. At the early stage of CoMERA, aggressively pruning ranks dramatically hurts the convergence. Hence, we start the training with the following linear scalarization formulation [8]

It is still hard to solve (7) since S ( D ) uses ∥ · ∥ 0 which is nonsmooth. Therefore, we replace ∥ · ∥ 0 by the ℓ 1 norm ∥ · ∥ 1 and get the convex relaxation

We note that ˆ S ( D ) can be arbitrarily close to 0 while keeping L ( G , D ) unchangeable, since the corresponding slices of TT factors can be scaled accordingly. Therefore, a direct relaxation of the scalarization (11) does not have a minimizer. To address this issue, we add an ℓ 2 regularization ∥ G ∥ 2 := ∑ P i =1 ∑ d i j =1 ∥ G i j ∥ 2 to the relaxation and get the formulation

The optimizer of Problem (9) is a Pareto point for a constrained problem, shown in the following.

Proposition 3.1. For all γ > 0 , β > 0 , there exists some constant C > 0 such that the solution to the problem (9) is a Pareto point of the following multi-objective optimization problem

Proof. See Appendix A.2 for the complete proof. Late Stage (Optional). The early-stage training can provide us with a Pareto point, but we cannot control where the Pareto point is. In the late stage of CoMERA, we may continue training the model towards a preferred loss L 0 and a preferred model size S 0 for deployment requirements. This can be achieved by the achievement scalarization [8] that leads to a Pareto point close to ( L 0 , S 0 ) :

Here w 1 , w 2 > 0 scale the objectives into proper ranges, and ρ > 0 is a small constant. After relaxing S ( D ) to ˆ S ( D ) and adding the regularization term, we get the following problem

where β > 0 is a positive constant. Note that the S ( D ) inside max is not relaxed now for accurate comparisons. When w 1 ( L ( G , D ) -L 0 ) ≥ w 2 ( S ( D ) -S 0 ) , we consider the following problem

Werun a step of a gradient-based algorithm on this problem. When w 1 ( L ( G , D ) -L 0 ) < w 2 ( S ( D ) -S 0 ) , we relax the S ( D ) again and get the following problem

and run a step of a gradient-based algorithm on this problem. The Algorithm 1 is summarized in Appendix A.3. The late stage optimization can be independently applied to a trained tensorcompressed model for further model size reductions.

4 Performance Optimization of CoMERA

While CoMERA can greatly reduce training variables and memory cost, the low-rank and small-size tensor operations in CoMERA are not efficiently supported by GPU. This often slows the training process. This section presents three methods to achieve real training speedup on GPU.

4.1 Performance Optimization of TTM Embedding Tables.

Embedding tables are widely used to transfer discrete features into continuous hidden space. The row size of embedding tables is usually much larger than the column size, making TTM compression more suitable than the TT format. In the following, we use an order-4 TTM embedding table to illustrate how to accelerate the lookup process.

Weconsider an embedding table T ∈ R m × n . Alookup operation selects the submatrix T [ I , :] ∈ R b × n for the index set I = { i k } b k =1 . This operation is fast and inexpensive. However, the full embedding table itself is extremely memory-consuming. Suppose that

m = m 1 m 2 m 3 m 4 , n = n 1 n 2 n 3 n 4 , then we reshape T into tensor T ∈ R m 1 × n 1 ×···× m 4 × n 4 and represent it in TTM format

The compressed embedding table does not have the matrix T explicitly. We convert each row index i k ∈ I to a tensor index vector ( z k 1 , z k 2 , z k 3 , z k 4 ) and denote Z t = { z k t } b k =1 , then T [ I , :] can be computed by contracting the tensors { G t [: , Z t , : , :] } 4 t =1 where each has size r t -1 × b × n i × n t . The G t [: , Z t , : , :] stores many duplicated values especially when the set I is large. Therefore, directly computing the tensor contractions can cause much computing and memory overhead.

(a) Speed-up

(b) Memory reduction

8

optimized order

uncompressed

proposed approach

6.87

8.16

6000

5000

5279

5284

5294

uncompressed

proposed approach

optimized order

S

unique indices

without optimization

B

4000

without optimization

unique indices

time

4

3.69

4.01

memory

3000

2000

17631700

2

0.06

0.48 0.47

2.01 2.12

0.11

0.82 0.82

0.22

1.42 1.61

1000

670 662

1167 1107

799811

896 939

10000

20000

40000

0

10000

20000

40000

batch size

batch size

We optimize the tensor contraction by eliminating the redundant computation at two levels. Rowindex level. We construct the index set I u = { i k } ˆ b k =1 containing all unique indices in I . We can easily obtain T [ I , :] from T [ I u , :] . Tensor-index level. The reduced index set I u leads to ˆ b associated tensor index vectors ( z k 1 , z k 2 , z k 3 , z k 4 ) , but at most m 1 m 2 pairs of ( z k 1 , z k 2 ) and m 3 m 4 pairs of ( z k 3 , z k 4 ) are unique. For instance, (2,3,1,3) and (2,3,2,4) are common in (2,3), so we only compute (2,3) entry once. Therefore, we can consider all unique pairs ( z k 1 , z k 2 ) and ( z k 3 , z k 4 ) and compute

For each i k ∈ I u , let ( j k 1 , j k 2 ) be the coordinate of i k for size ( m 1 m 2 , m 3 m 4 ) . We denote J 1 = { j k 1 } ˆ b k =1 and J 2 = { j k 2 } ˆ b k =1 , then compute the unique rows of T as

Figure 4 summarizes the whole process of TTM embedding table look-up. This approach can be easily applied to higher-order embedding tables by first grouping some small tensor cores to obtain intermediate tensors and then utilizing them to compute unique row vectors.

Performance. We demonstrate the optimized TTM embedding tables on a single RTX 3090 GPU. We consider an embedding table of TTM shape [[80 , 50 , 54 , 50] , [4 , 4 , 4 , 2]] and rank 32 , extracted from a practical DLRM model. As shown in Figure 5, our proposed method achieves about 4 -5 × speed-up and 2 -3 × memory saving than the standard TTM embedding without any optimization. The uncompressed embedding with sparse gradients is faster than our approach since our TTM embedding table requires extra computation, but it uses much more memory than the TTM embedding table.

4.2 Contraction Path Optimization for TT-Vector Multiplications

Next, we optimize the forward- and back- propagation of linear layers in the TT format. We consider the linear layer Y = XW , where Y ∈ R b × N 2 , W ∈ R N 1 × N 2 , X ∈ R b × N 1 . The W is compressed into the Tensor-Train format: W = [ [ G 1 , . . . , G 2 d ] ] ∈ R n 1 ×···× n 2 d ,where G i ∈ R r i -1 × n i × r i and N 1 = n 1 · · · n d , N 2 = n d +1 · · · n 2 d . The forward-propagation in the einsum form is

where X ∈ R b × n 1 ×···× n d is the reshaping of X and S i denotes r i -1 n i r i . Suppose that the gradient to Y is g Y , then the gradients to G i and X can be computed as follows:

In total, 2 d +2 contraction sequences are needed for the TT-format forward- and back- propagation. To reduce the computational costs, it is critical to find an optimal or near-optimal contraction path.

Large batch case. We denote A i := G 1 ×··· × G i , A -i = G d -i +1 ×··· × G d , B i := G d +1 × · · · × G d + i , B -i = G 2 d -i +1 ×··· × G 2 d , which are all computed sequentially. In practice, we only need to compute A d , A -d , B d , B -d and store the intermediate results. The forward-propagation (19) is then computed in the following way

In backward propagation, the gradients are computed in the following way:

The gradient g X is computed as

The gradients g G i for i ≥ d +1 can be computed as

Similarly, the gradients g G i for i ≤ d can be computed as

The contraction paths of forward- and back- propagation are summarized in Appendix A.5.

Analysis. The proposed empirical path is near-optimal for large batch sizes. The following result analyzes the contraction path for forward-propagation.

Proposition 4.1. Suppose that the TT ranks satisfy 1 = r 0 < r 1 ≤ · · · ≤ r d ≥ r d -1 > · · · ≥ r 2 d = 1 and the batch size b is large enough. There exist groups { S i } k i =1 where S i = { G j i +1 , . . . , G j i +1 } containing consecutive tensor cores for 0 = j 1 < · · · < j k < j k +1 = 2 d . Then, the contraction path with the least number of flops for the forward-propagation (19) first contracts the tensor cores in each S i to obtain V i with dimension r j i × n j i +1 ×··· × n j i +1 × r j i +1 and then contract the input tensor X with tensors { V i } k i =1 in the sequential order.

Proposition 4.1 implies that the optimal path first contracts some consecutive tensor cores and then contracts obtained tensors with the input tensor sequentially. The groups { S i } k i =1 depend on the dimensions, ranks, and batch size. The proposed contraction path satisfies the property shown in Proposition 4.1 and has flops roughly b ( n 1 · · · n d + n d +1 · · · n 2 d ) r d . The optimal contraction path has flops about bn 1 · · · n d c 1 + bn d +1 · · · n 2 d c 2 , where c 1 , c 2 are some constants. Hence, the proposed is near-optimal and has a comparable complexity to the optimal path. Suppose the optimal path is different from the proposed empirical path. Then the optimal path will likely involve a few more large intermediate tensors, which pose more memory costs during training and inference especially for static computational graphs. The empirical path is a good choice to balance time and memory consumption. Similar arguments can be applied to the contractions for back-propagation.

When the batch size is small, the optimal path may have much fewer flops. However, the execution time is almost the same as the proposed path since all the operations are too small. Hence, we can use the proposed path for most batch sizes. See Appendix A.6 for more analysis.

4.3 GPU Performance Optimization via CUDA Graph

While CoMERA consumes much less computing FLOPS than standard uncompressed training, it can be slower on GPU if not implemented carefully. Therefore, it is crucial to optimize the GPU performance to achieve real speedup. Modern GPUs are highly optimized for large-size matrix multiplications. However, the small-size tensor contractions in CoMERA are not yet optimized on GPU and require many small-size GPU kernels, causing significant runtime overhead. During the training, Cuda Graph launches and executes the whole computing graph rather than launching a large number of kernels sequentially. This can eliminate lots of back-end overhead and lead to significant training speedup. It is more suitable for CoMERA since tensor-compressed training has much more small kernels than uncompressed training. This is just an initial step of GPU optimization. We expect that a more dedicated GPU optimization can achieve a more significant training speedup.

5 Training Results

In this section, we test the performance of CoMERA on a few benchmarks, including a domainspecific LLM. Our experiments are run on a Nvidia RTX 3090 GPU with 24GB RAM.

5.1 A Medium-Size Transformer with Six Encoders

We first consider a six-encoder transformer. The embedding tables and all linear layers are represented as tensor cores in the training process as detailed in Appendix A.7. We train this model on the MNLI dataset [38] with the maximum sequence length 128 and compare the accuracy, resulting model size, and training time of CoMERA with the standard uncompressed training.

0.64

accuracy

0.62

0.60

1on

0.58

0.56

uncompressed training

CoMERA

0.54

0

2

4

6

8

10

12

14

16

18

20

mance. The early-stage training of CoMERA achieves 74 × compression ratio on all tensorized layers, and the validation accuracy is even higher than the uncompressed training. Figure 6 shows the validation accuracy of CoMERA. In the late stage of CoMERA, we set different target compression ratios for more aggressive rank pruning. The target compression ratios are for the tensorized layers rather than for the whole model. The late-stage training can reach the desired compression ratio with very little accuracy drop. The smallest model has a compression ratio of 80 × for the whole model due to a 361 × compression on the tensorized layers with slightly worse accuracy.

Architecture Search Capability of CoMERA.

A major challenge in training is architecture search: shall we keep certain layers of a model? Interestingly, CoMERA has some capability of automatic architecture search. Specifically, the ranks of some layers become zero in the training, and thus the whole layer can be removed. For the target compression ratio 0 . 2 , the whole second last encoder and some linear layers in other encoders are completely removed after late-stage rank-adaptive training. The change of ranks of layers in the 5th encoder is shown in Table 2.

30

CoMERA w/CUDAGraph

uncompressed w/ CUDAGraph

25

CoMERA w/o CUDAGraph

uncompressed w/o CUDAGraph

(min)

20

18.5

16.6

16.4

15.3

time

15

14.3

14.7

14.7

10

7.2

6.4

8.2

5.5

7.4

5

0

32

64

128

batchsize

Training Time. As shown in Figure 7, CoMERA with CUDAGraph achieves around 3 × speed-up than uncompressed training. CoMERA without CUDAGraph can take much longer time in small batch-size setting due to the launching overhead of too many small kernels. The uncompressed training with CUDAGraph takes longer time than the one without CUDAGraph. This is because CUDAGraph requires all batches to have the same sequence length, and the consequent computing overhead is more than the time reduction of CUDAGraph. In contrary, CoMERA has much fewer computing FLOPS and the computing part accounts for a much smaller portion of the overall

(a) Training time per epoch.

(b) Memory cost.

CoMERA w/ optimization

CoMERA w/o optimization

1500

1344

1182

1000

time

807

794

791

500

420

423

424

N/A

10000

20000

40000

batch size

25000

CoMERA w/ optimization

CoMERA w/o optimization

uncompressed

20459

(MB

20000

18261

1838519005

15000

memory

10000

9259

6629

5000

2612

3947

N/A

0

10000

20000

40000

batchsize

runtime. Empirically CoMERA is 2 -3 × faster in the whole training than uncompressed training for transformers on a single GPU, but we do not have theoretical guarantees about the number of epochs although they are similar in our experiments. Appendix A.8 provides more details about the run-time comparison on this benchmark, showing that CoMERA is still faster than standard training even if the compression ratio is close to 1.

5.2 A DLRM Model with 4-GB Model Size

We further test CoMERA on DLRM [32] released by Meta on Criteo Ad Kaggle dataset [23]. We compress the ten largest embedding tables into the TTM format as in Section 4.1. All fully connected layers with sizes > 128 are compressed into TT format. The model is trained for two epochs.

Effect of Optimized TTM Embedding. The training time per epoch and peak memory cost are shown in Figure 9. Our optimized TTM lookup speeds up the training process by around 2 × and remarkably reduces the memory cost by 4 -5 × .

Overall Performance of CoMERA. Table 3 shows the testing accuracy, testing loss (measured as normalized CE), memory costs, and model sizes of CoMERA and uncompressed training. CoMERA achieves similar accuracy as the uncompressed training, while CoMERA compresses the whole model by 99 × and saves 7 × peak

uncompressed

CE

COMERA

0.90

normalzied

0.85

0.80

2000

4000

6000

8000

iteration

memory cost (with consideration of the data and backend overhead) in the training process. The reduction of model size and memory cost mainly comes from the compact TTM tensor representation of large embedding tables. Standard uncompressed training is faster than CoMERA since DLRM is an embedding-intensive model, and the computation in the embedding table is look-up rather than matrix multiplications. However, CoMERA uses much less memory, saving 6.9X, 4.8X, and 3.1X memory for batch sizes 10000, 2000, and 4000, respectively. Furthermore, CoMERA has a similar convergence curve and needs fewer iterations than standard training for DLRM, as shown in Figure 8.

5.3 Comparison with GaLore and LTE

We compare our method with two recent low-rank compressed training frameworks: GaLore [41] and LTE [20]. GaLore [41] reduces the memory cost by performing SVD compression on the gradient, and LTE represents the weights as the sum of parallel low-rank matrix factorizations. We evaluate their memory costs and training times per

epoch on the six-encoder transformer model under different batch sizes. We do not compare the total training time because the training epochs of various methods are highly case-dependent. The CoMERA and GaLore achieve almost the same validation accuracy, 64%, on the MNLI dataset. However, the LTE approach does not converge on the task using its default setting.

Training Time Per Epoch. We use rank 128 for the low-rank gradients in GoLore, and rank 32 and head number 16 for the low-rank adapters in LTE. For a fair comparison, all methods are executed with CUDA graph to reduce the overhead of launching CUDA kernels. The runtimes per training epochs are reported in Figure 1(a). For the LTE, we only report the results for batch sizes 32 , 64 , 128 since it requires the batch size to be a multiple of the head number. Overall, our CoMERA is about 2 × faster than GaLore and 3 × faster than LTE for all batch sizes, because the forward and backward propagation using low-rank tensor-network contractions dramatically reduce the computing FLOPS.

12

CoMERA

10

CodeBERT

Loss

8

Phase 1

raining

6

Phase2

4

2

0

0

1000

2000

3000

4000

5000

6000

7000

8000

Steps

Memory Cost. Figure 1 (b) shows the memory cost of all three training methods. In the singlebatch setting as used in [41], our CoMERA method is 9 × more memory-efficient than Galore on the tested case (with consideration of data and back-end cost). As the batch size increases, the memory overhead caused by data and activation functions becomes more significant, leading to less memory reduction ratios. However, our proposed CoMERA still uses the least memory.

We run the experiments on the RTX 3090 GPU. The work GaLore[41] uses the RTX 4090 GPU for experiments, so we also compare them on the RTX 3090 GPU. The results are in Appendix A.9.

5.4 Preliminary LLM Pre-Training Results: Case Study on CodeBERT

To show the benefit of CoMERA in pre-training (domain-specific) LLMs, we follow the setup from CodeBERT [10] to pre-train a BERT-like model for code generation. The pre-training dataset is the CodeSearchNet [21], a collection of 2M (comment, code) pairs and 6M pure code sequences from open-source libraries with 6 types of programming languages. We pre-train CodeBERTLARGE ( 357 M) and its CoMERA ( 84 M) variant using the masked language modeling (MLM) objective and compare their training loss in Figure 10. We achieve up to 12 . 72 × compression on tensorized layers and 4 . 23 × overall compression with final loss of 0 . 40 vs 0 . 28 . There is a small gap between the final losses. However, this does not necessarily imply performance degradation on downstream tasks, based on our observation on BERTLARGE[9] shown in Appendix A.10. Furthermore, CoMERA is 2 . 3 × and 1 . 9 × faster than standard pre-training in Phase 1 and Phase 2 respectively, when evaluated on the Nvidia RTX 3090 GPU. Our current CoMERA implementation is still slower than standard pre-training on massive GPUs, since no performance optimization has been done on HPC.

6 Conclusions and Future work

This work has presented CoMERA framework to reduce the memory and computing time of training AI models. We have investigated rank-adaptive training via multi-objective optimization to meet specific model sizes while maintaining model performance. We have achieved real training speedup on GPU via three optimizations: optimizing the tensorized embedding tables, optimizing the contraction path in tensorized forward and backward propagation, and optimizing the GPU latency via CUDAGraph. The experiments on a transformer model demonstrated that CoMERA can achieve 2 -3 × speedup per training epoch. The model sizes of the transformer and a DLRM model have been reduced by 43 × to 99 × in the training process, leading to significant peak memory reduction (e.g., 7 × total reduction in large-batch training of DLRM on a single GPU). Our method has also outperformed the latest GaLore and LTE frameworks in both memory and runtime efficiency. More importantly, our method has demonstrated significant speedup and model compression in pre-training CodeBERT, a domain-specific LLM for automatic code generation. We have also observed further speedup by combining CoMERA with mixed-precision computation. The discussions and some preliminary results are in Appendix A.11.

Unlike large-size matrix operations, the small low-rank tensor operations used in CoMERA are not yet well-supported by existing GPU kernels. The performance of CoMERA can be further boosted significantly after a comprehensive GPU and HPC optimization. The existing optimizers, e.g. Adam, are well-studied for uncompressed training. However, CoMERA has a very different optimization landscape due to the tensorized structure. Therefore, it is also worth studying the optimization algorithms specifically for CoMERA in the future.

References

AI and compute. https://openai.com/blog/ai-and-compute/ . OpenAI in 2019.

ChatGPT and generative AI are booming, but the costs can be extraordinary. https://www.cnbc.com/2023/03/13/ chatgpt-and-generative-ai-are-booming-but-at-a-very-expensive-price. html . Accessed: 2023-03-15.

Update: ChatGPT runs 10K Nvidia training GPUs with potential for thousands more. https://www.fierceelectronics.com/sensors/ chatgpt-runs-10k-nvidia-training-gpus-potential-thousands-more . Accessed: 2023-03-15.

Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165 , 2020.

Giuseppe G Calvi, Ahmad Moniri, Mahmoud Mahfouz, Qibin Zhao, and Danilo P Mandic. Compression and interpretability of deep neural networks via tucker tensor layer. arXiv:1903.06133 , 2019.

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.

Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. Compressing neural networks with the hashing trick. In International Conference on Machine Learning , pages 2285-2294, 2015.

Kalyanmoy Deb, Karthik Sindhya, and Jussi Hakanen. Multi-objective optimization. In Decision sciences , pages 161-200. CRC Press, 2016.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proc. Association for Computational Linguistics: Human Language Technologies , pages 4171-4186, 2019.

Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. Codebert: A pre-trained model for programming and natural languages. arXiv preprint arXiv:2002.08155 , 2020.

Jiaqi Gu, Ben Keller, Jean Kossaifi, Anima Anandkumar, Brucek Khailany, and David Z Pan. Heat: Hardware-efficient automatic tensor decomposition for transformer compression. arXiv preprint arXiv:2211.16749 , 2022.

Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In International Conference on Machine Learning , pages 1737-1746, 2015.

Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00149 , 2015.

Cole Hawkins, Xing Liu, and Zheng Zhang. Towards compact neural networks via end-to-end training: A bayesian tensor approach with automatic rank determination. SIAM Journal on Mathematics of Data Science , 4(1):46-71, 2022.

Cole Hawkins and Zheng Zhang. Bayesian tensorized neural networks with automatic rank selection. Neurocomputing , 453:172-180, 2021.

Qinyao He, He Wen, Shuchang Zhou, Yuxin Wu, Cong Yao, Xinyu Zhou, and Yuheng Zou. Effective quantization methods for recurrent neural networks. arXiv preprint arXiv:1611.10176 , 2016.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2015.

Oleksii Hrinchuk, Valentin Khrulkov, Leyla Mirvakhabova, Elena Orlova, and Ivan Oseledets. Tensorized embedding layers for efficient model compression. arXiv preprint arXiv:1901.10787 , 2019.

Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio. Quantized neural networks: Training neural networks with low precision weights and activations. The Journal of Machine Learning Research , 18(1):6869-6898, 2017.

Minyoung Huh, Brian Cheung, Jeremy Bernstein, Phillip Isola, and Pulkit Agrawal. Training neural networks from scratch with parallel low-rank adapters. arXiv preprint arXiv:2402.16828 , 2024.

Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Marc Brockschmidt. CodeSearchNet challenge: Evaluating the state of semantic code search. arXiv preprint arXiv:1909.09436 , 2019.

Max Jaderberg, Andrea Vedaldi, and Andrew Zisserman. Speeding up convolutional neural networks with low rank expansions. arXiv preprint arXiv:1405.3866 , 2014.

Olivier Chapelle Jean-Baptiste Tien, joycenv. Display advertising challenge, 2014.

Valentin Khrulkov, Oleksii Hrinchuk, Leyla Mirvakhabova, and Ivan Oseledets. Tensorized embedding layers for efficient model compression. arXiv preprint arXiv:1901.10787 , 2019.

Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Compression of deep convolutional neural networks for fast and low power mobile applications. arXiv preprint arXiv:1511.06530 , 2015.

Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Review , 51(3):455-500, Aug. 2009.

Urs Köster, Tristan Webb, Xin Wang, Marcel Nassar, Arjun K Bansal, William Constable, Oguz Elibol, Scott Gray, Stewart Hall, Luke Hornof, et al. Flexpoint: An adaptive numerical format for efficient training of deep neural networks. In NIPS , pages 1742-1752, 2017.

Joseph M Landsberg. Tensors: geometry and applications. Representation theory , 381(402):3, 2012.

Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition. arXiv preprint arXiv:1412.6553 , 2014.

Jian-Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. In Proceedings of the IEEE international conference on computer vision , pages 5058-5066, 2017.

Pavlo Molchanov, Arun Mallya, Stephen Tyree, Iuri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11264-11272, 2019.

Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.00091 , 2019.

Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. In Advances in neural information processing systems , pages 442-450, 2015.

Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. In Advances in Neural Information Processing Systems 28 , pages 442-450, 2015.

Ivan V Oseledets. Tensor-train decomposition. SIAM Journal on Scientific Computing , 33(5):2295-2317, 2011.

Xiao Sun, Naigang Wang, Chia-Yu Chen, Jiamin Ni, Ankur Agrawal, Xiaodong Cui, Swagath Venkataramani, Kaoutar El Maghraoui, Vijayalakshmi Viji Srinivasan, and Kailash Gopalakrishnan. Ultra-low precision 4-bit training of deep neural networks. NIPS , 33, 2020.

Andros Tjandra, Sakriani Sakti, and Satoshi Nakamura. Compressing recurrent neural network with tensor train. CoRR , abs/1705.08052, 2017.

Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) , pages 1112-1122. Association for Computational Linguistics, 2018.

Yifan Yang, Jiajun Zhou, Ngai Wong, and Zheng Zhang. LoRETTA: Low-rank economic tensor-train adaptation for ultra-low-parameter fine-tuning of large language models. arXiv preprint arXiv:2402.11417 , 2024.

Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-train recurrent neural networks for video classification. In Proc. the 34th International Conference on Machine Learning , volume 70, pages 3891-3900, Sydney, Australia, 06-11 Aug 2017.

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient LLM training by gradient low-rank projection, 2024.

Yequan Zhao, Xinling Yu, Zhixiong Chen, Ziyue Liu, Sijia Liu, and Zheng Zhang. Tensorcompressed back-propagation-free training for (physics-informed) neural networks. arXiv preprint arXiv:2308.09858 , 2023.

Yu-Bang Zheng, Xi-Le Zhao, Junhua Zeng, Chao Li, Qibin Zhao, Heng-Chao Li, and Ting-Zhu Huang. Svdinstn: A tensor network paradigm for efficient structure search from regularized modeling perspective. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 26254-26263, 2024.

A Supplementary Material

A.1 Comparison with existing works

SVDinsTN [43] uses sparse diagonal matrices and the ℓ 1 regularization to control tensor ranks for a compact tensor structure of a given tensor. In contrast, our work compresses weights during end-to-end training without any prior information on the tensor (i.e., model parameters). Both works use sparse diagonal matrices and ℓ 1 terms to control tensor ranks. Using diagonal matrices to control the ranks of matrices is very common, like SVD. The ℓ 1 norm is also widely used to induce sparsity in various models, like compressed sensing and Lasso regression. It is natural to combine these two techniques to control tensor ranks, regardless of tensor formats. Moreover, our work formulates the problem as a more generic multi-objective problem and uses a two-stage algorithm to solve it. The formulation in SVDinsTN is similar to linear scalarization approach in our early stage. Our work further uses the achievement scalarization in the late stage to find a model close to our preferred model performance and size.

HEAT [11] also considers contraction optimization for post-training model compression of trained models. In contrast, CoMERA considers end-to-end tensor-compressed training, where no model parameters are known prior to training. In addition, HEAT only discusses the single path optimization for forward propagation in CP format. We have optimized d +2 contraction paths jointly in both forward- and back- propagation in TT format. Since these contractions can be coupled, we have also minimized the overall computation costs by reusing intermediate results.

A.2 Proof of Proposition 3.1

Proof. The objective function in (9) is bounded below by 0 . Hence, the problem (9) has a finite infimum value f ∗ . Let { G k , D k } ∞ k =1 be a sequence such that lim k →∞ L ( G k , D k ) + γ ˆ S ( D k ) + β ∥ G k ∥ 2 = f ∗ . The sequence must be bounded because of the l 1 regularization of D and the ℓ 2 regularization of G . As a result, the sequence has a cluster point ( G ∗ , D ∗ ) which is a minimizer of the (9). Let C := ∥ G ∗ ∥ 2 . The relaxation (9) is equivalent to the constrained optimization problem.

It implies that the solution to the training problem (9) is a Pareto point of the multi-objective optimization problem min G , D ( L ( G , D ) , ˆ S ( D )) .

A.3 Algorithm for Late Stage Optimization in Section 3.2

The algorithm for the late stage optimization in Section 3.2 is summarized in Algorithm 1.

A.4 Proof of Proposition 4.1

Proof. For convenience, let V i be a string of characters to specify the dimension of V i , C i be the set of tensor cores used to obtain V i , and X i be the tensor by contracting X with V 1 , . . . . V i , denoted by the string X i .

̸

We first show that V 1 must be in the proposed format. Suppose otherwise for contradiction. Let G i be the first tensor core used to obtain V 1 . If i = 1 , then we write V 1 = V 1 1 V 2 1 where the tensor V 1 1 corresponding to V 1 1 is obtained by contractions of longest consecutive tensor cores containing G i in the set C 1 . Let Z 1 = V 1 i ∩ ˜ X , Z 2 = V 2 i ∩ ˜ X . The number of flops for the contraction between ˜ X and V 1 is π ( X ) π ( V 1 1 ) π ( V 2 1 ) π ( Z 1 ) π ( Z 2 ) . If we first contract ˜ X with V 1 1 and then contract the obtained tensor with V 2 1 , the number of flops is π ( X ) π ( V 1 1 ) π ( Z 1 ) + π ( X ) π ( V 1 1 ) π ( V 2 1 ) π ( Z 1 ) 2 π ( Z 2 ) which is less than π ( X ) π ( V 1 1 ) π ( V 2 1 ) π ( Z 1 ) π ( Z 2 ) since Z 2 = V 2 1 . It contradicts our assumption that this is the optimal path. If d ≥ i > 1 , let S be the tensor generated by the longest consecutive tensor cores containing G i -1 and used in the optimal path. A better path is to first contract V 1 with S to obtain W , then contract W with X and all other unused parts in the optimal path. It is better because the number of flops for contracting W and X is no greater than that for contracting V 1 and X and the new path reduces the number of flops in the remaining contractions. The reduction in the remaining contractions is more than the potential flop increase in obtaining W when the batch size b is big enough. Finally, we consider the case that i > d . Let G j be the first tensor core used to obtain V 2 . If j > d , then first contracting V 1 and V 2 and then all other parts is a better choice. Otherwise, if j ≤ d , let V 1 2 be the tensor contracted by the consecutive tensors containing G j in C 2 . The tensor V can be represented by the contraction of V 1 2 and another tensor V 2 2 generated by the remaining tensor cores in C 2 . In this scenario, we can contract V 1 2 with V 1 to get S , then S with X to get W , then V 2 2 with W , and finally the obtained tensor with V 3 , . . . , V k . It is not hard to verify contracting V 1 2 and V 1 , S and S and X , and W and V 2 2 uses less flops than directly contracting X with V 1 and V 2 directly when the batch size b is large. Summarizing everything above, we can conclude that V 1 must be in the proposed format.

The contraction of X i and V i +1 , . . . , V k has the similar structure to the contraction of X and V 1 , . . . , V k . By applying the same proof, we conclude that the tensors V i 's must be in the format stated in the proposition and we will contract the input tensor X with the tensors tenV 1 , . . . , V k in the sequential order.

A.5 Algorithm for Contraction Path in Section 4.2

The empirical near-optimal contraction path for tensor-compressed training is shown in Algorithm 2.

Figure 11 presents the tensor diagrams for contraction paths of TT forward- and back- propagation as discussed in Section 4.2.

Algorithm 2 Empirical path for tensor-compressed forward- and back- propagation

Forward Input: Tensor cores G 1 , . . . , G 2 d and input matrix X Forward Output: Output matrix Y , and intermediate results.

1: Reshape the matrix X to the tensor X .

2: Compute A d , A -d , B d , B -d in the sequential order and store intermediate results of { A i , A -i , B i , B -i } d i =1 for back-propagation.

3: Compute T 1 as in (22) to store it for back-propagation.

4: Compute Y as in (23) and reshape it to the appropriate matrix Y .

Backward Input: Inputs of Forward , stored results from Forward , and output gradient g Y . Backward Output: Gradients g , g , . . . , g .

X G 1 G 2 d

1: Reshape the gradient g Y to the tensor g Y .

2: Compute U 1 and g X as in (24), (25) and store U 1 for future use.

3: Compute g G i for i ≥ d +1 as in (26), (27) using stored tensors.

4: Compute g G i for i ≤ d as in (22), (29) using stored tensors.

A.6 Small Batch Case for Contraction Path in Section 4.2

Small batch case. The empirical contraction path in Algorithm 2 eliminates the batch size dimension b early, so it is nearly optimal when the batch size is large. We may search for a better path using a greedy search algorithm to minimize the total operations. In each iteration, we prioritize the pairs that output the smallest tensors. Such a choice can quickly eliminate large intermediate dimensions to reduce the total number of operations. When the batch size is large, the searched path is almost identical to the empirical path in Algorithm 2 which eliminates the batch size dimension b

.

(a) TT-vector forward-propagation

(b) TT-vector back-propagation: compute

g

X

(c) TT-vector back-propagation: compute

g

G

i

,

1

i

d

Pu

I+Pu

nd+1...n2d

P6

I+P6

92d

nd+1

n2d

nd+1

6

gy

L2

Pu

d+1

n2d

r2d-

yd

I+P6

G2d

y

n1.

nd+1'

gx

gy

Pu

nd+1

n2d

+

r'd

rd+1

r2d

Gi+1

Gd+1

g2d

n

5%

gy

nd+1

ni+1

Td

rd+1

r2d-

Gd+1

Gi+1

G2d

mi

nd+1

n2d

.n2d

5%

B

B_(2d-i)

T'i

ri

G

(d) TT-vector back-propagation: compute

g

i

, d

+1

i

2

d

early. The searched path may differ from Algorithm 2 for small batch sizes, but their execution times on GPU are almost the same. This is because the tensor contractions for smaller batch sizes have a minor impact on the GPU running times. Consequently, despite certain tensor contractions in the empirical path being larger than those in the optimal path, the actual GPU execution times between them exhibit only negligible differences. Therefore, the empirical contraction path in Algorithm 2 is adopted for all batch sizes in CoMERA.

A.7 Compression Settings for the Experiment in Section 5.1

The compression settings for the experiment in Section 5.1 are shown in Table 4.

25

50× compression

2.2× compression

1.1× compression

4.9× compression

1.5× compression

uncompressed

20

18.5

17.46

16.6

16.4

n

15.53

15

13.94

14.35

12.44

11.16

11.56

m

e

10

9.7

9.07

7.2

7.79

7.79

6.4

5.5

6.21

5

32

64

128

batch size

A.8 Per Epoch Training Time of CoMERA on MNLI for Various Compression Ratios

Table 12 shows the per-epoch training time of CoMERA on MNLI dataset for different compression ratios. The acceleration is more obvious for larger compression ratios. When the compression ratio is greater than 1, CoMERA always has speedup. When the compression ratio approaches 1, the time of CoMERA approaches that of uncompressed training.

A.9 Comparison with GaLore and LTE on a single RTX 4090 GPU

Since GaLore[41] uses a single RTX 4090 GPU for experiments in the original paper, we also run the experiments on the RTX 3090 GPU and compare the results. Figure 13 presents the training time and peak memory consumption. Compared to RTX 3090, the training on RTX 4090 uses similar memory and takes less training time, and CoMERA is still the fastest method and consumes the least memory among all three techniques. The memory savings are almost the same as the results reported in Figure 1 in our paper. The speed-up factors are almost identical for batch sizes 32, 64, and 128. For batch size 1, our method is 1 . 2 × faster and 1 . 7 × faster than GaLore on RTX 3090, respectively. The difference is that RTX 4090 GPU significantly accelerates matrix multiplications of batch size 1, while it does not accelerate that much for smaller tensor contractions. We find that r = 30 matrix multiplication on RTX 3090 has a similar speedup for both batch sizes, whereas the same multiplication on RTX 4090 only has speedup for batch 32 and does not have any speedup for batch 1. We would like to note that it might be caused by that different GPU platforms have different backend overhead, which can become more dominant as computation decreases to batch=1. We will continue optimizing GPU-level kernels to accelerate small tensor contractions and expect to see a similar speedup.

A.10 CoMERA Pretraining Result on BERTLARGE

CoMERA on Original BERTLARGE. Our results on CodeBERT is rather preliminary as only pre-training loss is available. For the original BERTLARGE[9] (336M) and its CoMERA (125M) variant that we trained by using Wikipedia (2500M words), we achieve up to 6 . 36 × compression on tensorized layers and 2 . 69 × overall compression, with final loss of 1.45 vs 1.26. On downstream tasks, CoMERA ourperforms BERTLARGE on SST-2 (accuracy: 92.10% vs 91.74%) and MRPC (accuracy: 86.82% vs 86.00%), underperforms BERTLARGE on SQuAD (f1: 88.76% vs 90.68%).

(a) Training time per epoch.

(b) Peak memory consumption.

60

CoMERA

GaLore

LTE

CoMERA

GaLore

LTE

10964

44.8

B

10000

(min

40

37.1

8048

memory

6628

7002

time

20

5000

3632

4636

3784

4682

N/A

3.4

6.8

11.1

3.4

6.3

9.1

3.4

6.0

8.2

182

1674

N/A

1780

32

64

128

32

64

128

batch size

batch size

A.11 Discussion: Mixed-Precision CoMERA

Modern GPUs offer low-precision computation to speed up the training and inference. It is natural to combine low-rank tensor compression and quantization to achieve the best training efficiency. However, CoMERA involves many small-size low-rank tensor contractions, and a naive low-precision implementation may even slow down the training due to the overhead caused by precision conversions.

To resolve the above issue, we implement mixed-precision computation in CoMERA based on one simple observation: large-size contractions enjoy much more benefits of low-precision computation than small-size ones. This is because the overhead caused by precision conversions can dominate the runtime in small-size contractions. In large-batch tensor-compressed training, small- and large-size tensor contractions can be distinguished by whether the batch size dimension b is involved. In general, a contraction with the batch b is regarded as large and is computed in a low precision. Otherwise, it is regarded small and is computed in full-precision. The actual mixed-precision algorithm depends on the contraction path used in the forward- and back- propagations of CoMERA.

Runtime. We evaluate the mixedprecision forward and backward propagations of CoMERA in a FP8 precision on the NVIDIA L4 GPU. We consider a single linear layer. The shapes (1024 , 1024) and (1024 , 4096) are converted to the TT shapes (16 , 8 , 8 , 8 , 8 , 16) and (16 , 8 , 8 , 16 , 16 , 16) respectively, and the ranks are both 32 . The total execution time for 1000 forward and backward propagations are shown in Table 5. The FP8 tensor-compressed linear layer has about 3 × speed-up compared to the FP8 vanilla linear layer when the batch size

alidation accuracy

0.625

0.600

0.575

Mixed-precision CoMERA

val

0.550

FP32CoMERA

0

2

4

6

8

10

12

14

16

18

20

and layer size are large. When the batch size is small, the FP8 vanilla linear layer is even faster. This is because the tensor-compressed linear layer consists of a few sequential computations that are not well supported by current GPU kernels. We expect to see a more significant acceleration after optimizing the GPU kernels.

Convergence. We use the mixed-precision CoMERA to train the DLRM model and the six-encoder transformer. The result on DLRM is shown in Table 6. The convergence curve of the six-encoder transformer is shown in Figure 14. The experiments demonstrate that the accuracy of FP8 training is

similar to FP32 training. However, we did not see much acceleration of using FP8 in the experiments. This is mainly because of 1 ⃝ the computation overhead of slow data type casting between FP32 and FP8; 2 ⃝ the sequential execution of small tensor contractions that are not well supported by current GPUs; 3 ⃝ the relatively small sizes of linear layers in the tested models. We will investigate these problems in the future, and are optimistic to see significant acceleration on larger models.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction accurately reflect the methods in Section 3 and Section 4 and the numerical results in Section 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We mentioned that current GPU-level optimization is not optimal yet and there still exist lots of opportunities for improvement.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The complete assumptions and proofs are provided in the Appendix, e.g. Appendix A.2 and Appendix A.6.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper clearly describes the proposed methods in Section 3 and Section 4. The experimental settings are provided in Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We used the open-sourced data for experiments. An implementation of CoMERA is available at https://github.com/ziyangjoy/CoMERA .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental settings are provided in Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The paper focuses on the computation-efficiency training. We mainly report the resource usage and training results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computer resources are provided in Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our paper is about new computational algorithms for efficient training of large models. It conforms with the NeurIPS code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Our paper proposes new methods for efficient training of large models. The research has the possibility to significantly reduce the training costs of large models and also accelerate the inference. It is discussed in Section 1 and Section 6.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper is about efficient training methods and poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Yes, they are properly credited, mentioned, and respected.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper is about methods and algorithms for efficient training. No new assets are released.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper is about methods and algorithms for efficient training. It does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: he paper is about methods and algorithms for efficient training. It does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces CoMERA, a novel training method for large AI models that focuses on optimizing both computing and memory efficiency through rank-adaptive tensor compression. CoMERA aims to reduce training costs and environmental impact by achieving high compression ratios and maintaining accuracy. Key contributions include a multi-objective optimization formulation, performance optimization techniques for tensor-compressed training on GPUs, and empirical validation showing significant speedup and memory savings compared to existing methods.\n\nSTRENGTHS:\n1. The experimental results demonstrate significant improvements in training speed and memory efficiency, outperforming recent methods like GaLore and LTE.\n2. By addressing both computing and environmental costs, the method has practical implications for large-scale AI model training, making tensor models more practically useful in machine learning.\n3. In the part of multi-objective optimization, the formulation balances compression ratio and model accuracy, providing a customizable framework for different resource requirements.\n\nWEAKNESSES:\nWhile the method shows impressive results on tested models, scalability to even larger models and diverse architectures remains an area for further exploration.\n\nQUESTIONS:\n1. For learning TT-ranks, this work imposes diagonal matrices  D  (shown in Eq. 6) with  l_1  norm regularization. This technique for tensor network structure search was also utilized in the recent paper:\n\nZheng, Yu-Bang, et al. “SVDinsTN: A Tensor Network Paradigm for Efficient Structure Search from Regularized Modeling Perspective.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n\n2. In line 122, could you provide more explanation of what the linear scalarization means? Does it refer solely to Eq. (7), or is there additional context?\n\n3. Could you provide more intuition regarding Eq. (11), i.e., the achievement scalarization? While the two scalarizations in the paper may be standard methods in multi-objective optimization, a more intuitive explanation would help general readers unfamiliar with this field grasp the main idea quickly.\n\n4. In lines 171-173, the design of the “tensor-index level” is unclear. Please improve the clarity of this part if possible.\n\n5. In the contraction section, how many core tensors are involved in the contraction? If the number of involved tensors in the contraction is not large, why not use the default contraction path searching algorithms integrated into einsum, e.g., ‘dp’ or ‘optimal’? What benefits can be achieved from handcrafting a new contraction path for the current model compression task?\n\n6. Could you provide more explanation in Sec. 4.3? Why can CUDA graphs improve tensor computation? What specific efforts were made in this work? Is using CUDA graphs straightforward, or does it require advanced programming techniques, such as designing a more GPU-suitable contraction order?\n\n7. In the comparison with GaLore, why was a 3090 used instead of a 4090, which was used in the original GaLore paper?\n\n8. Regarding the training time, only the time per epoch is given. Is it possible to provide the overall running time? Are more epochs required if tensor networks are utilized?\n\n9. Will the codes be released once the paper is accepted?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscript presents techniques for efficient training from scratch based on tensor decompositions. The authors propose several modifications to the basic training approach to improve accuracy as well as several optimizations for tensor-compressed training, achieving training speedup. In the experiments, the proposed methodology is used to compress a transformer model and a recommender system model.\n\nSTRENGTHS:\nThe authors present a new approach for training a model with tenzorized weights, which allows for adaptation of ranks during training. They also managed to optimize the process and achieve a real-time speed-up, which is not an easy task when working with tensorized models. The described methods are useful and have the potential to significantly contribute to further advancements in efficient operation for tensor decomposition formats.\n\nWEAKNESSES:\n1) There are several places where a comparison with the baseline time (e.g., uncompressed) would be appropriate, such as in Figures 5 and 8.\n\n2) Experiments with the transformers are limited. Training a model on MNLI from scratch may not accurately reflect the behaviour of the method for language models,  as MNLI is a relatively simple classification dataset that might not require a large number of parameters. It would be beneficial if the paper included an experiment or a simulation of some kind in a pre-training setting for language models, such as pretraining Roberta on C4\n\nQUESTIONS:\n1) I am a little confused about how to adapt the rank to achieve further compression. It is not clear from the text what the procedure is and when it is applied during training. Specifically, it is not explained how the ranks are chosen and at what point of time the procedure is applied. I would recommend incorporating a new paragraph into the paper that provides a more detailed explanation of the procedure, as well as some practical advice on how to apply it in various scenarios. This could help readers better understand the potential uses and benefits of this approach.\n\n2) As for the transformer model, I am also interested about the model upper compression limit that still maintains training speedup? I would imagine maintaining good performance can only be possible with lower compression ratios (say x2-5).\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes using low-rank tensor train decomposition to accelerate deep learning model training and save memory usage. In the algorithm, both embedding tables in recommendation systems and large linear weights are written as a tensor train, and rank-adaptive optimization is used to adaptively reduce the rank without sacrificing the model accuracy. Multiple techniques, including TT embedding lookup, contraction path optimization, and CUDA graphs are combined to improve the efficiency. Experimental results show good memory saving and speedup without sacrificing model accuracy.\n\nSTRENGTHS:\noriginality: the idea to adaptively change the TT rank during training looks new. In addition to the training algorithm, multiple performance optimization techniques make the paper solid. In particular, the TT embedding lookup algorithm that does sampling and contraction in a interleaved way is also useful.\n\nWEAKNESSES:\npresentation: section 4.2 is hard to understand. For contraction path optimization, it would be good to visualize the process using tensor diagrams.\n\nlimitation of the proposed algorithm: it would be good to discuss the limitation of the algorithm in detail. In particular, I believe low-rank weights can help only under relatively small datasets/tasks, under which the low-rankness is a good regularization. Under large dataset/foundation model setting, I suspect the work can beat the baseline algorithms. In particular, the algorithm is also different from previous algorithms that use low-rank. In previous works, low-rank approximation is only applied on gradients rather than weights. To claim that the work can help foundation model training, we need larger datasets.\n\nQUESTIONS:\n1. line 42: replies -> relies\n\n2. How to choose a reasonable initial TT rank?\n\n3. how does accuracy compare between CoMERA, GaLore, and LTE?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper leverages the tensor decomposition concept in model training and tries to address two questions: the first is how to enable rank-adaptive, and the second one is how to obtain real efficiency. This paper achieves a 2 − 3× speedup per training epoch compared with standard training with on-par accuracy.

STRENGTHS:
* An easy paper to follow
* Although tensor decomposition is not new, they successfully enable real speed-up with several performance optimization techniques.
* Real speed up and on-par accuracy on small models.

WEAKNESSES:
* The problem formulation in 3.1 and how to solve the l0 norm is not quite new. 
* Contraction Path Optimization is also not a new topic. The authors should cite previous work, e.g. [1] and discuss differences, and highlight the uniqueness.
* Does this training take more time in terms of convergence? And could the author discuss more about the training difficulty of tensorized layers and normal layers.
[1] Gu, Jiaqi, et al. "Heat: Hardware-efficient automatic tensor decomposition for transformer compression." arXiv preprint arXiv:2211.16749 (2022).

QUESTIONS:
See weakness before.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13