DataFlex: Dynamic LLM Data Optimization

• DataFlex is a unified framework for dynamic data-centric training of large language models that integrates sample selection, domain mixture, and sample reweighting.
• It replaces isolated implementations with modular, drop-in trainer abstractions compatible with LLaMA-Factory and DeepSpeed ZeRO-3, thus streamlining the training workflow.
• Empirical results indicate that dynamic data control enhances accuracy and training efficiency across multi-domain benchmarks and large-scale distributed environments.

Searching arXiv for the DataFlex paper and a few directly related methods referenced in the provided material. DataFlex is a unified framework for data-centric dynamic training of LLMs that treats training data not as a fixed stream but as a controllable optimization variable. Built on top of LLaMA-Factory and compatible with PyTorch’s DeepSpeed ZeRO-3, it provides a common infrastructure for three major paradigms of dynamic data optimization: sample selection, domain mixture adjustment, and sample reweighting. Its central contribution is to replace isolated implementations and inconsistent interfaces with extensible trainer abstractions and modular components that preserve the original training workflow while enabling online and offline data-centric control during optimization (Liang et al., 27 Mar 2026).

1. Motivation and problem setting

In conventional LLM pretraining, data are treated as a fixed resource: the entire corpus is shuffled and streamed through for a fixed number of epochs, domain proportions are predetermined, and every sample carries equal weight. The framework is motivated by the observation that such static pipelines cannot adapt to changing model needs. Early in training, the model may benefit from easier or more densely sampled data; later, it may benefit from rarer or harder examples. The stated consequence is that static pipelines can waste compute on low-utility examples and under-exploit high-utility ones (Liang et al., 27 Mar 2026).

Data-centric dynamic optimization seeks to close this loop by using signals from the model, including loss, gradient norms, validation feedback, and embeddings, to steer data usage on the fly. In the formulation adopted by DataFlex, the operative question is not only how to update model parameters, but also which samples should be seen, how heterogeneous corpora should be mixed, and how much each example should contribute to gradient updates (Liang et al., 27 Mar 2026).

The framework is positioned against a background of prior work distributed across multiple paradigms. For sample selection, the paper distinguishes offline methods such as NEAR and TSDS, which precompute scores in embedding space, from online methods such as LESS, NICE, loss-based selection, and $\Delta$-loss, which adapt selection throughout training. For data mixture optimization, it separates offline mixture methods such as DoReMi, DoGE, and REGMIX from online methods such as ODM, Aioli, and Adaptive Data Optimization. For sample reweighting, it notes loss-based schemes that assign per-sample coefficients $w_i \propto f(\ell_i)$, boosting hard examples while downweighting easy ones. The stated difficulty is that these implementations are scattered across repositories with inconsistent interfaces, which hinders reproducibility, fair comparison, and large-scale deployment (Liang et al., 27 Mar 2026).

2. System architecture and programming model

DataFlex is organized as a layered system. The base layer inherits from LLaMA-Factory and provides model loading, data tokenization, optimizer and scheduler setup, mixed-precision, and distributed support. The trainer layer replaces the standard trainer with three modes: SelectTrainer for dynamic sample selection, MixTrainer for dynamic domain mixture, and WeightTrainer for dynamic per-sample reweighting. The component layer loads pluggable selectors, mixers, or weighters through a centralized registry. Components expose a minimal interface, exemplified by methods such as .score(...) and .update(...), and may request model inference, embedding extraction, or gradient information (Liang et al., 27 Mar 2026).

Trainer	Function	Pluggable component
SelectTrainer	Dynamic sample selection	Selector
MixTrainer	Dynamic domain mixture	Mixer
WeightTrainer	Dynamic per-sample reweighting	Weighter

A central design property is drop-in compatibility with the original training workflow. Configuration extends LLaMA-Factory by adding a short dataflex section specifying the training type, component name, and schedule parameters such as warmup_step, update_step, and update_times. At the Python level, trainer instantiation is correspondingly minimal, with SelectTrainer(model, dataset, optimizer, selector_name='less', warmup_step=100, update_step=50, update_times=30) followed by trainer.train() given as the canonical example (Liang et al., 27 Mar 2026).

