模型合并之 PMA 策略：大型语言模型预训练的 “加速引擎”-优快云博客

本文链接：https://blog.youkuaiyun.com/jike007gt/article/details/148773181

大型语言模型发展迅猛，但预训练成本高昂、训练不稳定等问题需要解决。本文探讨了预训练模型合并技术，尤其是 PMA 策略，为模型开发带来新思路。在模型性能、训练效率及成本优化等方面均有显著优势，虽有局限但前景广阔，引发对 AI 模型优化与资源高效利用的深刻思考。

大家好，我是肆〇柒。在 AI 领域，大型语言模型（LLM）在自然语言处理、智能问答系统和代码生成等众多任务中展现出卓越能力。然而，这些模型的发展面临诸多挑战，包括预训练成本高昂、特定领域后训练效果有限、性能扩展预测不确定以及大规模训练不稳定等问题。而模型合并技术作为一种新兴方法，为解决这些挑战提供了新思路。尽管模型合并技术在后训练阶段已显示出显著优势，但在预训练阶段的研究相对较少。我看到一篇关于在预训练阶段模型合并的研究论文，这篇论文是由 ByteDance Seed 发表的，此文深入探讨了预训练过程中的模型合并技术，并重点介绍了一种创新的预训练模型平均（PMA）策略。下面我们就一起来看看这篇论文说了啥。

研究背景

大型语言模型虽在性能上表现出色，但其发展面临多重阻碍。预训练成本高昂，例如 GPT-3 的预训练成本高达数百万美元，对中小研究团队构成巨大经济压力。特定领域后训练效果不理想，模型在通用领域的优异表现难以完全复制到特定专业领域。性能扩展预测的不确定性使得模型开发方向难以确定。大规模训练的不稳定性可能导致训练进程失败。

然而，模型合并技术在后训练阶段的应用为我们提供了借鉴。DARE 方法可通过合并不同模型，在 GSM8K 数据集上的得分从 2.2 提升至 66.3，展现出模型合并挖掘模型潜力的巨大优势。但预训练阶段的模型合并研究较少，LAWA 方法虽尝试利用模型合并加速 LLM 训练，但随着模型和数据规模的爆炸式增长，独立研究人员面临无法获取大规模预训练中间 check point 的困境，限制了预训练模型合并的研究进展。

研究方法：预训练模型平均（PMA）策略

PMA 策略作为本文的核心创新点，主张在预训练的不同阶段，将多个模型 check point 的权重进行融合，实现模型层面的深度整合。

在实施过程中，研究者们选择了 Dense 模型和 MoE 架构作为实验载体。Dense 模型能够充分捕捉数据的复杂模式，而 MoE 架构在处理大规模数据时展现出独特优势。学习率调度策略采用 Warmup-Stable-Decay（先在 warmup 阶段缓慢提升学习率，接着在稳定阶段保持恒定学习率，最后在 decay 阶段逐步降低学习率）策略。评估指标涵盖多个维度，包括模型在下游任务中的准确率和训练过程中的损失变化等。

实验设计全面而巧妙。研究者们在恒定学习率阶段对 PMA 策略进行了深入探索，还踏入了余弦退火阶段这一复杂领域。他们对不同规模的模型进行了全面训练，从小型的Seed-MoE-1.3B/13B 到大型的 Seed-MoE-10B/100B 等。所使用的数据集规模宏大，内部预训练语料包含数万亿个 token，涵盖多种语言和领域，为模型训练提供了丰富资源。

实验结果与分析：性能提升

在恒定学习率阶段，PMA 策略为模型性能带来了显著提升。以 Seed-MoE-1.3B/13B 模型为例，在 Humaneval 基准测试中，其性能从 31.1 提升至 36.6；Seed-MoE-10B/100B 模型在 GSM8K 数据集上的表现也从 54.3 提升至 61.6。这些提升代表着模型在代码生成和数学问题解决等复杂任务上能力的显著增强。

如下图所示，不同大小的 MoE 模型在稳定训练阶段经过模型合并后，在多个下游任务上的性能得到了显著提升，充分展示了 PMA 策略的有效性。

不同大小的Mixture of Experts（MoE）模型在稳定训练情况下的下游任务性能比较，以及模型合并前后的性能对比

在余弦退火阶段，PMA 策略展现出了惊人效果。随着学习率的逐步降低，模型表现出良好的收敛性。更令人惊喜的是，PMA 策略在早期退火阶段就能取得与退火过程结束时相当甚至更优的性能。这表明 PMA 策略能够更早地找到接近最优解的参数组合，加速模型的收敛过程。

如下图所示，不同大小的 MoE 模型在余弦退火阶段经过模型合并后的整体性能变化。

不同大小MoE模型退火训练下，合并前后整体性能对比

仅仅使用 Warmup-Stable 阶段与 PMA 策略相结合，跳过退火阶段的实验设计，结果有力地论证了这种简化训练流程的可能性。这种创新的训练方式在加速模型验证的同时，还显著降低计算资源的需求，为大规模模型训练节省大量时间和金钱成本。

