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

arXiv:2603.26164 · cs.LG, cs.CL · Submitted 2026-03-27 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models".

Tom: DataFlex presents a unified data-centric dynamic training framework built upon LLaMA-Factory, addressing the fragmentation and reproducibility issues in existing data selection, mixture optimization,

Jane: First, who's behind it and why it matters.

Title and authors: Tom: Let's talk about the title and who wrote this, "DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models." It’s a mouthful, but it really captures the essence of what they’re proposing.

Jane: They are essentially showing how to bring together data selection, mixture optimization, and sample reweighting into one cohesive system. It's about giving researchers tools that work together seamlessly instead of having them juggle separate pipelines for each task.

Lu: The authors are from Peking University and various other places, which suggests a really broad perspective on this problem; they aren't just looking at one niche area but trying to solve the general problem of data-centric training inconsistency.

Meng: That breadth is important because if it’s truly unified, it should mean that the underlying model operations—like embedding extraction or gradient computation—are handled through a single interface, which simplifies integration into existing infrastructure.

Lalam: For us, having a clear framework like this makes deploying advanced training techniques much safer; we can trust the results because they follow a standardized structure.

The paper's summary: Tom: Moving on to what DataFlex actually does, the paper summarizes it as a system that treats data selection, domain mixture adjustment, and sample reweighting as first-class optimization variables. That’s the core idea we need to grasp here.

Jane: That means instead of just picking a fixed dataset once at the start, you can dynamically change *what* data is used, *how* different sources are mixed during pretraining, or even adjust how much influence each individual training sample has on the final model weights.

Lu: The paper details three specific trainer abstractions: Select Trainer for dynamic subset selection, Mix Trainer for adjusting domain proportions, and Weight Trainer for modifying per-sample contributions during backpropagation. This structure is what makes it so flexible.

Meng: So it’s not just one algorithm; it’s a set of modular components where you can plug in different selectors or mixers without having to rebuild the entire training layer from scratch, which sounds quite efficient for testing new ideas.

Lalam: It really means we get to experiment with these dynamic strategies systematically, and the framework ensures that whatever strategy we choose is compatible with the larger system.

The paper's improvements: Tom: Now, let's talk about what they improved over existing methods. They focused heavily on solving the fragmentation issue by providing unified abstractions and allowing new algorithms to be added as self-contained components through a centralized registry.

Jane: That centralization is key; it means a researcher doesn't have to write custom code every time they want to try a new data selection technique, they just plug their component into the system. It’s about extensibility without breaking things.

Lu: A specific improvement mentioned is that DataFlex standardizes shared model-dependent operations, like gradient computation and embedding extraction, making it suitable for large-scale training settings like DeepSpeed ZeRO-three. That addresses a huge practical hurdle in distributed computing.

Meng: That standardization is what makes it viable for production; if the framework handles the complex parts of distributed gradient collection via interfaces like "safe get full grad," then we can actually run these dynamic strategies on our multi-node setups without massive headaches.

Lalam: It simplifies things immensely because we don't have to worry about mismatched code when scaling up; it’s designed to work within the existing LLaMA-Factory ecosystem while adding this powerful data control layer.

Conclusion: Tom: Wrapping up our discussion on "DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models," the main implication is that we can now systematically study and deploy dynamic training strategies in a way that was previously inconsistent and hard to reproduce.

Jane: The paper shows how treating data as an optimization variable allows us to dynamically select, mix, or reweight data during the entire training process, which should lead to better model performance than relying on static datasets alone.

Lu: The real potential lies in the ability for researchers to compare these different dynamic methods fairly because they all use the same interface and structure within DataFlex.

Meng: For practical impact, this framework makes it easier to run complex data-centric algorithms at scale across distributed training environments, which is a necessity for training the next generation of large models efficiently.

Lalam: I think the biggest cultural shift here is moving away from fixed pipelines toward a more adaptable approach where we can quickly inject new data optimization techniques into our workflow whenever they emerge.

