Memory-Efficient LLM Training with Online Subspace Descent
This addresses memory constraints in large language model training for researchers and practitioners, offering a more flexible and efficient optimizer with theoretical guarantees.
The paper tackles the problem of memory-efficient LLM training by providing the first convergence guarantee for arbitrary projection matrix update rules in subspace descent optimizers, and proposes Online Subspace Descent which replaces SVD with online PCA. It shows that Online Subspace Descent achieves lower perplexity and better downstream task performance than state-of-the-art low-rank methods on LLaMA models up to 7B parameters, narrowing the gap with full-rank baselines.
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 \emph{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.