This paper presents research on domain reweighting, which adjusts the relative weights of various data sources to improve the efficiency and effectiveness of LLM pretraining. Specifically, we highlight that data blending that performs well in small-scale experiments may not maintain its benefits at scale. To address this, we propose AutoScale, a two-stage scale-aware data composition framework. AutoScale first fits a parametric model that predicts the model's loss under various data configurations and then uses this model to find the optimal allocation within a smaller budget. Then, leveraging novel theoretical analysis of how the optimal configuration evolves with scale, it extrapolates this configuration to larger budgets without additional retraining. AutoScale accelerates convergence and improves downstream performance. It achieves 28% faster perplexity reduction than existing methods and up to 38% faster than unweighted training when pretraining the GPT-2 Large model. Furthermore, it achieves the best average performance across various downstream tasks.
