Clustering improves sampling for faster balanced fine-tuning of large language models
Clustering-Based Balanced Sampling and Allocation with Data Parallelism for High-Performance Fine-Tuning
Machine LearningArtificial Intelligence
Summary
Training large language models with instruction data can waste time by using too many repetitive examples and missing rare but important ones. The authors propose a method called CluSTER that groups similar training examples based on their influence on learning and then carefully distributes these groups across multiple GPUs. This balanced approach avoids repeatedly training on the same overrepresented examples, reducing training time a lot without hurting the model’s performance. CluSTER helps models learn faster and more reliably from large and diverse instruction datasets.
What this means in practice
- •For machine learning engineers: Accelerate distributed fine-tuning of large language models by reducing redundant training data while maintaining accuracy.
- •For cloud infrastructure teams: Cut GPU usage costs by shortening large language model training times through intelligent data sampling across multiple devices.
Authors
Hyunjin Kim, Youngeun Nam, Jaemin Han, Wonhyeok Choi, Jae-Gil Lee
Abstract
Instruction-tuning datasets for large language models (LLMs) are often large, redundant, and imbalanced, limiting efficient adaptation. Naive large-batch fine-tuning repeatedly includes overrepresented sample groups while weakly covering underrepresented but informative ones, especially under data parallelism (DP) across multiple GPUs. We propose CluSTER, a Cluster-aware balanced Sampling framework for Training Efficient data Reduction in DP instruction tuning. CluSTER curates a representative reduced dataset through gradient-space clustering and DP-aware balanced allocation, ensuring dual-level coverage across clusters and workers, while preserving the original data distribution by weighted update. As a result, CluSTER reduces redundant computation and improves training stability without compromising model quality. Across multiple instruction-tuning datasets, CluSTER reduces training time by up to 69.6% with almost no accuracy loss compared to prior sampling and data reduction methods. Code is available at https://github.com/kaist-dmlab/CluSTER.