Source-linked AI summary
DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models
Hao Liang, Zhengyang Zhao, Meiyi Qiang, Mingrui Chen, Lu Ma, Rongyi Yu, Hengyi Feng, Shixuan Sun, Zimo Meng, Xiaochen Ma, Xuanlin Yang, Qifeng Cai, Ruichuan An, Bohan Zeng, Zhen Hao Wong, Chengyu Shen, Runming He, Zhaoyang Han, Yaowei Zheng, Fangcheng Fu, Conghui He, Bin Cui, Zhiyu Li, Weinan E, Wentao Zhang
TL;DR
Data-centric methods are fragmented across inconsistent implementations, limiting reproducibility, fair comparison, and integration. DataFlex unifies selection, mixture optimization, and reweighting through modular trainers built on LLaMA-Factory, and its implementations improve benchmark performance and runtime over static or original baselines. The framework therefore provides infrastructure for reproducible and practical dynamic data optimization.
Problem
Inconsistent interfaces, protocols, and codebases across data-centric methods hinder reproducibility, fair comparison, and practical deployment.
Method
DataFlex integrates data selection, mixture optimization, and reweighting through modular trainer abstractions and standardized model-dependent operations within LLaMA-Factory.
Results
DataFlex methods consistently outperform static full-data training on reported performance measures and improve runtime over original implementations.
Takeaways & Limitations
DataFlex provides a reproducible infrastructure for systematically comparing and deploying data-centric dynamic training methods.
Takeaways & Limitations
LESS and DataFlex are restricted to a single-GPU setting for fair comparison because of LESS’s implementation constraints.
Abstract
from arXiv · showhide
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.
1 Introduction
DataFlex addresses fragmented data-centric training implementations with a unified framework that integrates selection, mixture optimization, and reweighting into standard LLM workflows. Its modular abstractions support reproducible comparison, scalable implementation, and dynamic control of data usage.
- Existing data-centric methods use inconsistent interfaces, training protocols, and codebases, hindering reproducibility, fair comparison, and practical deployment.
- DataFlex treats data as an optimization variable by dynamically controlling sample selection, source mixing, and sample weighting across online and offline methods.
- Built on LLaMA-Factory, DataFlex replaces the training layer while preserving compatibility with existing model management, data processing, and optimization pipelines.
- Three trainer abstractions—Select Trainer, Mix Trainer, and Weight Trainer—map to selection, mixture optimization, and reweighting paradigms.
- Unified operations for embedding extraction, inference, and gradient computation support reproducible development and large-scale training.
- DataFlex methods consistently outperform static baselines while improving model performance and training efficiency across data-centric paradigms.
2 Related Work
Related work optimizes which data are selected, how domains are mixed, and how samples are weighted, using offline or online decisions. Despite these methods, the field lacks a unified framework for systematic integration and management.
- Data Selection: Data selection identifies useful training subsets to improve efficiency, reduce computational cost, and enhance downstream performance.
- Data Selection: Offline selection determines subsets before training using signals such as data quality, diversity, or difficulty.
- Data Selection: Online selection updates policies during training from current-model feedback, including gradient-based influence and black-box optimization signals.
- Data Mixture: Data mixture allocates training probability across heterogeneous domains, with offline methods estimating weights beforehand and online methods adapting proportions during training.
- Data Reweighting: Online reweighting changes each sample’s importance from current loss, assigning higher weights to difficult or informative samples and lower weights to easy or redundant ones.
- Framework Gap: Existing data–model interaction algorithms lack a unified, organized, and scalable framework for systematic management and integration.
3 DataFlex: System Design and Core Abstractions
DataFlex makes data a first-class optimization variable within a unified, compatible, and extensible LLaMA-Factory-based training framework. Its trainer–component architecture supports dynamic selection, mixing, and reweighting while standardizing model-dependent operations and scalable execution.
- Goals and Design Philosophy: DataFlex dynamically controls sample selection, source mixture proportions, and per-sample contributions to gradient updates.
- Compatibility: DataFlex replaces the LLaMA-Factory training layer while preserving model management and optimization components, requiring only lightweight data-loading extensions for mixture optimization.
- Model-Dependent Operations: The framework unifies embeddings, inference outputs, and gradients, including full-gradient reconstruction compatible with DeepSpeed ZeRO-3.
- Trainer Layer: Its three trainer modes—Select, Mix, and Weight—map respectively to dynamic sample selection, domain mixing, and per-sample weighting.
- Component Layer: Pluggable selectors, mixers, and weighters encapsulate algorithm-specific logic behind common interfaces for controlled comparison and extension.
- Unified Trainer–Component Interaction: A standardized trainer–component interaction observes model state, emits data-centric control signals, and invokes components at configurable or per-step frequencies.
- Configuration and Efficiency: Configuring dynamic training requires appending a dataflex section to the existing YAML workflow, while interval scheduling, caching, proxy signals, and distributed execution reduce overhead.
4 Experiments
Experiments evaluate DataFlex across data selection, domain mixture optimization, and runtime efficiency using multiple models, datasets, baselines, and training scales. The study includes seven selection algorithms, two mixture algorithms, and one reweighting algorithm, with mixture experiments spanning 6B and 30B tokens.
- Experimental Scope: Seven data selection algorithms, two data mixture algorithms, and one data reweighting algorithm are evaluated entirely within DataFlex.
- Data Selection: Data selection uses 100,000 Open-Hermes-2.5 examples and evaluates MMLU validation and test splits on Mistral-7B-v0.1 and Llama-3.2-3B.
- Data Selection: Selection comparisons include online and offline methods, reweighting, random sampling, and static full-data training.
- Data Mixture: SlimPajama mixture experiments use seven domains and randomly sampled 6B-token and 30B-token subsets preserving the corpus’s natural proportions.
- Data Mixture: DoReMi trains reference, proxy, and target Qwen2.5 models, then uses optimized domain weights for target-model pretraining against a static default-proportion baseline.
- Data Mixture: In the 30B setting, DoReMi changes CommonCrawl from 54.1% to 34.1% and C4 from 28.7% to 33.6% relative to default proportions.
- Efficiency: Runtime evaluation compares DataFlex implementations of online LESS and offline TSDS with their original codebases using Open-Hermes-2.5 subsets of up to 100,000 examples.
- Efficiency: Offline TSDS scoring encodes tokenized samples into sentence-level embeddings before nearest-neighbor retrieval and KDE-based density estimation.
4.2 Data Selection and Reweighting
Dynamic data methods generally outperform static full-data training, with method and model size shaping the strongest results. DataFlex’s mixture experiments show complementary gains in accuracy and perplexity across 6B- and 30B-token scales.
- Dynamic methods generally outperform the static full-data baseline, with the strongest gains appearing for LESS and loss-based reweighting.
- Mistral-7B: 0.452 MMLU accuracy: LESS exceeds the Mistral-7B static baseline of 0.394 by 5.8 percentage points.
- Llama-3.2-3B: 0.453 MMLU accuracy: Reweight leads Llama-3.2-3B, while all online methods exceed the static baseline of 0.319.
- Data mixture: On SlimPajama-6B, ODM reaches the highest MMLU accuracy at 26.04%, while DoReMi achieves the best overall log perplexity at 4.134.
- Data mixture: DoReMi favors high-resource domains, whereas ODM more strongly explores minority domains and gains greater overall perplexity advantages at 30B tokens.
- Data mixture: At 30B tokens, DoReMi achieves the highest MMLU accuracy at 25.97%, while ODM produces the best overall log perplexity at 3.429.
4.4 Efficiency of DataFlex
DataFlex reduces runtime relative to original data-selection implementations across sampling ratios and dataset scales. Its distributed configuration further demonstrates substantial scalability.
- DataFlex consistently achieves lower training latency than LESS while maintaining or improving MMLU accuracy across sampling ratios.
- LESS comparison: 7.09% runtime reduction: DataFlex saves over 1,000 seconds at a 0.5 sampling ratio in the single-GPU comparison.
- Distributed scalability: 57.13% time reduction: DataFlex lowers runtime from 28,734 s on one GPU to 12,965 s on 8×H20 GPUs at the 1.0 ratio.
- TSDS comparison: 1–3% improvement: DataFlex is faster than TSDS across training-set scales from 5,000 to 100,000 with validation fixed at 1,000.
- TSDS comparison: 1.5–3.5% improvement: DataFlex is faster than TSDS as validation queries increase from 50 to 1,000 with training fixed at 10,000.
- The consistent, moderate TSDS speedup makes DataFlex more suitable for repeated selection experiments under changing configurations.
5 Conclusion
DataFlex unifies data selection, mixture optimization, and reweighting within LLaMA-Factory while making shared model-dependent operations reusable. Experiments report gains over static training and original implementations.
- DataFlex treats data as a first-class optimization object and unifies selection, mixture, and reweighting under one framework.
- Extensible trainers and modular components replace LLaMA-Factory’s training layer while preserving integration with existing workflows.
- Unified embedding extraction, inference, and gradient computation support scalable and reusable dynamic-training implementations.
- Experiments report higher MMLU accuracy for dynamic selection, improved MMLU accuracy and corpus-level perplexity for mixture optimization, and runtime gains over original codebases.
A Author Contributions
The author contributions span project leadership, system and algorithm design, implementation, experimentation, writing, and supervision.
- Hao Liang led the project, algorithmic and system design, large-scale data selection extensions, and manuscript preparation.
- Zhengyang Zhao founded the project and designed and implemented the data selection and reweighting interface.
- Meiyi Qiang contributed to NICE, improved LESS, and led selector and reweighting experiments and analysis.
- Mingrui Chen designed and implemented data mixture interfaces and contributed to mixture methods and experiments.
- Lu Ma and Rongyi Yu implemented the first LESS selector and offline data-selection algorithms, respectively.
- Multiple contributors handled experiments, manuscript revision and writing, system development, and project supervision.
B Comparison with Original Implementations
This section compares DataFlex with the original LESS and DoReMi codebases at the engineering level. It emphasizes DataFlex as a unified reimplementation while acknowledging the pioneering algorithmic contributions of the earlier works.
- DataFlex is compared with the original LESS and DoReMi implementations on practical engineering differences.
- The comparison focuses on engineering improvements provided by DataFlex as a unified reimplementation.
- The discussion explicitly preserves recognition of LESS and DoReMi as pioneering works that introduced important algorithmic ideas.
B.1 Data Selection
DataFlex addresses practical limitations in the original LESS-based data-selection implementations by supporting distributed training, improving compatibility, and consolidating fragmented workflows.
- DataFlex extends single-GPU LESS implementations with distributed training support through DeepSpeed ZeRO and FSDP.It supports gradient capture with partitioned parameters, enabling reconstruction of full-rank gradients and Adam optimizer states.
- DataFlex addresses dependency and model-compatibility challenges caused by pinned versions and hard-coded internal APIs.The original codebase may encounter compilation errors or incompatibilities when dependencies and model architectures change.
- DataFlex consolidates data preparation, gradient extraction, and training into a unified pipeline with a single configuration entry point.This reduces setup effort and supports experiment reproduction across environments.
B.2 Data Mixture
DataFlex improves the official DoReMi implementation while preserving its core algorithmic logic, adding scalable distributed training, decoupled model compatibility, and a more flexible data pipeline.
- DataFlex introduces these implementation improvements while maintaining consistent core DoReMi algorithmic logic.
- DataFlex adds native multi-node and multi-GPU support through integration with large-scale distributed training frameworks such as ZeRO and FSDP.The implementation targets stable execution across distributed infrastructures.
- DataFlex abstracts DoReMi as a pluggable Mixer component, decoupling data preprocessing from model interfaces.This improves compatibility with evolving open-source models, including Qwen and LLaMA.
- DataFlex enables on-the-fly data integration and sampling adjustments through a unified mixture_manager and dynamically reconstructed distributed dataloaders.This replaces strict offline preprocessing and fragmented script modifications for new datasets.