The framework’s unification is not only syntactic. It standardizes key model-dependent operations, including embedding extraction, inference, and gradient retrieval, into one code path. This reduces the implementation fragmentation that previously separated data selection, mixture optimization, and reweighting into distinct and often incompatible experimental stacks. A plausible implication is that methodological comparisons become less sensitive to engineering confounders, because identical scheduling and evaluation pipelines can be reused across paradigms.

3. Core algorithms and mathematical formulation

All three paradigms in DataFlex follow the same abstract pattern: at selected steps $t$, observe model state $\theta_t$ through losses, gradients, or embeddings; compute a control signal $c_t$; and apply that signal either to the data loader or to the loss function (Liang et al., 27 Mar 2026).

For sample selection, the framework instantiates multiple scoring rules. In the loss-based variant,

$s(x_i;\theta)=\ell(x_i;\theta),$

and the selected subset is the top-$k$ index set

$\{i_1,\dots,i_k\}=\arg\max_i s(x_i;\theta).$

In the $\Delta$-loss variant,

$s(x_i;\theta)=\ell_t(x_i)-\ell_{t-\Delta t}(x_i),$

which prioritizes samples whose loss is rising. In gradient-based selectors such as LESS and NICE, influence is approximated through gradient norms,

$w_i \propto f(\ell_i)$0

These selectors are called every $w_i \propto f(\ell_i)$1 steps, and the top-$w_i \propto f(\ell_i)$2 or weighted sample subsets feed the next $w_i \propto f(\ell_i)$3 training steps (Liang et al., 27 Mar 2026).

For domain mixture adjustment, the framework assumes $w_i \propto f(\ell_i)$4 domains with proportions $w_i \propto f(\ell_i)$5 and states the high-level objective as

$w_i \propto f(\ell_i)$6

where $w_i \propto f(\ell_i)$7 is the per-domain loss and $w_i \propto f(\ell_i)$8 is a regularizer, for example an entropic regularizer. In DoReMi, which is treated as an offline method, a proxy model measures excess loss $w_i \propto f(\ell_i)$9 relative to uniform sampling, and domain weights are updated by exponentiated gradient:

$t$0

followed by normalization such that $t$1. In ODM, treated as an online method, each domain is a multi-armed bandit arm and Exp3 is used:

$t$2

where $t$3 is the negative batch loss reward for domain $t$4, $t$5 is the bandit learning rate, and $t$6 is the normalization term (Liang et al., 27 Mar 2026).

For sample reweighting, per-sample weights are assigned according to current loss or gradient alignment:

$t$7

These weights are then applied multiplicatively in the loss:

$t$8

The framework presents this as the third major mechanism by which data usage can be dynamically controlled during optimization (Liang et al., 27 Mar 2026).

The paper describes the interaction layer that supports all of these methods as a unified data–model interface. Selectors and weighters invoke standardized methods to fetch embeddings or full gradients, while mixers sample from domain-tagged dataloaders using updated $t$9. This suggests that the main abstraction in DataFlex is not the individual algorithm but the control loop connecting model state to data policy.

4. Distributed implementation and scalability

DataFlex is designed for large-scale training settings and explicitly supports DeepSpeed ZeRO-3. When a selector or weighter requires full gradients or optimizer states, the framework invokes DeepSpeed’s safe_get_full_grad and safe_get_full_optimizer_state to reconstruct them from shards. To limit overhead, gradient-based operations run at configurable intervals determined by parameters such as warmup_step, update_step, and update_times, and decisions such as selected indices or weights are cached across multiple training batches (Liang et al., 27 Mar 2026).

The implementation emphasizes overhead minimization. It uses cached decisions, interval-based updates, and lightweight loss-only proxies when gradients are unnecessary. Component logic runs in parallel on each worker, and the system does not depend on a centralized coordinator. The reported distributed gradient and state reconstruction under ZeRO-3 allows gradient-based selectors to work on billion-parameter models across 8–32 GPUs (Liang et al., 27 Mar 2026).

Memory efficiency is addressed through mixed-precision, FlashAttention-2, and DeepSpeed ZeRO Stage-3, which the paper states keep GPU memory within bounds for Qwen2.5-1.5B at 30B tokens. At the same time, the paper identifies a practical challenge: frequent gradient reconstruction can be costly, so the choice of update frequency and caching strategy is crucial for efficiency. This is not presented as a theoretical limitation of dynamic data optimization itself, but as a systems constraint that directly shapes usable training schedules (Liang et al., 27 Mar 2026).

