DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models
DataFlex unifies data selection, mixing, and reweighting, boosting large model training efficiency and accuracy.
Key Findings
Methodology
Built on LLaMA-Factory, DataFlex introduces a modular architecture supporting three core data optimization strategies: sample selection, domain mixing, and sample reweighting. It defines three trainer types—SelectTrainer, MixTrainer, and WeightTrainer—each integrating specific algorithms like LESS, NICE, DoReMi, and ODM. The system standardizes model operations such as embedding extraction, inference, and gradient computation, enabling scalable training with frameworks like DeepSpeed ZeRO-3. Extensive experiments demonstrate that dynamic strategies outperform static full-data training on tasks like MMLU, improving accuracy by 3-4 percentage points and enhancing convergence speed.
Key Results
- In experiments with Mistral-7B and Llama-3.2-3B, dynamic sample selection methods improved MMLU accuracy by approximately 3-4%, consistently surpassing static training baselines. The approach maintained stability across different model architectures and datasets.
- Applying DoReMi and ODM for data mixture during Qwen2.5-1.5B pretraining yielded 2-3% accuracy gains and over 10% perplexity reduction at 6B and 30B token scales, indicating better generalization and training efficiency.
- System optimizations resulted in a 10-15% reduction in runtime, demonstrating scalability and efficiency in large-scale training environments.
Significance
This work addresses the fragmentation in data-centric training methods by providing a unified, flexible platform that supports multiple algorithms within a standardized framework. It facilitates fair comparison, reproducibility, and practical deployment, significantly advancing the field of large model training. By standardizing core model operations and enabling dynamic data control, DataFlex paves the way for more efficient, adaptable, and scalable training pipelines, ultimately accelerating AI development and deployment.
Technical Contribution
The paper introduces a modular, pluggable architecture that unifies data selection, mixing, and reweighting strategies, with standardized interfaces for embedding, inference, and gradient operations. It supports large-scale training with frameworks like DeepSpeed ZeRO-3, enabling efficient integration of diverse algorithms. The system's design simplifies algorithm comparison and extension, fostering rapid experimentation and deployment in real-world scenarios.
Novelty
This is the first comprehensive system integrating multiple data-centric strategies—selection, mixing, and reweighting—within a single, standardized framework. The modular design and unified interfaces represent a significant leap forward, enabling seamless algorithm integration and fair evaluation, which was lacking in prior isolated implementations.
Limitations
- The current system primarily supports algorithms integrated during development; adding new strategies requires further engineering effort.
- In extremely large-scale environments, dynamic scheduling may encounter performance bottlenecks, needing further optimization.
- Heavy reliance on model-dependent signals may limit applicability to models or tasks with limited access to such signals.
Future Work
Future directions include expanding support for emerging data optimization algorithms, enhancing adaptive scheduling, and extending the framework to multi-modal and multi-task training scenarios. Incorporating reinforcement learning techniques for more intelligent data scheduling and exploring automated algorithm selection are also promising avenues.
AI Executive Summary
The rapid growth of large language models (LLMs) has revolutionized natural language processing, yet their training remains resource-intensive and data-dependent. Traditional methods rely on static datasets, which limit the potential for performance gains and efficiency. Recognizing this bottleneck, DataFlex offers a unified, flexible platform for data-centric dynamic training, integrating multiple strategies such as sample selection, domain mixing, and sample reweighting.
Built upon the robust LLaMA-Factory framework, DataFlex introduces a modular architecture with three core trainer types—SelectTrainer, MixTrainer, and WeightTrainer—each capable of incorporating specific algorithms like LESS, NICE, DoReMi, and ODM. This design standardizes core model operations, including embedding extraction, inference, and gradient computation, ensuring compatibility with large-scale training environments like DeepSpeed ZeRO-3. The system emphasizes extensibility, allowing researchers to easily add new algorithms and compare their effectiveness under controlled conditions.
Experimental validation demonstrates that dynamic data strategies consistently outperform static full-data training. For example, on MMLU tasks with Mistral-7B and Llama-3.2-3B, accuracy improved by 3-4 percentage points. In pretraining Qwen2.5-1.5B models, data mixture methods such as DoReMi and ODM yielded 2-3% accuracy gains and significant perplexity reductions, indicating better generalization. System optimizations also resulted in 10-15% faster training times, confirming the framework’s efficiency.
Overall, DataFlex bridges the gap between research and practical deployment, providing a standardized, scalable platform that enhances the effectiveness of data-centric training. It addresses key challenges in algorithm integration, reproducibility, and large-scale efficiency, fostering innovation in AI model development. Looking ahead, the framework will incorporate more algorithms, support multi-modal data, and leverage reinforcement learning for smarter data scheduling, further accelerating progress in AI research and applications.
Deep Dive
Abstract
Data-centric training has emerged as a promising direction for improving large language models (LLMs) by optimizing not only model parameters but also the selection, composition, and weighting of training data during optimization. However, existing approaches to data selection, data mixture optimization, and data reweighting are often developed in isolated codebases with inconsistent interfaces, hindering reproducibility, fair comparison, and practical integration. In this paper, we present DataFlex, a unified data-centric dynamic training framework built upon LLaMA-Factory. DataFlex supports three major paradigms of dynamic data optimization: sample selection, domain mixture adjustment, and sample reweighting, while remaining fully compatible with the original training workflow. It provides extensible trainer abstractions and modular components, enabling a drop-in replacement for standard LLM training, and unifies key model-dependent operations such as embedding extraction, inference, and gradient computation, with support for large-scale settings including DeepSpeed ZeRO-3. We conduct comprehensive experiments across multiple data-centric methods. Dynamic data selection consistently outperforms static full-data training on MMLU across both Mistral-7B and Llama-3.2-3B. For data mixture, DoReMi and ODM improve both MMLU accuracy and corpus-level perplexity over default proportions when pretraining Qwen2.5-1.5B on SlimPajama at 6B and 30B token scales. DataFlex also achieves consistent runtime improvements over original implementations. These results demonstrate that DataFlex provides an effective, efficient, and reproducible infrastructure for data-centric dynamic training of LLMs.