Tom: That’s what it is, folks; DataFlex gives us a clear, unified path forward for optimizing LLMs based on the quality and composition of the data used during training.

Peking University · Institute for Advanced Algorithms Research, Shanghai · OriginHub Technology

cs.LG, cs.CL

Submitted: 2026-03-27

Updated: 2026-09-30

Code: https://github.com/OpenDCAI/DataFlex

Project page: https://opendcai.github.io/DataFlex-Doc

Importance score: 91/100

The gist: DataFlex presents a unified data-centric dynamic training framework built upon LLaMA-Factory, addressing the fragmentation and reproducibility issues in existing data selection, mixture optimization,

Key concepts

Data-Centric Dynamic Training
This approach treats the data itself as something that can be actively optimized during the model's training process, rather than just a fixed input. Instead of using one set of data for everything, DataFlex allows the system to dynamically change which samples are used, how different data sources are mixed, and how important each specific sample is in real-time.
Unified Trainer Abstractions
The framework provides three distinct interfaces: Select Trainer, Mix Trainer, and Weight Trainer. These abstractions allow users to plug in different algorithms for selecting data subsets (like LESS), adjusting domain mixtures (like DoReMi), or modifying sample weights. This modular design ensures compatibility across various training paradigms.
First-Class Optimization Variable
This philosophy means that the data is not just a passive input but an active variable that the training process can manipulate. By treating data selection, mixing, and reweighting as optimization tasks, DataFlex enables dynamic control over the data usage throughout the entire training lifecycle to achieve better results.
Gradient Collection Mechanism
DataFlex standardizes how model-dependent operations like gradient computation work in distributed settings. It uses a specific mechanism compatible with DeepSpeed ZeRO-3 to handle gradient acquisition, even when model parameters are partitioned across multiple devices, ensuring scalability for large models.

Terminology

Summary

DataFlex presents a unified data-centric dynamic training framework built upon LLaMA-Factory, addressing the fragmentation and reproducibility issues in existing data selection, mixture optimization, and reweighting methods for Large Language Models (LLMs). By treating data as a first-class optimization variable and providing unified abstractions for these paradigms—data selection, domain mixture adjustment, and sample reweighting—DataFlex enables researchers to systematically study and deploy dynamic data-centric training strategies in a reproducible manner. This framework is crucial because it unifies common model-dependent operations like embedding extraction, model inference, and gradient computation under a single interface, making complex data-centric methods scalable for large-scale training settings.

Framework Design and Philosophy

DataFlex is conceived as a Data-Centric Dynamic Training System that treats data as a first-class optimization variable. Its design is guided by three core principles: unification, compatibility, and extensibility. The framework replaces the training layer of LLaMA-Factory with extensible trainer abstractions and modular algorithm components rather than introducing an external pipeline. This ensures that DataFlex serves as a drop-in replacement for the training layer, preserving compatibility with existing model management and optimization pipelines while enabling dynamic control over data usage throughout the training lifecycle.

Unified Trainer Abstractions

The core of DataFlex is its modular architecture, which introduces three distinct trainer abstractions corresponding to the major paradigms:

  1. Select Trainer: Dynamically selects a subset of samples according to a specified strategy.

  2. Mix Trainer: Dynamically adjusts mixture ratios across domains or data sources during training.

  3. Weight Trainer: Dynamically modifies per-sample training weights during backpropagation.

Each trainer is coupled with pluggable algorithm components—selectors, mixers, and weighters—managed through a centralized registry, allowing new algorithms to be implemented as self-contained components without modifying the rest of the system. This design facilitates a unified interface for diverse methods, regardless of whether they are online or offline.

Algorithm Integration and Extensibility

DataFlex supports a wide range of data-centric methods across its three paradigms:

Data Selection: Supports gradient-based (LESS, NICE), loss-based (Loss, Delta Loss), and distribution-based methods (NEAR, TSDS).

Data Mixture: Includes offline DoReMi and online ODM. These methods allow for dynamic adjustment of domain proportions during pretraining.

