DomainPilot：基于领域级损失引导的两阶段数据混合优化，用于高效语言模型微调

# DomainPilot：领域级损失引导的两阶段数据混合优化，实现高效语言模型微调
来源：https://arxiv.org/html/2607.22769

###### 摘要

大型语言模型（LLMs）的训练效果从根本上受限于训练数据的质量与构成。现有的动态数据调度方法在工业规模的预训练和监督微调（SFT）中面临关键限制：数据选择在TB级语料库上会产生昂贵的O(N)成本；混合优化方案会引入严重的I/O瓶颈或需要训练辅助参考模型；而样本级重加权策略依赖的损失信号则混淆了噪声、难度和新颖性。

我们提出DomainPilot，一个领域级损失引导的两阶段数据混合优化框架。DomainPilot引入了**令牌级领域损失监控**，在不中断数据流水线的情况下捕获训练过程中每个领域的学习动态。基于这些信号，我们提出一个**缩放定律引导的粗优化**阶段，拟合领域特定的收敛曲线，并推导出混合调整的有原则先验。随后的**混合定律引导的细优化**阶段通过受控扫描实验建模跨领域交互效应，进一步细化混合比例。整个机制通过**补丁架构**实现，将领域感知的损失计算注入现有训练框架（如MindSpeed/Megatron-LM），仅需约30行框架特定适配器代码。

我们在Qwen3-1.7B模型的SFT过程中验证了DomainPilot。与原始数据混合相比，我们的优化混合在MMLU-Redux上提升+2%，在AIME24上提升+1.8%，在LiveCodeBench v5上提升+3.8%，在BFCL v3上提升+3.6%，且不增加总数据量或训练成本。这些结果表明，领域级训练信号为混合优化提供了一种有效、轻量级的替代方案，无需昂贵的数据选择或辅助模型训练。

## 1 引言