如下图所示，在稳定训练阶段使用 PMA 策略与真实退火模型的性能对比。

稳定训练下的模型融合结果与真实退火模型在下游任务性能上的比较

不同合并方法对比

在模型合并方法的探索中，研究者们对 Simple Moving Average（SMA）、Weighted Moving Average（WMA）和 Exponential Moving Average（EMA）这三种常见的方法进行了深入研究和比较，以应对不同的训练场景和需求。

这三种方法各有特点。SMA 方法在所有模型 check point 上均匀分配权重，对待每个 check point 都一视同仁。WMA 方法更注重近期的模型 check point ，认为它们更具参考价值，因此赋予它们更大的权重。而 EMA 方法通过指数递减的方式分配权重，对最近的模型 check point 更为敏感，能够快速响应模型在训练后期的变化。

从数学原理来看，WMA 方法在训练早期表现更优的原因在于其权重分配策略能够更好地捕捉模型参数的变化趋势。在训练初期，模型参数更新较快，近期的 check point 包含了更多的新信息。WMA 通过赋予这些 check point 更大的权重，能够更有效地利用这些新信息，从而更快地逼近最优解。随着训练的进行，模型参数逐渐趋于稳定，各方法之间的权重差异对性能的影响逐渐减弱。这一过程可以通过对模型参数更新的数学建模来进一步验证。假设模型参数在训练过程中遵循某种动态变化规律，如随机梯度下降（SGD）中的参数更新公式：

在 Seed-MoE-1.3/13B 模型的实验中，这些方法展现出了不同的优势。在训练早期，模型的参数尚未稳定，变化较大，此时 WMA 方法因其对近期 check point 的重视，在性能上表现更为出色。然而，随着训练的推进，模型参数逐渐趋于稳定，各方法之间的性能差异逐渐缩小。基于其实现的简单性和在大多数情况下的稳定性，SMA 方法在后续实验中成为了优先选择的对象，它在各种复杂的环境中都能保持稳定的发挥。

如下图所示，不同模型合并方法（WMA、SMA、EMA 不同参数）对最终模型性能的影响在训练的不同阶段各有优势，这为选择合适的合并方法提供了依据。

不同模型融合方法对最终模型性能的影响

模型合并的优化应用

确定最佳合并参数

在模型合并的优化应用中，确定最佳合并参数至关重要。研究者们通过对不同规模模型的消融研究，探索了模型合并时的最佳间隔（V）和模型数量（N）这两个关键参数。以 Seed-MoE-1.3/13B 模型为例，在训练早期，较小的间隔和较少的模型数量会因包含过多不稳定的权重而导致性能下降。而随着训练的进行，模型的权重逐渐稳定，较大间隔和较多模型数量的组合则能更好地融合不同阶段的知识，提升模型性能。研究发现，对于不同规模的模型，最佳参数选择呈现出一定的规律性。小型模型如 1.3B/13B 模型，最佳间隔约为 8B tokens；而大型模型如 10B/100B 模型，最佳间隔则需要扩大到 80B tokens 左右。

如下图所示，不同间隔（V）和模型数量（N）对模型性能的影响揭示了确定最佳合并参数的重要性。

不同模型融合超参数对最终模型性能的影响

下游训练的促进作用（PMA - init）

PMA - init 作为 PMA 策略在下游训练阶段的创新应用，将模型合并后的权重作为初始化权重，以改善下游模型性能并稳定训练过程。在持续训练（CT）阶段，实验结果显示，采用 PMA - init 的模型在初始训练阶段的损失值普遍低于基线模型。尽管在训练后期，两种初始化方式的损失值逐渐趋于一致，但 PMA - init 模型在早期的快速收敛为整个训练过程节省了宝贵时间。在监督微调（SFT）阶段，虽然 PMA - init 的优势不像在 CT 阶段那样明显，但它仍然在某些情况下能够带来性能提升，并且不会对最终性能造成负面影响。

下表展示了在不同学习率调度下，PMA-init 对模型性能的具体提升数据。可以对照前文关于 PMA-init 在 SFT 阶段作用的描述来理解，帮助大家更直观地理解 PMA-init 在不同学习率设置下的效果差异。

模型	Open-Benchmark					In-house Evaluation
指标	MMLU	LiveBench	AMC-2023	GPQA	LiveCodeBench	OOD	Reasoning	IF
Baseline2e-5- >2e-6	86.8	50.5	61.0	55.2	39.7	32.6	32.1	36.3
PMA2e-5- >2e-6	87.1	52.0	64.0	54.0	39.4	34.7	34.0	38.8
PMA1e-5- >2e-6	87.2	53.2	65.5	54.4	39.7	33.8	33.2	37.3
PMA4e-5->2e-6	87.0	51.3	61.4	54.0	39.2	31.8	32.6	37.2

