Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures

Authors: Dang Nguyen, Wenhan Yang, Rathul Anand, Yu Yang, Baharan Mirzasoleiman
Primary area: foundation or frontier models, including LLMs
Keywords: Data selection, Memory efficient, Large Language Models, Data Mixtures

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.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730403722006

thread_type:
official_review

<thread_object>
creator: reviewer_twby

<content>
summary:
The paper introduces a method in using smaller mini-batches to simulate the performance of training with larger mini-batches, in the context of Large Language Models (LLMs). In particular, the authors propose coresets for training LLMs (CoLM) which involves the following aspects:

1. Including all examples from small sources in the mini-batch coresets
2. Adapt the mini-batch coreset selection process with ADAM, selecting medoids of big sources.
3. Use zero-order optimization technique based on MeZO in combination with LoRA to estimate the gradients.

Experiments show that the proposed method achieves comparable performance to regular fine-tuning that uses a much larger batch size, while reducing the memory requirements.

strengths:
S1: Well written paper with nice motivation to enable both memory efficiency and performance gains by looking at a data perspective.

S2: Strong experimental results that demonstrate the benefits of their approach compared to baselines. The increase in training time due to mini-batch selection also justified given significant improvement in accuracy.

S3: Ablation studies give insight to the relative importance to each of the three components proposed, with emphasis on gradient approximation, which helps to better understand the benefits of the method.

weaknesses:
W1: There should be some further clarification on how "all" examples of small sources can be included in the mini-batch coreset. Based on experiments, figure 4 shows that the data count for even small data sources can be over 1000 examples. Including all the sources in a single mini-batch does not seem to be plausible, but implied from the methodology and Algorithm 1.

W2: Dependency on having access to the true source data distribution may not always be possible. In these cases, it is difficult to distinguish between small and large sources. In addition, how to determine the boundary between "small" and "large" partitions could be discussed in more detail with experiments.

questions:
Q1: Related to W1 - Is there some sampling or representative subset selection method for the small sources going on behind the scenes?

Q2: Related to W2 - How to determine the boundary between "small" and "large" partitions, especially with skewed distributions where average count may not serve as a good cutoff?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730491587625

thread_type:
official_review

<thread_object>
creator: reviewer_5bea

<content>
summary:
The paper identifies three key difficulties when applying previous coreset findings methods (methods for finding small mini-batches that have gradients whose mean closely approximates the full gradient) to fine-tuning LLMs: imbalanced data, Adam's update, and large gradient dimensionality. The authors systematically formalize the difficulties and design methods to address them: by treating small-source examples and large-source examples differently, by comparing gradients in the space of Adam updates, and by using zeroth order gradient estimations methods, sparsifying the gradient, and using an appropriate distance metric. The combination of these improvements makes up their method,  Coresets for Training LLMs (CoLM). In their empirical study, authors demonstrate that their method leads to stronger performance compared to a number of baselines (including strong ones) and uses less memory than comparable methods while adding a reasonable overhead in training time.

strengths:
- The paper is very well written. It provides clear motivations of the problem and in thoughtfully structured around the three key difficulties of applying coreset finding methods to LLM fine-tuning. This structure makes the paper clear and very easy to follow. 
quality, 
- The authors provide extensive ablation analyses about the great majority (if not all) of the improvements they propose, which is greatly appreciated and helps show they are worthwhile.
- The authors empirically demonstrate the effectiveness of CoLM for different pre-trained language models and across many datasets (Math Instruct, SST-2, CB, MultiRC).
- The method often outperforms fine-tuning with twice the batch size.
- The authors demonstrate that the method achieves a smaller memory overhead than a number of representative baselines.

weaknesses:
- Using gradient accumulation + LoRA trivially reduces the memory overhead. The paper would benefit from experiments making an explicit comparison to this case. It seems from Fig 2 (a) that you should be able to claim CoLM is still faster in this case, but not showing it and claiming memory efficiency makes the paper seem weak.

- While the memory overhead is reduced due to the smaller batch size, there is no figure showing how much the total memory decreased by making the batch size smaller when fine-tuning with LoRAs using CoLM. Does the batch size reduction account for a substantial portion of the memory consumption or is the memory consumption dominated by the model parameters given the smaller optimizer states required for one LoRA? Providing a figure with the memory consumption of CoLM compared to fine-tuning at different batch sizes is warranted given the author's claims about memory efficiency.

- Do the methods to estimate gradient distances scale to larger vocabulary sizes? Recent SOTA models (Llama 3 has 128k vocab size whereas phi2,phi3,and zephyr have 50k at most) have been using much larger vocabulary sizes and recent work "Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies" shows this is an important axis along which to scale LLMs. Simply including one experiment with a larger tokenizer 100k vocab+ would strengthen this aspect.

- *I believe the above concerns are easily addressable with small additional experiments. I would be happy to raise my score if these concerns are addressed. *