大型语言模型（LLMs）的性能越来越多地不仅取决于架构创新，还取决于训练数据的策展、组成和权重[3 (https://arxiv.org/html/2607.22769#bib.bib1),2 (https://arxiv.org/html/2607.22769#bib.bib2)]。在以数据为中心的AI范式中，即使对数据混合比例的微小调整也能带来与参数翻倍相当的改进[10 (https://arxiv.org/html/2607.22769#bib.bib3)]。尽管如此，工业规模的训练流水线仍严重依赖多领域数据混合的手工启发式方法，导致大量潜在收益未实现。

现有的动态数据调度方法大致分为三类，在预训练或大规模SFT场景中各有根本性局限：

**数据选择**方法，如LESS[8 (https://arxiv.org/html/2607.22769#bib.bib4)]、DSIR[9 (https://arxiv.org/html/2607.22769#bib.bib5)]和Quad[7 (https://arxiv.org/html/2607.22769#bib.bib6)]，通过计算每个样本的梯度或重要性分数来识别高价值子集。虽然对微调有效，但这些技术需要对整个语料库进行O(N)前向传播——成本接近一个完整的训练周期——这使得它们在TB级预训练中不可行。

**数据混合优化**方法，如DoReMi[10 (https://arxiv.org/html/2607.22769#bib.bib3)]、CLIMB[1 (https://arxiv.org/html/2607.22769#bib.bib7)]和ScaleBiO[5 (https://arxiv.org/html/2607.22769#bib.bib8)]，根据模型反馈调整领域比例。DoReMi训练一个辅助参考模型（3倍计算开销），且仅在300亿令牌上得到验证，比生产预训练低两个数量级。CLIMB依赖聚类自动发现领域，但当工业数据团队已维护显式领域分类时，这一步骤并非必要。ScaleBiO的批次内动态重平衡将顺序读取转换为随机访问，引入严重I/O瓶颈，使预取缓存失效并降低GPU利用率。

**样本重加权**方法，以DataFlex[11 (https://arxiv.org/html/2607.22769#bib.bib9)]和RHO-1[4 (https://arxiv.org/html/2607.22769#bib.bib10)]为代表，在训练过程中修改每个样本的损失贡献。这种范式避免了额外的数据遍历，是大规模下唯一可行的轻量级选项。然而，现有加权策略将**样本级损失**作为数据质量的代理，忽略了高损失可能意味着噪声、困难领域内容、新颖有价值信息或短文本统计不稳定性——这些现象无法从单个标量中区分。

除了这些算法限制，还有**框架可移植性**障碍。DataFlex与HuggingFace Trainer生态系统紧密耦合（修改_inner_training_loop和compute_loss），而工业预训练主要在基于Megatron-LM的框架（如MindSpeed）上运行，这些框架拥有完全独立的训练循环、数据加载器和损失计算路径。Select和Mix模式需要重建DataLoader，在Megatron的内存映射二进制数据流水线下不可行。只有Weight类干预——仅修改损失计算——才实际可迁移。

### 1.1 核心动机

这些观察结果引导出五个核心研究问题，指导我们的工作：

1.  缺乏数据驱动的混合优化。多领域训练混合通常靠手动调节设定，且在整个训练过程中保持静态，忽略了不同领域表现出不同的学习动态。
2.  样本级损失作为质量信号不够充分。标量损失无法区分噪声、难度、新颖性和不稳定性；仅基于此信号的加权决策可能导致放大损坏数据。
3.  现有调度器与工业规模训练不兼容。在线选择、参考模型训练和动态批次重平衡都会产生不可接受的成本或基础设施变更，不适用于生产预训练。
4.  微调中周期边界损失下降原因不明。在重复微调训练中观察到的周期边界“阶梯状”损失下降缺乏系统解释，模糊了泛化与记忆之间的区别。
5.  训练反馈与数据清洗之间缺失桥梁。离线基于规则的过滤缺乏机制来利用实时训练信号识别和移除低质量源数据。

### 1.2 贡献

我们提出DomainPilot，一个领域级损失引导的框架，通过以下贡献解决上述限制：

1.  令牌级领域损失监控。我们实现了一个非侵入式补丁，在预处理阶段为每个令牌标记其领域标识，在前向传播过程中按领域聚合损失，并在数据并行秩之间同步，开销可忽略。
2.  两阶段混合优化流水线。阶段1拟合领域特定**缩放定律**，提取学习动态先验（收敛损失、收敛速度、初始幅度），并计算粗略重新分配。阶段2以阶段1输出为中心，进行**混合定律**扫描实验，建模跨领域交互，并在±20%局部邻域内细化比例。
3.  补丁架构。框架采用两层设计——框架无关的算法层（纯PyTorch）和薄薄的框架特定适配器层（每个框架约30行）。这使得可以在不修改底层训练代码库的情况下部署到MindSpeed/Megatron-LM。
4.  在Qwen3-1.7B上的实证验证。我们的优化混合在LiveCodeBench v5上最高提升+3.8%，在BFCL v3上提升+3.6%，在MMLU-Redux（+2%）和AIME24（+1.8%）上也有持续提升，且不增加额外数据或计算成本。

## 2 相关工作

我们按照第1节 (https://arxiv.org/html/2607.22769#S1)引入的三个轴组织相关工作：数据选择、混合优化和样本重加权。表1 (https://arxiv.org/html/2607.22769#S2.T1)提供了高层次比较。

表1：动态数据调度范式的比较。“预训练可行”表示适用于TB级预训练或大规模SFT。

### 2.1 数据选择

数据选择旨在识别最有价值的训练子集，而不改变保留数据的混合比例。LESS[8 (https://arxiv.org/html/2607.22769#bib.bib4)]使用低秩梯度相似性选择对目标任务影响最大的样本。DSIR[9 (https://arxiv.org/html/2607.22769#bib.bib5)]将选择表述为针对目标分布的重要性重采样。Quad[7 (https://arxiv.org/html/2607.22769#bib.bib6)]联合优化质量和多样性。虽然理论上吸引人，但三者都需要评估每个候选样本——成本随语料库大小线性增长，当N超过10^9令牌时变得不切实际。

### 2.2 数据混合优化

混合优化调整领域级数据源的**比例**，而非选择单个样本。DoReMi[10 (https://arxiv.org/html/2607.22769#bib.bib3)]训练一个小型代理模型以及一个参考模型，并对代理模型表现出更高超额损失的领域进行上加权。该方法有原则，但需要从头训练参考模型（3倍总计算量），且仅已在300亿令牌上得到验证，远低于生产预训练规模。CLIMB[1 (https://arxiv.org/html/2607.22769#bib.bib7)]通过聚类发现领域并迭代搜索最优比例；然而，当领域标签已由数据工程团队显式管理时，自动聚类提供的额外价值有限。ScaleBiO[5 (https://arxiv.org/html/2607.22769#bib.bib8)]提出批次内动态重平衡，但改变每批次采样分布会破坏内存映射数据流水线中的顺序读取假设，导致预取缓存未命中及GPU空闲时间，可能抵消模型质量提升。

### 2.3 样本重加权

样本重加权修改每个训练样本对全局损失的贡献，而不改变看到哪些样本或它们所占的比例。动态损失重加权[11 (https://arxiv.org/html/2607.22769#bib.bib9)]系统研究了基于每个样本损失的线性、二次和极值加权策略。RHO-1[4 (https://arxiv.org/html/2607.22769#bib.bib10)]将其扩展到令牌粒度的选择性损失。由于重加权仅触及损失计算，几乎不引入I/O或同步开销，使其成为万亿令牌预训练唯一可行的范式。

现有重加权方法的关键限制是依赖**样本级损失**作为唯一信号。如第1节 (https://arxiv.org/html/2607.22769#S1)所述，样本损失至少混淆了四种不同现象（新颖性、噪声、领域难度、短文本不稳定性）。DataFlex通过其Warmup策略部分解决了这一点——在模型部分适应之前禁用加权——但并未解决信号的根本模糊性。

### 2.4 框架可移植性

DataFlex实现在LLaMA-Factory[12 (https://arxiv.org/html/2607.22769#bib.bib12)]内，大量重写HuggingFace Trainer内部（compute_loss、_inner_training_loop）。工业预训练框架如Megatron-LM[6 (https://arxiv.org/html/2607.22769#bib.bib11)]和MindSpeed构建了自己的数据加载器、训练循环和分布式通信模式。在Megatron的二进制mmap流水线下重建Select或Mix模式的数据加载器不可行。因此，只有Weight类干预——仅需拦截每个样本损失和领域标签——是实际可移植的。DomainPilot的补丁架构正是围绕这一约束设计的。

## 3 方法论

DomainPilot包含四个相互关联的组件：（1）令牌级领域损失监控，在训练时捕获每个领域的学习动态；（2）缩放定律引导的粗优化，推导混合重分配的有原则先验；（3）混合定律引导的细优化，通过建模跨领域交互来细化比例；（4）补丁架构，在不修改底层训练框架的情况下实现这些机制。我们依次描述每个组件。

### 3.1 领域损失监控

DomainPilot的基础是能够在每个训练步骤测量每个领域贡献多少损失。与样本级损失——混淆了噪声、难度和新颖性——不同，领域级损失聚合了领域内数千个令牌的信号，平均掉了样本特定的异常值，揭示了真正的学习趋势。

#### 令牌级领域识别。

在预处理阶段，MultiDomainPackedHandler对来自多个领域的子样本进行背包打包。在打包过程中，它为每个令牌标记其源domain_id，生成与标准令牌索引对齐的配对二进制文件（_packed_domain.bin/.idx）。由于打包可能将来自不同领域的子样本合并到单个序列中，领域标签保持在**令牌粒度**，产生一个长度等于序列长度的领域id向量。

#### 数据加载和前向传播。

DecoderPackedMTFDataset加载令牌级领域ID以及输入ID，并将它们注入样本字典。SFT训练器通过标准批次通信路径广播领域ID，将其作为关键字参数传入模型的forward方法。

#### 按领域损失聚合。

在模型前向传播内部，标准logits计算之后，一个轻量级后处理步骤按domain_id对令牌进行分组，并累积每个领域的(loss_sum, token_count)对。这些统计信息存储在self._last_domain_losses中。在每个日志记录间隔，一个训练日志补丁执行数据并行秩间的全局All-Reduce，以获得全局一致的领域损失，然后写入TensorBoard和控制台日志。

整个流水线增加的吞吐开销<1%，因为它重用了现有损失张量，仅引入了轻量级索引和规约操作。

### 3.2 缩放定律引导的粗优化

有了领域损失轨迹，阶段1拟合**领域特定缩放定律**，以刻画每个领域的损失如何随训练进度演变。函数形式为：

L_i(D)=a_i D^{-α_i} + b_i  (1)

其中D是训练步数（数据量的代理），a_i是初始损失幅度，b_i是不可约（收敛）损失，α_i是收敛速度。将这条曲线拟合到监控轨迹中，每个领域得到三个可解释的参数：

- a_i：反映初始领域难度和预训练熟悉度。a_i≈0表示模型已对该领域拥有较强的先验知识。
- b_i：渐近损失；更高的值意味着固有更难领域。
- α_i：学习效率；更小的值意味着收敛较慢，通常表明数据量不足。

#### 四因子奖励分数。

我们通过结合这些参数得出一个粗略的重分配分数：

score_i = (b_i / min_j b_j) × (max_j α_j / α_i) × (L_i_curr - b_i) / (max_j (L_j_curr - b_j)) 。