如下图所示，PMA-init 技术与基线模型在不同学习率调度下的损失曲线和性能指标对比，突出了 PMA-init 在 CT 阶段的优势。

损失曲线（上）和性能指标（下）在CT阶段随不同学习率调度策略的变化，其中采用余弦调度器将学习率从学习率峰值衰减到学习率终点（表示为lrpeak→lrend）的对比情况

对训练稳定性的提升

在大规模 LLM 训练中，训练不稳定性现象如损失突增或发散等问题，可能破坏训练进程，浪费大量计算资源。PMA - init 在提升训练稳定性方面表现出色。通过在小模型上使用大比率学习率模拟训练不稳定性的实验，研究者们发现，在监督微调（SFT）阶段应用 PMA - init 能够显著稳定 GradNorm 指标，并减少损失突增的频率。当模型遭遇不可恢复的损失突增时，PMA - init 提供了一种可靠的恢复方法。通过合并前几个保存的 check point ，训练可以重新回到正确轨道上，继续沿着原有轨迹前进，避免了资源的巨大浪费。

如下图所示，PMA-init 对初始化与 SFT 阶段模型的 GradNorm 指标以及在小模型上使用 PMA-init 恢复训练前后的损失曲线的影响，突出了其对训练稳定性的提升作用。

上：使用PMA初始化进行SFT训练的GradNorm对比。下：从PMA初始化恢复训练与原始训练的预训练损失曲线对比

模型合并的机制探究

从理论层面深入探究模型合并的机制，研究者们发现，不同模型 check point 的权重平均能够有效降低模型的损失。这是因为不同 check point 在参数空间中探索了不同的区域，它们的权重偏差在合并后相互补充，使得合并后的模型更接近最优解。通过可视化方法，选取 Seed-MoE-1.3B/13B 预训练中的多个 check point ，研究者们绘制了特定层中两个参数的平均分布情况以及 MMLU（Massive Multitask Language Understanding，一个综合性的基准测试，用于评估模型在多种语言理解任务上的表现）分数的等高线图。这些图像清晰地展示了各个模型权重位置以及合并后的模型权重位置如何更接近高 MMLU 分数区域，进一步佐证了模型合并的效果和原理。

如下图所示，通过可视化方式展示了模型合并后权重位置更接近高 MMLU 分数区域的现象，揭示了模型合并的原理。

MMLU分数等高线的可视化，用于比较原始模型与合并模型的权重

总结以及自己的实践

这篇论文通过深入研究和实验验证，得出了一个明确的结论：PMA 策略在预训练模型合并领域具有重大意义和价值。它在稳定训练阶段合并 check point ，不仅能够带来显著的性能提升，还可以准确预测退火行为，为模型训练提供了一种强大的模拟工具，从而简化开发流程并降低训练成本。PMA - init 在下游训练中的成功应用，尤其是在提升训练稳定性方面的重要作用，为模型训练提供了一种可靠的恢复方法。

但是这篇论文的研究也存在一些局限性。在学习率对模型合并影响方面的探索仍显不足，不同学习率设置下模型合并效果的差异可能对实际训练过程中的学习率调度策略选择产生重要启示，这需要进一步深入挖掘。此外，模型合并技术在 RL 训练等post train场景中的应用前景广阔，但也面临着诸多挑战。RL 训练中模型与环境交互的复杂性和特殊性，如何将模型合并技术与 RL 算法相结合，以提高模型在强化学习过程中的稳定性和性能，也可以是未来关注火探索的方向。

在写本文的过程中，我对模型合并技术在预训练阶段的应用有了一定的了解。我认识到，模型合并是一种蕴含着深刻数学原理和智能策略的技术。它能够在不同的模型 check point 之间寻找到一种微妙的平衡，将各个阶段的知识巧妙融合，从而创造出更加强大和稳定的模型。

我通过阅读字节的这篇关于预训练阶段模型合并的研究论文，让我回想起去年有一段时间，我因为项目需要关注过模型合并技术，并且实操过。当然，这与本文所介绍的合并阶段不同，本文是在 pretrain 阶段，而我实践的是在 post-train 阶段进行的模型合并。

回想当时 Qwen2 的开源模型发布，让模型性能有了飞跃，甚至 Qwen2 的 72B 很长一段时间在 Huggingface 处于 No 1 霸榜的地位。但有个问题出现了，Qwen2 并不包含 32B 模型（Qwen1.5 有）。这让我有点郁闷，所以就打起了模型合并的主意。我想通过不同的合并算法，结合 finetune 来实现一个高性能的 30B 左右的模型。但后来，Qwen 团队发布了 Qwen2.5，这次有了 32B，模型合并这件事我就放下了。

到了去年年底的时候，因为建立自己的开源模型仓库的原因，我又一次关注到了 huggingface 上的 Qwen2.5 32B 的 finetune + merge 的模型，这与我当时的思路有点像，但又不完全一样。我们的区别在于：