PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zi Yang, Ziyue Liu, Samridhi Choudhary, Xinfeng Xie, Cao Gao, Siegfried Kunzmann, Zheng Zhang

arXiv:2405.14377v2cs.LGcs.AI

TL;DR

Training large AI models requires substantial hardware and time, and robust rank-adaptive tensor training with practical GPU speed remains challenging. CoMERA combines multi-objective rank-adaptive tensor compression with optimized tensor computations and GPU implementation. It reports faster training and reduced model size across transformer, recommendation, and preliminary code-generation pre-training evaluations.

  • Problem

    Training large models remains costly, while prior rank-adaptive tensor-compressed training relies on heuristic fixed-rank warm-up and tensor contractions can incur GPU overhead.

    Source: §1 · pp. 1–2

  • Method

    CoMERA jointly optimizes training loss and compressed model size through rank-adaptive tensor training, with optimized tensor computations and GPU execution.

    Source: §1 · p. 2

  • Results

    CoMERA achieves 2–3× faster epochs than standard training on the six-encoder transformer, with 80× transformer and 99× DLRM compression while maintaining standard-training accuracy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    On the tested six-encoder transformer, CoMERA is 2–3× faster per epoch than GaLore and LTE and uses 9× less memory than GaLore in single-batch training.

    Source: §1 · p. 3

    Small low-rank tensor operations are not yet well-supported by existing GPU kernels, and CoMERA’s tensorized optimization landscape warrants further study.

    Source: §6 · p. 10

Abstract

from arXiv · show

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 Computing- and Memory-Efficient training method via Rank-Adaptive tensor optimization. CoMERA achieves rank-adaptive 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\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.

1 Introduction

Training large models demands substantial hardware and time, while tensor compression has faced challenges in robust rank adaptation and practical GPU speed. CoMERA combines multi-objective rank-adaptive training with GPU optimizations to reduce memory and runtime.

Paper Contributions

  • 2–3× faster per epoch than GaLore and LTE, CoMERA also uses 9× less memory than GaLore in single-batch training on the tested six-encoder transformer.Figure 1 compares training time and total memory across varying batch sizes on an RTX 3090.Source: Figure 1 · p. 2, §1 · p. 3

Motivation and prior work

  • Prior rank-adaptive tensor training relies on heuristic fixed-rank warm-up, while small tensor contractions can incur GPU runtime overhead despite low theoretical FLOPS.The paper identifies robust rank adaptation and practical GPU speed as open challenges.Source: §1 · p. 2

Paper Contributions

  • CoMERA’s multi-objective formulation balances compression ratio and model accuracy while adapting the model to resource requirements.Rank-adaptive training can identify unnecessary layers and remove them.Source: §1 · p. 2
  • 80× and 99× compression on the six-encoder transformer and DLRM, respectively, while maintaining standard uncompressed training accuracy.The transformer also shows 2–3× speedup per epoch; preliminary CodeBERT pre-training reports 1.9×–2.3× speedup and 4.23× overall model reduction.Source: §1 · p. 2
  • Three optimizations target tensorized embedding lookups, tensor-network contractions in forward and backward propagation, and GPU overhead via CUDA Graph.The introduction also notes that low-rank parameters and gradients can reduce distributed communication cost by orders of magnitude.Source: §1 · pp. 2–3

2 Background

The background introduces tensor notation and contractions, then describes TT and TTM decompositions used to represent compressed neural-network weights and embeddings.

Tensor notation

  • A tensor is indexed by its entries and characterized by its order and the dimensions along each mode.The passage gives an order-d tensor with dimensions n1 through nd as an example.Source: §2 · p. 3

Tensor contractions

  • Tensor contraction combines tensors by summing over paired dimensions, producing a tensor over the remaining dimensions.The passages also describe einsum notation for specifying input dimensions and retained output dimensions.Source: §2 · p. 3

Tensor decomposition

  • TT decomposes a tensor into small cores, while TTM represents a higher-order tensor using cores suited to matrix-shaped weights.TT is used for linear-layer weights because of efficient tensor-vector multiplication; TTM suits highly unbalanced embedding tables.Source: §2 · p. 3

3 The CoMERA Training Framework