*List of Typos*:
 "to iteratively selecting small mini-batches from larger random batches." selecting --> select

questions:
(Q1) For Figure 3 (c), how large does the batch size need to be to match your method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730509987751

thread_type:
official_review

<thread_object>
creator: reviewer_dq66

<content>
summary:
This paper proposes a method called Coresets for Training Large Language Models (CoLM) to enable memory-efficient training of large language models (LLMs) by constructing mini-batch "coresets." These coresets are small, representative subsets of mini-batches that approximate the gradient of larger batches, allowing effective model training with reduced memory usage. CoLM addresses challenges such as the high memory demands of large mini-batches, the imbalanced distribution in language data, and the complexities of using the Adam optimizer in high-dimensional gradient spaces.

strengths:
- The presentation is clear, and the paper is easy to follow, with only a few minor typos.
- The proposed method, CoLM, is straightforward and demonstrates strong empirical performance.

weaknesses:
- **Limited Base Models**: While the authors mention using Phi-2, Phi-3, and Zephyr, the main results in Tables 1 and 7 only report Phi-2. To more fully demonstrate CoLM’s performance across various models, I recommend including Phi-3 in these tables and consider adding state-of-the-art models, such as LLaMA-3 8B and LLaMA-3.1.

- **Insufficient Discussion of Related Work and Novelty**: One of CoLM’s contributions is gradient normalization for the Adam optimizer, which improves mini-batch approximation by using historical averages. However, previous studies, such as [1], have explored renormalization techniques with mini-batches. It would strengthen the paper if the authors could discuss distinctions between [1] and CoLM to clarify CoLM’s unique contributions.



Reference:

[1] Batch Renormalization: Towards Reducing Minibatch Dependence in Batch-Normalized Models

questions:
- **Computational Complexity**: CoLM introduces additional computational costs to find mini-batches. Could the authors discuss the complexity and time requirements for this estimation algorithm and expand on it in the main paper?

- **Convergence Rate Comparison**: Figure 3 compares CoLM and fine-tuning (FT) by plotting accuracy against training iterations to highlight CoLM’s faster convergence. Could the authors include a comparison of CoLM with LoRA in Figure 3? It would also be valuable to add perplexity or training loss versus iterations to provide further insight into convergence rates.

- **Time vs. Accuracy Comparison**: Given that CoLM involves additional computation to determine mini-batches, comparing accuracy against time would offer a fairer evaluation alongside FT and LoRA in Figure 3(a).

- **Quantitative Memory Cost Analysis**: While the paper claims CoLM uses 4x less memory, it lacks a quantitative breakdown of memory usage. Could the authors add specific memory costs for LoRA, other low-rank adaptation methods, and CoLM in a table? This would provide stronger evidence of CoLM’s memory efficiency.

- **Clarification on Motivation**: 1) In the introduction, the authors state that fine-tuning a model like Phi-2 with 2.7B parameters and a batch size of 128 requires at least 44 GB of GPU memory. However, fine-tuning LLMs primarily incurs memory costs due to model weights and optimizer states, which together cost approximately three times the model weight. Since Phi-2’s model size is under 7 GB, the total memory cost should be around 22 GB rather than 44 GB. Furthermore, batch size typically only affects activation memory and does not significantly impact memory for optimizer states or model weights, and the activation memory cost is not an issue for PEFT methods such as LoRA. 2) **Comparison to LoRA**: LoRA also reduce optimization memory significantly. Besides, LoRA also achieve significantly speedup since it has less trainable parameters. CoLM achieve similar wall clock time with full fine-tuning, but when it applies to full fine-tuning, it will be much slower than LoRA. Could the authors clarify the motivation of the paper from these aspects? 



- **Title Consistency**: Please align the submission title, "Memory-efficient Training of Large Language Models with Larger Mini-batches," with the manuscript title, "MINI-BATCH CORESETS FOR MEMORY-EFFICIENT TRAINING OF LARGE LANGUAGE MODELS."

I would like to discuss the questions I raised regarding the weaknesses and concerns with the authors. If my concerns are adequately addressed, I would be willing to reconsider my rating.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731041113038

thread_type:
official_review

<thread_object>
creator: reviewer_akmp

<content>
summary:
The paper introduces CoLM (Coresets for Training LLMs), a method for memory-efficient training of large language models (LLMs) using mini-batch coresets that approximate the gradients of larger batches. CoLM addresses specific challenges in LLM training, such as data imbalance, high gradient dimensionality, and compatibility with the Adam optimizer. CoLM addresses them by including all examples from smaller data sources, normalizing gradients for Adam, and using zeroth-order gradient estimates. It reduces memory requirements by 2x while maintaining or surpassing performance compared to larger mini-batches.

strengths:
1. The approach takes a unique and innovative perspective on memory-efficient training. Unlike typical methods in field, it focuses on data selection and reduces batch size to achieve memory-efficient training.
2. The paper is clear to understand and well motivated. 
2. The paper presents comprehensive experiments on diverse datasets and models, substantiating the benefits of CoLM over various baselines.

