Natural GaLore: Accelerating GaLore for memory-efficient LLM Training and Fine-tuning

TL;DR

Natural GaLore is a memory-efficient optimizer that accelerates large language model training and fine-tuning by applying second-order information with minimal memory overhead, outperforming existing methods.

Contribution

The paper introduces Natural GaLore, a drop-in replacement for AdamW that efficiently incorporates inverse Fisher information for faster, memory-efficient LLM training and fine-tuning.

Findings

Achieves lower perplexity on large Llama models without extra memory.

Reduces fine-tuning gap on GLUE benchmark by over 86%.

Outperforms 16-bit LoRA and surpasses GPT-4 Turbo in accuracy with less memory.

Abstract

Training LLMs presents significant memory challenges due to growing size of data, weights, and optimizer states. Techniques such as data and model parallelism, gradient checkpointing, and offloading strategies address this issue but are often infeasible due to hardware constraints. To mitigate memory usage, alternative methods like Parameter-Efficient-Fine-Tuning (PEFT) and GaLore approximate weights or optimizer states. PEFT methods, such as LoRA, have gained popularity for fine-tuning LLMs, though they require a full-rank warm start. In contrast, GaLore allows full-parameter learning while being more memory-efficient. This work introduces Natural GaLore, a simple drop in replacement for AdamW, which efficiently applies the inverse Empirical Fisher Information Matrix to low-rank gradients using Woodbury's Identity. We demonstrate that incorporating second-order information speeds up optimization significantly, especially when the iteration budget is limited. Empirical pretraining on 60M, 130M, 350M, and 1.1B parameter Llama models on C4 data demonstrate significantly lower perplexity over GaLore without additional memory overhead. By fine-tuning RoBERTa on the GLUE benchmark using Natural GaLore, we demonstrate significant reduction in gap 86.05% vs 86.28% for full-finetuning. Furthermore, fine-tuning the TinyLlama 1.1B model for function calling using the TinyAgent framework shows that Natural GaLore achieving 83.09% accuracy on the TinyAgent dataset, significantly outperforms 16-bit LoRA at 80.06% and even surpasses GPT4-Turbo by 4%, all while using 30% less memory. All code to reproduce the results are available at: https://github.com/selfsupervised-ai/Natural-GaLore.git