CoMERA formulates tensor-compressed training as balancing loss against model size, with rank control and staged scalarization to target resource requirements. Its training strategy seeks Pareto-efficient solutions while allowing optional further compression toward deployment targets.

3 The CoMERA Training Framework

  • CoMERA addresses automatic rank selection under resource limits by jointly considering training loss and compressed network size.The framework seeks a Pareto point when the two objectives cannot be minimized simultaneously.Source: §3 · p. 4, §3.1 · p. 4

3.1 Multi-Objective Training Model

  • Modified TT compression controls core ranks using diagonal matrices, enabling the compressed network size to be adjusted.The network is parameterized by tensor cores and the training data is optimized against a loss function.Source: §3.1 · p. 4
  • The multi-objective model minimizes loss and compressed size, seeking a Pareto point when no solution improves both objectives.The size term aggregates the size contributions across the model.Source: §3.1 · p. 4

3.2 Training Methods

  • Early training uses linear scalarization because aggressive rank pruning can harm convergence, and replaces the nonsmooth size penalty with an ℓ1 relaxation.The method uses different scalarization approaches at early and late training stages.Source: §3.2 · p. 4
  • An ℓ2 regularizer on the TT cores addresses the fact that the relaxed size term can approach zero without changing the loss.The regularized solution is shown to be a Pareto point for a constrained multi-objective problem.Source: §3.2 · pp. 4–5
  • Optional late-stage achievement scalarization moves training toward preferred loss and model-size targets for deployment.The late stage can be applied independently to a trained tensor-compressed model for further size reduction.Source: §3.2 · p. 5

4 Performance Optimization of CoMERA

CoMERA’s low-rank, small tensor operations reduce memory but can run slowly on GPUs; this section presents three optimizations to achieve practical training speedup.

4.1 Performance Optimization of TTM Embedding Tables.

  • TTM lookup reduces contraction overhead by deduplicating row indices and reusing unique tensor-index pairs before reconstructing requested rows.The lookup approach also extends to higher-order embedding tables by grouping tensor cores to form intermediate tensors.Source: §4.1 · pp. 5–6
  • 4–5× speedup and 2–3× memory saving versus standard TTM embedding without optimization; uncompressed embedding is faster but uses more memory.Figure 5 compares the proposed method with uncompressed embedding and several optimization variants.Source: Figure 5 · p. 6, §4.1 · p. 6

4.2 Contraction Path Optimization for TT-Vector Multiplications

  • For large batches, the near-optimal path contracts consecutive TT cores into groups, then combines these with the input sequentially, balancing FLOPs against intermediate-tensor memory costs.The grouping depends on dimensions, ranks, and batch size.Source: §4.2 · p. 7
  • For small batches, another path may use fewer FLOPs, but execution time is almost the same because the operations are small.The proposed path can therefore be used for most batch sizes.Source: §4.2 · p. 7

4.3 GPU Performance Optimization via CUDA Graph

  • CUDA Graph executes the computing graph together, reducing backend overhead from the many small kernels used by CoMERA’s tensor contractions.The authors describe this as an initial GPU-optimization step and expect further dedicated optimization to improve speedup.Source: §4.3 · p. 7

5 Training Results

Experiments compare CoMERA with uncompressed training and recent low-rank methods across a six-encoder transformer, DLRM, and CodeBERT pre-training. Results report compression, memory, accuracy, and epoch-time trade-offs across these settings.

5.1 A Medium-Size Transformer with Six Encoders

  • On the six-encoder transformer, early-stage CoMERA achieved 74× compression on tensorized layers with higher validation accuracy than uncompressed training; the smallest model reached 80× whole-model compression.That smallest model used 361× tensorized-layer compression with slightly worse accuracy. Figure 6 shows early-stage validation accuracy.Source: §5.1 · p. 8, Figure 6 · p. 8
  • 2–3× faster per epoch than uncompressed transformer training, CoMERA reached its target compression ratio with little accuracy drop in late-stage rank pruning.At target compression ratio 0.2, the second-last encoder and some linear layers in other encoders were removed; CoMERA with CUDAGraph achieved around 3× speedup. Figure 7 reports training time per epoch.Source: §5.1 · pp. 8–9, Figure 7 · p. 8

