Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training

Determining an effective data mixture is a key factor in Large Language Model (LLM) pre-training, where models must balance general competence with proficiency on hard tasks such as math and code. However, identifying an optimal mixture remains an open challenge, as existing approaches either rely on unreliable tiny-scale proxy experiments or require prohibitively expensive large-scale exploration. To address this, we propose Decouple Searching from Training Mix (DeMix), a novel framework that leverages model merging to predict optimal data ratios. Instead of training proxy models for every sampled mixture, DeMix trains component models on candidate datasets at scale and derives data mixture proxies via weighted model merging. This paradigm decouples search from training costs, enabling evaluation of unlimited sampled mixtures without extra training burden and thus facilitating better mixture discovery through more search trials. Extensive experiments demonstrate that DeMix breaks the trade-off between sufficiency, accuracy and efficiency, obtaining the optimal mixture with higher benchmark performance at lower search cost. Additionally, we release the DeMix Corpora, a comprehensive 22T-token dataset comprising high-quality pre-training data with validated mixtures to facilitate open research. Our code and DeMix Corpora is available at https://github.com/Lucius-lsr/DeMix.

Large language models, such as chatbots and coding assistants, learn from huge collections of text before they are adapted to specific uses. What those collections contain matters: too little math or code can make a model weak on reasoning and programming, while too much of one type of data can hurt its general abilities. Today, finding the right recipe for mixing data is difficult because cheap small-scale tests can be misleading, while trying many recipes at full scale is extremely expensive. We introduce DeMix, a method that separates the search for a good data recipe from the costly process of training a large model. DeMix first trains separate models on different kinds of data, then combines these models in different proportions to cheaply estimate how a full model would perform under many possible data mixtures. This lets researchers test far more data recipes without repeatedly training large models from scratch. In our experiments, DeMix found mixtures that led to stronger benchmark performance while using less search cost. We also release DeMix Corpora, a 22-trillion-token collection of pre-training data and validated mixtures, to support open research on building better language models.