5. Empirical performance

The empirical study covers dynamic selection, dynamic mixture, and implementation efficiency. On dynamic selection, the benchmark setup uses an Open-Hermes subset with 100 K train examples and MMLU valid/test, evaluated with Mistral-7B and Llama-3.2-3B under LoRA PEFT with rank $\theta_t$0, $\theta_t$1, and 1 epoch. Dynamic selectors include LESS, NICE, Loss, and $\theta_t$2-Loss, with a reweighting baseline labeled “Reweight,” and all are compared against a static full-data baseline (Liang et al., 27 Mar 2026).

For Mistral-7B, the reported MMLU results are: LESS at 45.2% versus 39.4% for static training, corresponding to $\theta_t$3 percentage points; Reweight at 42.9%; TSDS at 42.9%; and NEAR at 41.9%. For Llama-3.2-3B, Reweight reaches 45.3% versus 31.9% for static training, corresponding to $\theta_t$4 percentage points; LESS reaches 45.0%; $\theta_t$5-Loss reaches 43.4%; and NICE reaches 42.8%. The paper states that offline methods underperform on the smaller model, highlighting the value of online, model-aware selection (Liang et al., 27 Mar 2026).

For dynamic mixture on SlimPajama pretraining, the experiments train Qwen2.5-1.5B from scratch at 6B and 30B token scales, using MMLU accuracy (5-shot) and corpus perplexity, both global and per-domain, as metrics. At Slim-Pajama-6B, the static baseline has accuracy 25.27% and perplexity 4.217; DoReMi reaches 25.84% with a gain of $\theta_t$6 percentage points and perplexity 4.134; ODM reaches 26.04% with a gain of $\theta_t$7 percentage points and perplexity 4.244. Per-domain, DoReMi is best on large domains (CC, C4), while ODM is best on niche domains (SE, ArXiv, Book). At Slim-Pajama-30B, the static baseline has accuracy 25.51% and perplexity 3.584; DoReMi reaches 25.97% with a gain of $\theta_t$8 percentage points and perplexity 3.562; ODM reaches 25.63% and perplexity 3.429, which is described as the largest perplexity improvement. ODM is reported to excel in 5/7 domains at 30B, showing that online exploration yields stronger gains given more training steps (Liang et al., 27 Mar 2026).

The efficiency comparisons are framed relative to original implementations. In Table 3, for LESS on a single GPU, DataFlex reduces runtime by 3.7–7.1% across sample ratios while matching or exceeding baseline accuracy. In 8-GPU scaling, runtime decreases from 28 734 s to 12 965 s, a 57.1% reduction, while accuracy increases from 42.4% to 43.0%, demonstrating multi-GPU scalability absent in original LESS. In Figure 1, the TSDS reimplementation is reported to be 1–3.5% faster than the original across training-set and validation-set scaling experiments, which the paper characterizes as making iterative selection convenient (Liang et al., 27 Mar 2026).

6. Reproducibility, usage guidance, and limitations

Reproducibility is a core design goal. The framework uses a modular registry in which new selectors, mixers, or weighters can be registered with two simple decorators and then invoked by name in the configuration. It is also presented as a drop-in replacement: users need only append a short dataflex section to existing LLaMA-Factory configurations, with no change to model, dataset, or optimizer definitions. The paper further states that unified evaluation, including an MMLU harness, identical scheduling parameters, and shared code paths guarantee fair comparison (Liang et al., 27 Mar 2026).

These design choices bear directly on methodological use. Because the same infrastructure supports sample selection, mixture adjustment, and reweighting, the framework reduces the degree to which empirical conclusions depend on separate codebases, divergent logging conventions, or incompatible distributed setups. This suggests that DataFlex functions not only as a training framework but also as an experimental control surface for data-centric LLM research.

The stated limitations are equally explicit. Current support covers only three paradigms. Extending the framework to multi-modal data, integrating reinforcement-learning signals, and designing adaptive update schedules are identified as future work. The paper therefore does not present DataFlex as a complete account of all possible data-centric optimization mechanisms; rather, it establishes a scalable and reproducible substrate for three major classes already active in the literature (Liang et al., 27 Mar 2026).