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

TL;DR

CoMERA tackles the high cost of training large AI models by introducing rank-adaptive tensor optimization to reduce memory and compute. The framework uses a modified tensor-train (TT) / tensor-train-matrix (TTM) representation with diagonal rank controls and a multi-objective optimization to balance accuracy and model size, including early-stage scalarization and optional late-stage refinement. Three GPU-oriented optimizations—optimized TT-embedding lookups, contraction-path optimization for TT-vector multiplications, and CUDA Graph-based execution—drive practical speedups, achieving 2-3x per-epoch improvements on transformers and up to 99x overall compression on DLRM, with notable gains in CodeBERT pretraining. These results suggest CoMERA can enable efficient training of large models on limited hardware, motivate further HPC- and kernel-level optimizations, and support automatic architecture search via rank adaptation.

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 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.

Figures (14)

• Figure 1: Training time and total memory cost of CoMERA, GaLore zhao2024galore and LTE huh2024training on a six-encoder transformer with varying batch sizes. The experiment is done on Nvidia RTX 3090 GPU.
• Figure 2: (a) Tensors. (b) Tensor contractions.
• Figure 3: Tensor networks for (a) tensor-train and (b) tensor-train-matrix decompositions.
• Figure 4: Optimized TTM embedding lookup.
• Figure 5: Performance of optimized TTM embedding table lookup. The labels uncompressed, proposed approach, optimized order, unique indices, without optimization represent standard embedding with sparse gradients, the new method in \ref{['sc:embedding']}, the method that only uses the unique order, the method that only uses the unique indices, and the method without optimization, respectively.

Theorems & Definitions (6)

• Proposition 3.1
• proof
• Proposition 4.1
• proof
• proof
• proof