Data Reweighting: Provides a loss-based weighter with multiple strategies to dynamically adjust the importance of each training sample based on its current loss.

The system standardizes shared model-dependent operations, including embedding extraction, model inference, and gradient computation, making it suitable for large-scale training settings such as DeepSpeed ZeRO-3.

System Efficiency and Scalability

DataFlex inherits support for mixed-precision training and distributed data parallelism from LLaMA-Factory. It addresses the challenge of gradient acquisition in distributed settings by leveraging a distributed gradient collection mechanism compatible with DeepSpeed ZeRO-3, utilizing interfaces like safe get full grad to reconstruct full gradients from partitioned shards. To reduce overhead, DataFlex executes operations at configurable intervals and caches decisions, ensuring scalability to multi-node, multi-GPU settings. For instance, the framework demonstrates a 57.13% reduction in time when leveraging 8×H20 GPUs for data selection compared to single-GPU implementations.

Experimental Validation

Comprehensive experiments confirm the efficacy of DataFlex:

** Data Selection: Dynamic methods consistently outperform static full-data training on MMLU across both Mistral7B and Llama-3.2-3B backbones. Online methods, such as LESS, often achieve the highest final accuracy on certain models. For example, LESS achieved 0.452 accuracy on Mistral-7B.**

** Data Mixture: Methods like DoReMi and ODM improve both MMLU accuracy and corpus-level perplexity over default data proportions when pretraining Qwen2.5-1.5B on SlimPajama at various token scales (6B and 30B). ODM showed a stronger per-domain perplexity on specialized text types like ArXiv, GitHub, and Book.**

** Overall, the results validate that dynamic data optimization provides consistent improvements in both model performance and training efficiency compared to static baselines.**

Comparison with Original Implementations

DataFlex highlights significant engineering improvements over original codebases. For instance, in Data Selection (LESS), DataFlex extends it to support gradient capturing when model parameters are partitioned across devices, enabling the selector to reconstruct full-rank gradients from sharded parameters—a capability absent in the original single-GPU implementation.

Improvements for AI systems

Based on the research presented in DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models, here are specific improvements that can be made to AI systems, detailing what those improved systems will be able to do:

The core improvement is the transition from static, fixed-dataset training pipelines to a dynamic, data-centric optimization framework.

Here are the specific improvements and capabilities:

  1. Enhance Model Generalization and Robustness through Dynamic Data Selection:

  2. Improve Training Efficiency and Reduce Computational Costs via Adaptive Sample Utilization:

  3. Optimize Multi-Domain Language Modeling Performance via Dynamic Corpus Composition:

  4. Enable Reproducible, Scalable, and Fair Comparison of Data-Centric Algorithms:

Detailed Specific Improvements and System Capabilities:

Detailed Specific Improvements and System Capabilities (Specific Actions):

  1. The system can dynamically identify and prioritize the most informative or high-utility data samples during the training loop (using methods like LESS or NICE) rather than relying on a fixed dataset, leading to models that are more robust to variations in real-world input distributions.

  2. The system can reduce training time and memory consumption by intelligently deciding which specific data points to use at each step, potentially allowing for faster convergence or enabling the training of larger models (e.g., Llama-3.2-3B) with better performance than static full-data training suggests.

  3. The system can automatically adjust the proportion of different data sources (e.g., web text vs. books vs. code) in real-time based on the model's current performance or loss observations, leading to specialized models that excel in specific sub-domains (e.g., better performance on ArXiv papers or StackExchange questions).

  4. The framework can serve as a universal testing ground where researchers can plug and play any new data selection, mixing, or reweighting algorithm without rewriting the entire training pipeline, significantly lowering the barrier to entry for novel data-centric research. Furthermore, by unifying model-dependent operations (like embedding extraction and gradient computation) under one interface, it ensures that complex algorithms work reliably across different large-scale distributed training environments (like DeepSpeed ZeRO-3).

Sources

Related papers