Memory-Efficient LLM Training with Online Subspace Descent
Authors: name: Kaizhao Liang; name: Bo Liu; name: Lizhang Chen; name: Qiang Liu
Published: 2024-08-23 05:54:53

Abstract: Recently, a wide range of memory-efficient LLM training algorithms have
gained substantial popularity. These methods leverage the low-rank structure of
gradients to project optimizer states into a subspace using projection matrix
found by singular value decomposition (SVD). However, convergence of these
algorithms is highly dependent on the update rules of their projection matrix.
In this work, we provide the first convergence guarantee for arbitrary
update rules of projection matrix. This guarantee is generally applicable to
optimizers that can be analyzed with Hamiltonian Descent, including most common
ones, such as LION, Adam. Inspired by our theoretical understanding, we propose
Online Subspace Descent, a new family of subspace descent optimizer without
SVD. Instead of updating the projection matrix with eigenvectors, Online
Subspace Descent updates the projection matrix with online PCA. Online Subspace
Descent is flexible and introduces only minimum overhead to training. We show
that for the task of pretraining LLaMA models ranging from 60M to 7B parameters
on the C4 dataset, Online Subspace Descent achieves lower perplexity and better
downstream tasks performance than state-of-the-art low-rank training methods
across different settings and narrows the gap with full-rank baselines.