CompAct: Compressed Activations for Memory-Efficient LLM Training

Shamshoum, Yara, Hodos, Nitzan, Sieradzki, Yuval, Schuster, Assaf

arXiv.org Artificial Intelligence 

We introduce CompAct, a technique that reduces peak memory utilization on GPU by 25-30% for pretraining and 50% for fine-tuning of LLMs. Peak device memory is a major limiting factor in training LLMs, with various recent works aiming to reduce model memory. However most works don't target the largest component of allocated memory during training: the model's compute graph, which is stored for the backward pass. By storing low-rank, Figure 1: Breakdown of memory components for various compressed activations to be used in the backward LLaMA model sizes, with batch size 256. Blue: pass we greatly reduce the required memory, linear operations compressed by CompAct; Red: nonlinear unlike previous methods which only reduce operations which CompAct doesn't compress; optimizer overheads or the number of Green: model parameters and non-linear operation's optimizer trained parameters. Our compression uses random states. Most of the memory is used by the computational projection matrices, thus avoiding additional graph. CompAct's compression gets more memory overheads.