Clustering-Based Balanced Sampling and Allocation with Data Parallelism for High-Performance Fine-Tuning

2026-09-11

AI Summary

为解决大规模语言模型微调数据集大、冗余和不平衡问题，提出CluSTER框架，通过梯度空间聚类和平衡分配减少冗余计算，提高训练稳定性。

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.

Problem

• instruction-tuning
• large language models
• data parallelism
• imbalanced dataset
• redundant computation

Innovation

• Clustering
• Data Parallelism
• Balanced Sampling
• Instruction Tuning
• Efficient Data Reduction