5.2 A DLRM Model with 4-GB Model Size

  • 99× whole-model compression and 7× lower peak memory accompanied similar accuracy on DLRM, although uncompressed training was faster for this embedding-intensive model.The optimized TTM lookup speeds training by around 2× and reduces memory by 4–5×; Figure 9 presents optimized-CoMERA performance. Figure 8 shows a similar convergence curve and fewer iterations for CoMERA.Source: §5.2 · p. 9, Figure 9 · p. 9, Figure 8 · p. 9

5.3 Comparison with GaLore and LTE

  • 9× more memory-efficient than GaLore in single-batch training, CoMERA is about 2× faster than GaLore and 3× faster than LTE per epoch on the six-encoder transformer.CoMERA and GaLore reach 64% validation accuracy on MNLI; LTE does not converge with its default setting. CoMERA uses the least memory across tested batch sizes.Source: §5.3 · pp. 9–10

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

  • 4.23× overall compression accompanied CodeBERT pre-training, with CoMERA 2.3× faster in Phase 1 and 1.9× faster in Phase 2 than standard pre-training.The reported final losses were 0.40 and 0.28; the authors note this gap does not necessarily imply downstream performance degradation. Figure 10 shows the pre-training loss curves.Source: §5.4 · p. 10, Figure 10 · p. 10

6 Conclusions and Future work

The conclusion presents CoMERA as a way to reduce training memory and time through rank-adaptive compression and GPU-oriented optimizations. It also identifies GPU-kernel support and optimization algorithms as future-work areas.

6 Conclusions and Future work

  • Transformer experiments showed 2–3× speedup per epoch, while transformer and DLRM model sizes were reduced by 43× to 99×.The conclusion reports 7× total peak-memory reduction for large-batch DLRM training on a single GPU and says CoMERA outperformed GaLore and LTE in memory and runtime efficiency.Source: §6 · p. 10
  • The reported speedups used tensorized-embedding and contraction optimizations plus CUDAGraph to reduce GPU latency.The authors also observed further speedup with mixed-precision computation.Source: §6 · p. 10
  • Small low-rank tensor operations lack adequate support from existing GPU kernels, and the authors identify comprehensive GPU/HPC optimization and CoMERA-specific optimizers as future work.They state that CoMERA's tensorized structure creates a different optimization landscape from uncompressed training.Source: §6 · p. 10

A.1 Comparison with existing works

The appendix contrasts CoMERA with prior tensor-compression approaches in rank control and contraction optimization. It emphasizes end-to-end training and joint optimization of forward and backward contractions.

A.1 Comparison with existing works

  • Unlike SVDinsTN, CoMERA compresses weights during end-to-end training without prior parameter information and uses a two-stage multi-objective procedure.Both methods use sparse diagonal matrices and ℓ1 terms to control tensor ranks; CoMERA adds late-stage achievement scalarization to target a preferred performance and model size.Source: §A.1 · p. 14
  • Unlike HEAT's post-training, single-path forward optimization in CP format, CoMERA jointly optimizes d + 2 contraction paths in forward and backward propagation in TT format.CoMERA also reuses intermediate results to reduce overall computation costs.Source: §A.1 · p. 14

A.2 Proof of Proposition 3.1

The proof establishes existence of a minimizer for the relaxed objective and connects its solution to a Pareto point of the multi-objective problem. It does so by bounding a minimizing sequence and identifying a cluster point.

A.2 Proof of Proposition 3.1

  • The objective in (9) is bounded below by zero, so its infimum is finite.A sequence approaching that infimum is bounded by the ℓ1 regularization of D and ℓ2 regularization of G, and therefore has a cluster point that minimizes (9).Source: §A.2 · p. 14
  • The proof identifies a minimizer of (9) and sets C to its squared G norm, making the relaxation equivalent to a constrained optimization problem.Source: §A.2 · p. 14
  • The solution to training problem (9) is a Pareto point of the multi-objective problem minimizing the loss and rank surrogate.Source: §A.2 · p. 14

A.3 Algorithm for Late Stage Optimization in Section 3.2

Algorithm 1 summarizes the late-stage optimization as an iterative procedure that updates tensor cores and rank-control parameters using optimization problems (13) and (14).

