DataFlex：一个面向大语言模型数据中心动态训练的统一框架 DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models
统一数据选择、混合和重加权的动态训练框架
前置知识
数据中心训练
数据中心训练是一种优化方法，不仅优化模型参数，还优化训练数据的选择、组合和加权。传统训练将数据视为静态资源，而数据中心训练将数据视为可动态调整的优化变量，通过智能选择高价值样本、调整不同领域数据的混合比例、以及根据样本难度动态调整权重，来提升训练效率和模型性能。
本文的核心贡献就是统一实现数据中心训练的三大范式，理解这一概念是把握论文创新点的基础。
动态数据选择
动态数据选择是在训练过程中根据当前模型状态动态调整使用哪些样本的方法。在线方法如LESS通过梯度近似估计每个样本对目标目标的影响，NICE通过黑盒优化处理不可微评估指标。这些方法相比离线预选择更具适应性，能根据模型学习演化优先级处理样本。
DataFlex统一实现了多种数据选择算法，理解其工作原理有助于评估框架的覆盖范围和实用性。
LoRA微调
LoRA是一种参数高效的微调方法，通过在模型权重矩阵上添加低秩分解矩阵来实现适应。具体地，对于预训练权重W，LoRA引入低秩矩阵A和B，使得前向计算变为h = Wx + BAx，其中BAx是可训练部分。本文实验中设置rank r=32，缩放因子alpha=64。
论文的所有数据选择和重加权实验都使用LoRA进行参数高效微调，理解这一技术有助于把握实验设置。
DeepSpeed ZeRO-3
ZeRO是DeepSpeed框架中的内存优化技术，通过分割优化器状态、梯度和参数来减少显存占用。ZeRO-3是最激进级别，将参数、梯度和优化器状态都分割到不同GPU上。这对梯度获取提出了挑战，因为完整梯度在模型并行时通常不会物化。DataFlex通过safe_get_full_grad接口实现兼容ZeRO-3的分布式梯度收集。
DataFlex的一个重要技术贡献就是解决分布式训练中的梯度获取问题，这对理解其可扩展性至关重要。
多臂赌博机算法
多臂赌博机是序贯决策的经典问题，算法需要在不同手臂间平衡探索和利用。Exp3是一种适用于对抗性设置的多臂赌博机算法，为每个臂维护权重并根据反馈更新。ODM使用Exp3进行在线数据混合，将每个领域视为一个臂，根据观察到的损失调整领域权重，指数加权移动平均实现平滑更新。
ODM是DataFlex实现的在线数据混合方法之一，理解其背后的多臂赌博机原理有助于把握该方法的决策机制。
研究动机
现有数据选择、数据混合优化和数据重加权方法通常在隔离的代码库中开发，接口不一致，阻碍了可复现性、公平比较和实际集成。这些方法大多作为特定算法的存储库发布，具有不一致的接口、异构的训练和评估协议，以及通常过时的代码库。例如，原始LESS实现只能单GPU执行，不原生支持多GPU并行，且依赖特定版本；DoReMi官方代码缺乏原生多节点训练支持，存在数据预处理和模型接口的结构性依赖。这种碎片化使得系统性评估、可复现实验和实际部署变得困难。
本文的目标是提出一个统一的数据中心动态训练框架，将数据选择、数据混合优化和数据重加权集成到一个训练范式中，实现系统性比较和实际部署。框架应与现有大规模模型训练基础设施无缝集成，而非引入孤立的工作流，同时允许研究者以最小的工程开销实现和比较新的数据中心算法。标准化共享的模型相关操作，包括嵌入提取、模型推理和梯度计算，支持可扩展的分布式训练。
与已有工作不同的是，提出数据中心动态训练系统的概念，将数据视为一等优化变量，提供统一接口来动态控制使用哪些样本、如何混合不同数据源以及每个样本对优化的贡献程度。现有方法通常被设计为孤立的算法组件，而非通用数据模型交互系统的一部分。许多方法都需要访问模型相关信号，没有统一抽象来管理这些交互，将多样化的数据中心训练算法集成到可扩展工作流中仍然具有挑战性。
核心方法
DataFlex采用三层模块化架构：Base Layer继承LLaMA-Factory提供模型管理、数据处理和优化基础设施；Trainer Layer用三种动态训练模式替换原始LLaMA-Factory训练器：Select Trainer根据策略动态选择样本子集，Mix Trainer动态调整领域或数据源的混合比例，Weight Trainer动态修改反向传播时的每样本训练权重；Component Layer为每种训练器关联可插拔策略组件。这些组件封装算法特定逻辑的同时共享训练框架内的公共接口，通过中心化注册表管理，新算法可通过装饰器注册并在运行时自动发现和实例化。框架设计故意轻量级：不在LLaMA-Factory上引入外部编排层，而是替换训练层，仅在其他地方引入最小扩展。
核心创新是提出模块化数据模型交互架构，将三种主要的数据中心优化范式统一到公共训练框架下，提供可插拔的算法组件接口。统一trainer-component交互模式：每种方法观察当前模型状态，计算数据中心决策，反馈到后续优化步骤。Select Trainer和Mix Trainer在可配置间隔调用组件，Weight Trainer在预热阶段后的每个训练步骤应用每样本权重。这种统一抽象将算法特定逻辑与训练流水线解耦。标准化共享模型相关操作，通过safe_get_full_grad接口实现兼容DeepSpeed ZeRO-3的分布式梯度收集。
方法步骤详情
使用流程分为配置和执行两步。配置方面，YAML文件复用LLaMA-Factory格式指定模型、数据集和训练超参数，仅添加dataflex部分指定数据中心策略和调度：train_type选择训练范式，component_name选择具体算法，调度参数控制数据中心策略在训练期间的更新频率。执行方面，通过dataflex-cli train命令启动。具体到不同范式：数据选择时，训练池由Open-Hermes-2.5的10万样本组成，使用LoRA进行参数高效微调，训练一个epoch，AdamW优化器，余弦学习率调度，预热比例0.1，学习率5.0乘以10的负7次方，全局批量大小8；数据混合时，在SlimPajama的6B和30B token子集上进行预训练，从随机初始化开始训练，使用DeepSpeed ZeRO-3进行内存高效分布式训练，FlashAttention-2进行高效注意力计算。DoReMi采用三步过程：Step 1用默认领域比例训练参考模型，Step 2用DoReMi算法训练代理模型，计算相对于Step 1参考模型的每领域超额损失，通过指数梯度上升更新领域权重，Step 3用Step 2最终优化的领域权重作为静态混合训练目标模型。ODM在单次训练传递中使用Exp3多臂赌博机算法和指数加权移动平均动态调整领域权重。
技术新颖性
技术新颖性体现在四个方面：一是系统设计层面，不引入孤立外部管道，而是用可扩展训练器抽象和模块化算法组件替换LLaMA-Factory训练层，直接将数据中心策略集成到训练循环中；二是接口统一层面，提供标准化trainer-component交互模式，将算法特定逻辑与训练流水线解耦，新算法可作为自包含组件实现并通过注册表注册，无需修改系统其余部分；三是操作标准化层面，统一管理数据中心方法共享的模型相关操作，包括嵌入提取、模型推理和梯度计算，使可复用实现和大规模训练部署变得容易；四是分布式支持层面，实现兼容DeepSpeed ZeRO-3的分布式梯度收集机制，通过safe_get_full_grad接口从分片参数重构完整梯度。相比原始实现，DataFlex在保持一致核心算法逻辑的同时引入了几个关键工程改进：原生支持多GPU并行，不锁定特定依赖版本，统一配置入口，原生支持多节点训练，解耦架构，灵活数据管道。
实验结果
数据选择实验在Mistral-7B和Llama-3.2-3B两个模型上使用Open-Hermes-2.5的10万样本子集进行，在MMLU数据集上评估。在Mistral-7B上，LESS达到最佳最终准确率0.452，比静态基线0.394提升5.8个百分点，Reweight排名第二0.429，离线方法TSDS和NEAR紧随其后，NICE达到0.418，Delta Loss达到0.412，Loss-only selector和Random selector与静态基线相当。离线方法在早期训练收敛更快，因为其预计算选择从一开始就与高价值样本对齐。在Llama-3.2-3B上，动态方法与静态基线的差距更加明显，静态基线仅达到0.319，而所有在线方法都超过0.427。Reweight达到最佳最终准确率0.453，LESS紧随其后0.450，Delta Loss达到0.434，NICE达到0.428，Random达到0.431，Loss-only selector达到0.429，离线方法NEAR和TSDS表现明显较差，表明当模型容量有限时，动态模型感知选择更为关键。数据混合实验在SlimPajama的6B和30B token子集上进行，使用MMLU准确率和语料库级别对数困惑度评估。在SlimPajama-6B上，两种动态方法在MMLU上都优于静态基线，ODM达到最高准确率26.04%，DoReMi 25.84%，基线25.27%。在整体对数困惑度方面，DoReMi达到最佳分数4.134，优于基线4.217和ODM 4.244。每领域细分显示两种动态方法具有互补优势：DoReMi在CommonCrawl和C4达到最低困惑度，ODM在StackExchange、ArXiv和Book达到最佳困惑度。在SlimPajama-30B上，扩展到30B token放大了动态数据混合的好处。DoReMi达到最高MMLU准确率25.97%，ODM 25.63%，基线25.51%。在整体对数困惑度方面，ODM产生最佳结果3.429，显著优于DoReMi 3.562和基线3.584。效率评估显示，在0.05采样率下，DataFlex比LESS减少3.72%运行时间，在0.5比率下增加到7.09%，节省超过1000秒。在1.0比率下的8GPU设置中，DataFlex将运行时间从28734秒减少到12965秒，时间减少57.13%。离线数据选择比较显示，DataFlex实现在所有测试规模下都更快，对应约1-3%的稳定改进。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数据选择（Mistral-7B） | MMLU准确率 | 0.452 (LESS) | 0.394 (静态全数据) | 提升5.8个百分点 |
| 数据选择（Llama-3.2-3B） | MMLU准确率 | 0.453 (Reweight) | 0.319 (静态全数据) | 提升13.4个百分点 |
| 数据混合（SlimPajama-6B） | MMLU准确率 | 26.04% (ODM) | 25.27% (默认比例) | 提升0.77个百分点 |
| 数据混合（SlimPajama-30B） | MMLU准确率 | 25.97% (DoReMi) | 25.51% (默认比例) | 提升0.46个百分点 |
| 数据混合（SlimPajama-30B） | 整体困惑度 | 3.429 (ODM) | 3.584 (默认比例) | 降低0.155 |
| 训练效率 | 训练时间减少 | 3.72%-7.09% (单GPU), 57.13% (8GPU) | 原始LESS实现 | 显著提升 |
局限与改进
作者承认框架当前主要关注文本数据，尚未扩展到多模态场景。从实验观察来看，离线数据选择方法在较小模型上表现明显不如在线方法，这表明离线方法可能无法适应有限的模型容量，当计算预算限制只能使用较小模型时，框架需要更强的模型感知能力。实验仅在特定数据集和模型架构上验证，框架在其他数据域和模型架构上的泛化能力尚未充分验证。虽然框架支持分布式训练，但梯度收集机制仍可能在大规模设置下引入额外开销，特别是在更新频率很高时。DoReMi的三步过程需要训练额外的代理模型，增加了计算成本，这在计算预算受限时可能限制其实用性。
独立分析的弱点
离线数据选择方法的可扩展性存在限制，当验证集很大时，嵌入提取和最近邻检索的内存和计算开销可能成为瓶颈，特别是在训练集达到数十万样本规模时。改进方向包括引入近似最近邻算法、批处理嵌入提取以及利用vLLM等高效推理后端。梯度获取在大规模分布式设置下的开销问题可以通过减少更新频率、使用轻量级代理信号以及缓存选择决策来缓解。DoReMi的三步训练过程引入了额外计算成本，可以考虑一步近似方法或与其他方法结合以减少代理模型训练开销。框架的配置接口虽然设计简洁，但在复杂场景下可能需要更多灵活性，比如支持多种策略组合以及更细粒度的调度控制。
未来方向
作者提出可以将框架扩展到多模态场景，支持图像、音频等多模态数据的动态选择和混合。基于实验结果可以延伸的方向包括：探索更复杂的数据策略组合机制，如在训练不同阶段自适应切换不同策略；引入元学习或强化学习来自动学习最优的数据策略调度策略；扩展框架支持持续学习场景，在模型演化过程中动态调整数据使用策略；研究跨模态的数据选择和混合方法，在视觉语言预训练中实现文本和图像数据的协同优化；开发更高效的近似算法和分布式协议，将框架扩展到更大规模的训练场景；集成自动超参数调优和神经架构搜索，实现数据和模型架构的联合优化。
复现评估
项目完全开源，代码仓库在GitHub，文档在DataFlex-Doc网站，数据集和演示在HuggingFace collections。实验使用8块NVIDIA H20 GPU进行数据选择实验，30B token实验使用4节点乘以8 H20 GPU总共32 GPU协调执行。框架支持从标准pip安装，与LLaMA-Factory兼容，可以无缝集成到现有训练流水线。配置文件使用标准YAML格式，实验参数详细记录在论文附录中。数据选择实验在Open-Hermes-2.5的10万样本子集上进行，使用MMLU验证和测试分割构建验证和测试集。数据混合实验使用SlimPajama的两个子集，保留原始SlimPajama语料库的自然token级领域比例。所有实验使用固定随机种子42，使用BFloat16混合精度训练。论文附录提供了与原始实现的详细对比以及作者贡献声明。总体而言，论文在开源、数据、算力要求方面都提供了充分的复现支持，复现难度中等。
论文图表