Accelerating Memory-Efficient LLM Training and Fine-Tuning via Tracking the Gradient Subspace

Abstract

Training and fine-tuning Large Language Models (LLMs) is often highly resource- and time-intensive due to their large model sizes. To address this issue and improve accessibility, several memory-efficient techniques have been developed, such as Low-Rank Adaptation (LoRA), which optimizes the weights in a low-rank subspace, and Gradient Low-Rank Projection (GaLore), which projects gradients onto a lower-dimensional space. In this paper, we introduce Gradient Subspace Tracking (SubTrack), a method that restricts the optimization process to a small core subspace of gradient matrices while dynamically tracking subspace changes. By leveraging estimation errors and previously detected subspaces, SubTrack adjusts the subspace estimation using a computationally efficient approach. Despite applying only rank-1 updates, SubTrack achieves performance comparable to, or better than, GaLore while reducing runtime by up to 20.56%.