A.3 Algorithm for Late Stage Optimization in Section 3.2

  • The algorithm outputs tensor cores G_T and rank-control parameters D_T.Its inputs include initializations G0 and D0, constants, weights, and an optimization algorithm O.Source: §A.3 · p. 14
  • The procedure repeats for t = 0, . . . , T −1, starting from initializations G0 and D0.The inputs also include constants L0 and S0, weights w1 and w2, ρ, β, and optimization algorithm O.Source: §A.3 · p. 14
  • Each iteration runs one step of optimization problem (13) or (14), depending on the branch.The supplied algorithm description does not specify the branch condition.Source: §A.3 · p. 14

A.4 Proof of Proposition 4.1

The proof compares tensor-contraction paths by tracking tensor dimensions, cores, and intermediate tensors. It argues that the proposed sequential path reduces FLOPs under the stated batch-size conditions.

A.4 Proof of Proposition 4.1

  • The proof defines V_i as dimension strings, C_i as the cores used to obtain V_i, and X_i as the contraction of X with V_1 through V_i.These notations are used to compare contraction paths.Source: §A.4 · p. 14
  • For sufficiently large batch size b, regrouping consecutive tensor cores before contracting with X can use fewer FLOPs than direct contraction.The proof uses this comparison to argue that the first tensor must have the proposed format.Source: §A.4 · p. 15
  • Applying the same argument to subsequent tensors yields the stated format and sequential contraction order for V_1 through V_k.The proof extends the path argument from the initial contraction to the remaining tensors.Source: §A.4 · p. 15

A.5 Algorithm for Contraction Path in Section 4.2

The appendix details CoMERA’s contraction-path implementation and supplementary experiments. It reports training-time and memory advantages in several settings, alongside limits to mixed-precision acceleration on the tested models.

A.5 Algorithm for Contraction Path in Section 4.2

  • Algorithm 2 gives empirical forward- and backward-propagation paths, including stored intermediate results used to compute outputs and gradients.Its forward procedure reshapes the input, computes intermediates, and forms the output; its backward procedure uses stored results to compute gradients.Source: §A.5 · p. 15

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

  • The empirical path eliminates the batch dimension early and is nearly optimal for large batches; for small batches, a greedy search finds a different path with almost the same GPU execution time.The empirical path is therefore used for all batch sizes.Source: §A.6 · pp. 15–16, Figure 11 · p. 16

A.7 Compression Settings for the Experiment in Section 5.1

  • Table 4 presents the tensorized settings used for the Transformer experiment in Section 5.1.The supplied passages identify the table’s scope but do not provide its settings.Source: §A.7 · p. 16, Table 4 · p. 17

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

  • For MNLI, CoMERA had a speedup at every compression ratio above 1, while its training time approached uncompressed training as the ratio approached 1.The acceleration became more pronounced at larger compression ratios.Source: Figure 12 · p. 17, §A.8 · p. 17

A.10 CoMERA Pretraining Result on BERTLARGE

  • CoMERA’s BERTLARGE variant achieved 2.69× overall compression, with final loss 1.45 vs 1.26; it outperformed BERTLARGE on SST-2 and MRPC accuracy but underperformed on SQuAD F1.The reported accuracy pairs were 92.10% vs 91.74% on SST-2 and 86.82% vs 86.00% on MRPC; SQuAD F1 was 88.76% vs 90.68%.Source: §A.10 · p. 17

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

  • CoMERA was fastest and used the least memory among CoMERA, GaLore, and LTE on the RTX 4090; at batch size 1, it was 1.7× faster than GaLore.The batch-size-1 comparison was across RTX 3090 and RTX 4090 runs, with speedups of 1.2× and 1.7×, respectively.Source: §A.9 · p. 17, Figure 13 · p. 18

A.11 Discussion: Mixed-Precision CoMERA

  • FP8 tensor-compressed linear layers were about 3× faster than FP8 vanilla layers for large batches and layers, but vanilla layers were faster at small batch sizes.The sequential tensor contractions are not well supported by current GPU kernels, which limits small-batch speed.Source: §A.11 · p. 18
  • FP8 training accuracy was similar to FP32, but the tested experiments showed little acceleration due to casting overhead, sequential small contractions, and relatively small linear layers.The convergence curve shown is for the six-encoder transformer.Source: Figure 14 · p. 18, §A.11 · pp. 18–19

Related papers

Loading 2405.14377v2…