Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking

Abstract: Fueled by their remarkable ability to tackle diverse tasks across multiple domains, LLMs have grown at an unprecedented rate, with some recent models containing trillions of parameters. This growth is accompanied by substantial computational challenges, particularly regarding the memory and compute resources required for training and fine-tuning. Numerous approaches have been explored to address these issues, such as LoRA. While these methods are effective for fine-tuning, their application to pre-training is significantly more challenging due to the need to learn vast datasets. Motivated by this issue, we aim to address the following questions: Can parameter- or memory-efficient methods enhance pre-training efficiency while achieving performance comparable to full-model training? How can the performance gap be narrowed? To this end, the contributions of this work are the following. (1) We begin by conducting a comprehensive survey that summarizes state-of-the-art methods for efficient pre-training. (2) We perform a benchmark evaluation of several representative memory efficient pre-training approaches to comprehensively evaluate their performance across model sizes. We observe that with a proper choice of optimizer and hyperparameters, full-rank training delivers the best performance, as expected. We also notice that incorporating high-rank updates in low-rank approaches is the key to improving their performance. (3) Finally, we propose two practical techniques, namely weight refactorization and momentum reset, to enhance the performance of efficient pre-training methods. We observe that applying these techniques to the low-rank method (on a 1B model) can achieve a lower perplexity than popular memory efficient algorithms such as GaLore and Fira, while simultaneously using about 25% less memory.

Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking

The paper under review addresses a pertinent challenge in the domain of large language models (LLMs): the substantial computational resources required for training and fine-tuning models with trillions of parameters. With the rapid expansion in model size, the paper seeks to explore and evaluate various methods for parameter and memory-efficient pre-training to bridge the gap between full-model training and more resource-efficient approaches.

The authors delineate their contributions into three primary aspects:

First, the paper presents a comprehensive survey of current methodologies aimed at enhancing pre-training efficiency. Techniques such as Low-Rank Adaptation (LoRA), which has been efficacious in fine-tuning, are scrutinized for their applicability to the more demanding pre-training phase. Herein, the challenge lies in the necessity to handle extensive datasets and maintain model expressiveness, which often falls short when restricted parameter spaces are employed.

Second, the paper embarks on a benchmarking exercise, evaluating these memory-efficient strategies across models of varying sizes. The comparison reveals several insights: full-rank training, under appropriate configurations, still offers optimal performance. However, memory-efficient methods show promise, wherein techniques like GaLore benefit from high-rank updates to preserve a richer gradient information and mitigate performance loss. Additionally, innovative methods such as SLTrain, which combine sparse and low-rank representations, demonstrate noteworthy achievements in reducing parameter count without significant detriments to learning efficacy.

Third, two novel techniques—weight refactorization and momentum reset—are introduced to ameliorate the performance of efficient pre-training approaches. Weight refactorization involves restructuring the model's parameters during training to optimize the conditioning and convergence properties of the training algorithms. Momentum reset is posited as a mechanism to stabilize training dynamics by periodically nullifying the accumulated optimizer's momentum states, thereby enhancing the robustness and convergence rate of learning processes.

The benchmark results elucidate the competitive landscape of efficient pre-training methods, highlighting the trade-offs involved in memory usage, parameter count, and model performance. Besides presenting strong numerical results, the authors also posit the intriguing possibility that by combining the proposed innovations, low-rank and sparse methods could potentially match or even exceed the performance of some full-rank approaches, while offering substantial memory savings—45% less memory in the case of SLTrain compared to standard methods for select configurations.

The paper furthermore speculates on future pathways for AI development, suggesting that overcoming the resource constraints imposed by traditional LLM pre-training could foster the proliferation of sophisticated models in environments previously constrained by inadequate computational infrastructure. This possibility brings theoretical and practical implications where AI models might be more widely accessible and deployed on a broader scale, thus democratizing advanced language model functionalities.

In sum, the research pushes the boundary on efficient LLM pre-training, introducing significant advancements that emphasize both theoretical insights and practical applicability, and paving the way for more scalable AI systems in resource-constrained contexts.