DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models

1. Problem Statement

The training of Large Language Models (LLMs) has shifted from a purely model-centric approach to a data-centric paradigm, where optimizing the selection, composition, and weighting of training data is as critical as optimizing model parameters. While numerous algorithms have been proposed for data selection (choosing useful samples), data mixture (adjusting domain proportions), and data reweighting (modulating sample importance), the current landscape suffers from significant fragmentation:

• Isolated Implementations: Most methods exist in separate, method-specific codebases with inconsistent interfaces and training protocols.
• Reproducibility Issues: The lack of a unified framework hinders fair comparison and reproducibility across different algorithms.
• Integration Barriers: Existing methods often require external pipelines or complex modifications to standard training loops, making them difficult to integrate into scalable, production-grade training infrastructures (e.g., those using DeepSpeed ZeRO-3).
• Missing Abstractions: There is no unified system to manage the shared, model-dependent operations required by these methods, such as embedding extraction, inference, and gradient computation.

2. Methodology: DataFlex Framework

DataFlex is a unified, drop-in replacement framework built upon LLaMA-Factory that treats data as a "first-class optimization variable." It unifies three major paradigms of data-centric dynamic training:

A. System Architecture

DataFlex replaces the standard training layer of LLaMA-Factory with a modular, three-layer architecture:

1. Base Layer: Inherited from LLaMA-Factory, handling model management, data processing, and optimization (e.g., DeepSpeed, FSDP).
2. Trainer Layer: Introduces three unified trainer abstractions corresponding to the three paradigms:
   • Select Trainer: Dynamically selects a subset of samples (e.g., LESS, NICE).
   • Mix Trainer: Dynamically adjusts mixture ratios across data domains (e.g., DoReMi, ODM).
   • Weight Trainer: Dynamically modifies per-sample weights during backpropagation (e.g., Loss-based reweighting).
3. Component Layer: A pluggable registry of algorithm-specific components (selectors, mixers, weighters) that share a common interface. This allows researchers to swap algorithms with minimal code changes.

B. Key Design Features

• Unified Interface: Supports both online methods (adapting decisions during training) and offline methods (precomputing strategies) under a single configuration format (YAML).
• Shared Operations: Standardizes model-dependent operations like embedding extraction, inference, and gradient computation, ensuring efficiency and compatibility.
• Scalability: Addresses the challenge of acquiring full parameter gradients in distributed settings (e.g., DeepSpeed ZeRO-3) by leveraging safe_get_full_grad to reconstruct gradients from partitioned shards.
• Drop-in Replacement: Users can switch from standard training to dynamic training by simply adding a dataflex section to their existing LLaMA-Factory configuration file, without altering model or optimizer settings.

3. Key Contributions

1. Conceptual Framework: Introduced the concept of a Data-Centric Dynamic Training System, providing a unified infrastructure for controlling data usage throughout the training lifecycle.
2. Modular Implementation: Developed a scalable system built on LLaMA-Factory that integrates data selection, mixture optimization, and reweighting into a single training paradigm with extensible, pluggable components.
3. Standardized Operations: Unified the management of shared model-dependent operations (embeddings, inference, gradients), facilitating efficient development and large-scale deployment.
4. Comprehensive Evaluation: Conducted extensive experiments covering seven data selection algorithms, two data mixture methods, and one reweighting method, demonstrating consistent improvements over static baselines.

4. Experimental Results

The authors evaluated DataFlex across various backbones (Mistral-7B, Llama-3.2-3B, Qwen2.5-1.5B) and datasets (Open-Hermes-2.5, SlimPajama).

A. Data Selection & Reweighting

• Performance: Dynamic methods consistently outperformed static full-data training.
  • On Mistral-7B, the gradient-based LESS method achieved 0.452 MMLU accuracy, beating the static baseline (0.394) by 5.8 percentage points.
  • On Llama-3.2-3B, the gap was even larger; the static baseline plateaued at 0.319, while dynamic methods (e.g., Reweight at 0.453) significantly surpassed it.
• Efficiency: DataFlex implementations were faster than original codebases. For the LESS algorithm, DataFlex reduced training time by 3.7% to 7.1% on single GPUs and achieved a 57.13% reduction in time when scaled to 8 GPUs (a capability absent in the original LESS implementation).

B. Data Mixture (Pretraining)

• Setup: Pretrained Qwen2.5-1.5B on SlimPajama (6B and 30B token scales) using DoReMi (offline) and ODM (online).
• Results:
  • MMLU Accuracy: Both dynamic methods improved accuracy over the static baseline. On the 6B scale, ODM reached 26.04% vs. baseline 25.27%.
  • Perplexity (PPL): DoReMi achieved the best overall PPL on the 6B scale (4.134), while ODM excelled on the 30B scale (3.429).
  • Domain Specificity: DoReMi improved performance on high-resource domains (CommonCrawl, C4), while ODM's online exploration mechanism effectively upweighted minority domains (StackExchange, GitHub, Book), leading to superior per-domain perplexity.

C. Efficiency

• Offline Selection: The DataFlex re-implementation of the TSDS algorithm showed consistent runtime improvements (1–3.5% faster) across varying dataset sizes compared to the original implementation, making it more suitable for iterative experimentation.

5. Significance

• Reproducibility & Fairness: DataFlex eliminates the fragmentation in data-centric research by providing a standardized environment, allowing for direct, fair comparisons between diverse algorithms (online vs. offline, gradient-based vs. loss-based).
• Practical Deployment: By integrating seamlessly with LLaMA-Factory and supporting large-scale distributed training (DeepSpeed ZeRO-3), DataFlex bridges the gap between theoretical data-centric algorithms and practical, large-scale LLM training.
• Infrastructure for Future Research: The modular design lowers the engineering barrier for researchers to implement and test new data-model interaction strategies, fostering further innovation in data-centric optimization.

In conclusion, DataFlex demonstrates that treating data as a dynamic optimization variable is not only theoretically sound but also practically achievable at scale, leading to significant gains in model performance and training efficiency.