Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures
Dang Nguyen, Wenhan Yang, Rathul Anand, Yu Yang, Baharan Mirzasoleiman
TL;DR
The paper introduces CoLM, a data-centric approach to memory-efficient fine-tuning of large language models by constructing mini-batch coresets that approximate gradients from larger batches. It tackles three core challenges—imbalanced data sources, Adam optimization, and high gradient dimensionality—via including all small-source examples, normalizing gradients with historical averages for Adam, and using zeroth-order gradient estimates with sparsification to enable efficient medoid selection. Empirical results on MathInstruct and SuperGLUE show CoLM reduces activation memory by about 2x and can outperform training with 4x larger mini-batches, while also speeding up training and integrating with LoRA. The method is demonstrated to generalize across multiple model families and data settings, offering a practical path to memory-efficient, high-performance fine-tuning of LLMs.
Abstract
Training with larger mini-batches improves the convergence rate and can yield superior performance. However, training with large mini-batches becomes prohibitive for Large Language Models (LLMs), due to the large GPU memory requirement. To address this problem, an effective approach is finding small mini-batch coresets that closely match the gradient of larger mini-batches. However, this approach becomes infeasible and ineffective for LLMs, due to the highly imbalanced mixture of sources in language data, use of the Adam optimizer, and the very large gradient dimensionality of LLMs. In this work, we address the above challenges by proposing Coresets for Training LLMs (CoLM). First, we show that mini-batch coresets found by gradient matching do not contain representative examples of the small sources w.h.p., and thus including all examples of the small sources in the mini-batch coresets is crucial for optimal performance. Second, we normalize the gradients by their historical exponential to find mini-batch coresets for training with Adam. Finally, we leverage zeroth-order methods to find smooth gradient of the last V-projection matrix and sparsify it to keep the dimensions with the largest normalized gradient magnitude. We apply CoLM to fine-tuning Phi-2, Phi-3, Zephyr, and Llama-3 models with LoRA on MathInstruct and SuperGLUE benchmark. Remarkably, CoLM reduces the memory requirement of fine-tuning by 2x and even outperforms training with 4x larger mini-batches. Moreover, CoLM seamlessly integrates with existing memory-efficient training methods like LoRA, further reducing the memory requirements of training LLMs. Our code is available at https://github.com/BigML-CS-UCLA/CoLM.