weaknesses:
1. Computational overhead: Figure 2(b) suggests that the computational overhead for CoLM is notably high, doubling the training time in some cases, whereas typical memory-efficient methods usually incur 5-20% overhead. A comparison of "memory vs. computational overhead vs. performance" with other methods would be beneficial.
2. The paper mainly evaluates CoLM in a fine-tuning context, where model and optimizer state memory are the main bottlenecks. It would be valuable to see CoLM tested in pre-training or continual pre-training settings, where activation memory dominates due to large batch sizes. This could pose additional challenges given complex data structures and extensive amount of training data.
3. Not enough details of algorithm and discussion on hyper-parameters.

questions:
1. Is zeroth-order necessarily needed? How about estimating using standard first-order methods?
2. What hypeperparametrs involved in CoLM? Like how many small sources and large sources needed? How much is the cost of tuning those hyper parameters? Are they sensitive to different data structures and types of data?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732327383657

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment

comment:
We thank all the reviewers for their valuable feedback and supporting our work. In our general comment, we discuss the memory and time requirement of our method CoLM and when it yields memory reduction, speedup, and performance improvement in different settings: (i) stacked with LoRA and gradient accumulation, (ii) large batch size, (iii) applied to state-of-the-art models (Llama3). We have added the discussion to our revised version.

**CoLM Effectively Reduces the Activation Memory.** The memory required for training an LLM can be decomposed into three parts: activation memory + weight memory + optimizer state memory. Memory-efficient methods often reduce the weight or optimizer-state memory. For example, LoRA reduces the optimizer state memory but slightly increases the weight and activation memory by adding low-rank matrices. Orthogonal to such methods, CoLM effectively reduces the activation memory by reducing the batch size. Hence, stacked with memory-efficient methods such as LoRA, it can further reduce memory, particularly when batch size is large. Notably, while memory-efficient methods harm the performance, CoLM effectively improves the performance over training with small and larger batches. For reference, full-parameter fine-tuning of Phi-2 on MathInstruct with bs=128 achieves 56% average accuracy, while LoRA fine-tuning with bs=128 achieves 52.6%, and CoLM (128->64) with LoRA achieves 56.5%.

**CoLM Easily Stacks with LoRA and Gradient Accumulation**. In all our experiments in the paper, we used LoRA (lines 365-366 and caption of table 1), and a gradient accumulation step of 8 for all methods, as our A40 GPUs didn’t have enough memory to hold bs=128 (we added the detailed explanation to our revised Appendix E). As explained above, CoLM easily stacks with existing memory-efficient methods.

Below, we first discuss the negligible memory overhead of CoLM and then clarify its benefit in reducing the memory requirement, speeding up the training, and improving the performance in different settings.

**Memory overhead of CoLM**. The memory overhead of CoLM stems from three main sources the last layer zeroth-order gradient $\hat{g}_{i,t}^{vp}$ in Eq (6), the historical terms of Adam in Eq (8), and the pairwise dissimilarity matrix when solving Eq (8). Because with LoRA, the last layer dimension is 327K and 2560 before and after sparsification, for CoLM (128->64), the memory overhead is less than 200MB.

The following table summarizes the memory, wall-clock time, and average accuracy, when fine-tuning Phi-2 on MathInstruct (setting of Table 1), with 4xA40 GPUs each with 45G memory. Using LoRA, each GPU can have a maximum device batch size of 5, so we trained all models with a gradient accumulation step of 8, with LoRA on 4 GPUs. The total batch size = #GPUs x device batch size x gradient accumulation step.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces CoLM, a memory-efficient method that constructs small, representative mini-batch coresets to approximate the gradients of larger batches. CoLM addresses three key challenges in applying coreset methods to LLMs: (1) imbalanced data distributions, (2) the complexities of Adam's update rule in high-dimensional gradient spaces, and (3) computational difficulties due to large gradient dimensionality.

The proposed method tackles these issues by including all examples from smaller data sources, adapting gradient comparisons to the Adam update space, and employing zeroth-order optimization techniques with sparsified gradients. CoLM achieves comparable or superior performance to standard fine-tuning methods with large mini-batches while reducing memory requirements and adding only modest computational overhead.

While this paper has several strengths, **it also has several weaknesses**, as pointed out by Reviewer dq66. (1) The three key components of the proposed method have been explored in prior work, and the paper lacks detailed comparisons and explanations to highlight its unique contributions. (2) The main results are based on fine-tuning with LoRA, which inherently has low memory requirements; demonstrating a 2x memory reduction would be more compelling in the context of full fine-tuning. (3) Appropriate ablation studies are needed to evaluate the contributions of each component independently.

Despite these limitations, the paper’s strengths outweigh its weaknesses. Most reviewers supported acceptance, with several providing high scores, reflecting confidence in the paper’s contributions. Therefore, I recommend accepting this paper.