大语言模型(LLMs)在各类任务中的卓越表现,相信大家都有目共睹。但要想高效部署和运用这些模型,一个关键环节不可忽视——超参数优化。过去,这一过程很大程度上依赖“盲猜”或“经验试错”,不仅成本高昂,结果也往往不尽如人意。
为了彻底解决这一难题,我们团队开展了一项大规模的实证研究。通过在多种配置下进行系统的网格搜索,最终揭示了一套通用的最优超参数缩放定律(Optimal Hyperparameter Scaling Law)。
先看看核心结论:最优学习率与模型参数规模及数据规模呈幂律关系,而最优批量大小则主要受数据规模影响,与模型参数规模关联不大。 这相当于为超参数调优提供了清晰的导航图。
进一步研究发现,当固定模型参数和数据规模后,观察超参数的损失函数曲面,会发现一个有趣特性——凸性特征。这意味着存在一个稳定的超参数最优区间。换句话说,只要超参数落在这个区间内,模型性能就能稳定在接近最佳的水平,调整起来更有把握,无需追求那个“物理极限上的完美点”。
基于这些发现,我们为社区准备了一套通用的即插即用超参数优化工具。测试结果令人信服:使用该工具估算的超参数配置,与通过穷举搜索找到的全局最优LLM性能相比,误差仅为 0.09%。这表明该工具的可靠性和有效性已得到充分验证,甚至在某些场景下可以替代成本极高的穷举搜索。
更令人惊喜的是,这套缩放规律在不同模型稀疏度、训练数据分布以及模型结构变化下,都展现出惊人的鲁棒性。据我们所知,这是首个同时适用于不同模型结构(比如专家混合模型Mixture-of-Experts和稠密Transformer结构)以及不同数据分布的最优超参数缩放定律研究。
为了得出这些结论,实验规模是空前的。我们消耗了近百万小时的NVIDIA H800 GPU计算资源,从零开始训练了 3,700个不同模型尺寸和超参数配置的LLM,总计处理了约 100万亿个tokens。每一组数据背后都代价高昂,也因此,结论的含金量足够高。
最优超参数缩放法则:「Step Law」
下面我们具体来看这个最优超参数缩放法则——我们称之为 「Step Law」。
具体公式是这样的:最优学习率会随模型参数规模与数据规模呈幂律变化,而最优批量大小则仅与数据规模相关。这意味着,当模型变大或训练数据增多时,学习率需要以特定方式调整,而非线性缩放。
图一展示了一个10亿参数的模型在1000亿个tokens上训练的超参数空间。那个红色的“全局最优点”,就是每对学习率和批量大小组合下,最终训练损失最低的那个点。等高线直观地告诉我们,靠近这个点的区域损失变化平缓,而偏离远了损失会急剧增加。我们的Step Law在预测这个最优点方面,准确性最高,几乎与真实全局最优点重合。
图一:在400M的 Dense LLM上训练40B Token(左)和在1B的 Dense LLM上训练100B Token(右)的超参-损失等高线图
(图片保留)
接下来看实验细节。当固定模型与数据规模时,超参数优化的损失函数曲面呈现出明显的凸性特征。这意味着存在一个稳定且易寻的最优超参数区域。
图二:Learning Rate 与 Batch Size在1B模型训练100B Token上的损失分布。散点图(左)与3D曲面(右)图
(图片保留)
这里每一张图背后的数据都来自120个从头训练的小模型。训练结束后的收敛Loss,在三维空间中构成了一个类似“山谷”的形状。可以看到,这个山谷的底部,也就是Loss最低的那片区域,非常平坦。这说明最优Learning Rate和Batchsize很可能是一个较大的区间,而不是一个针尖般的点,这给实际应用带来了很大便利。
为了便于学界和业界直接使用,我们推出了一个最优超参数估算工具(https://step-law.github.io)。该工具的预测结果与穷举搜索出的全局最优超参数相比,性能差距仅有 0.09%。同时,网站上还公开了所有超参数组合的loss热力图。
图三:1B模型、100B Token训练上的LR与BS热力图
(图片保留)
在这张图中,每一个点位上的数字,都是对应小模型训练结束后的真实Train Smoothed Loss。红点就是我们公式的预估值所对应的BS、LR位置,可以看到它就在那片平坦的“最优盆地”里。
(所有热力图见:https://step-law.github.io/)
相关研究梳理
要在大规模训练中找到最优超参数,前提是固定的条件很多:固定的模型结构、固定的数据分布、固定的模型参数规模以及固定的数据规模。
对比现有的最优超参数估算公式,我们的研究进行了全覆盖的网格搜索——充分覆盖了模型参数规模、训练数据规模、批量大小 (BS) 和学习率 (LR)。最终得出的缩放法则,在适用性和准确度上都有大幅提升。具体对比请看下表。
表一:不同方法的最佳超参数缩放定律比较
(表格保留)
(注:Data Recipe指是否在不同预训练语料配比下研究过最优超参;Model Sparsity指是否同时支持MoE和Dense模型,以及不同稀疏度下的MoE模型;LR指峰值Learning Rate;BS指Token Wise的Batch Size。)
「普适性」超参数缩放法则的三大性质
算法不能只在实验室里好用,还要经得起考验。我们的Step Law在以下三个方面展现出了很强的普适性。
1. 跨模型形状的稳定性
我们专门探究了不同模型形状(比如宽度大一些还是深度大一些,或者两者平衡)对缩放规律的影响。结果发现,无论模型是什么形状,Step Law的表现都非常稳定。这说明这条规律不只针对某一种特定模型,在更广泛的架构设计空间里也适用。
图四:最优超参在不同Model Shape下的拓扑不变性
(图片保留)
这里固定了模型的非词表参数量大小和训练Token数,但改变了层数、Hidden Dimension和FFN倍数。红色五角星是Step Law预测的点位,可以看到它在6种不同的Shape上都预测到了Global Minimum附近。当然,不同Shape下底部那片“最优区域”的位置确实会发生微小位移,但规律本身是稳定的。
2. 跨模型架构的泛化性
更进一步,这套缩放规律不仅适用于稠密模型,还能很好地推广到不同稀疏度的MoE模型。对于不同的模型结构,它表现出了极强的泛化能力。
图五:不同稀疏比下MoE模型的超参-损失等高线图
(图片保留)
无论低稀疏度、中等稀疏度,还是不同D/N比的配置,Step Law的预测位置都在Global Minimum + 0.5%的范围之内,大多数情况下甚至只在0.25%以内。这充分验证了Step Law的鲁棒性。
3. 跨数据分布的稳定性
最后,我们验证了不同数据分布下的规律一致性。不管数据是英语主导、中英双语、Code和英语混合,还是代码主导,Step Law都稳定如一。
表二:实验的不同数据分布
(表格保留)
在这几种数据分布下,实验结果同样非常理想。
图六:不同数据分布下的超参-损失等高线
(图片保留)
每一个图下方都是从头训练的45个模型,总共135个模型。Step Law预测的最优Batch Size/Learning Rate,都在最低Loss +0.125%或0.25%的范围内。
研究细节解读
1. 学习率调度策略优化
学习率如何衰减,对最优超参的选择影响很大。传统做法是把学习率从一个最大值逐渐减小到最小值的十分之一;而我们的方案则采用恒定的绝对最小值(10^-5)。
图七:不同学习率策略的比较
(图片保留)
从等高线图可以清楚看到,传统衰减方法(蓝色)的最优学习率区域会向左偏,即向较低的学习率区间偏移。原因是传统方法的高初始学习率会拉高最低学习率阈值,导致训练末期学习率偏大,参数更新幅度过大,损失函数在收敛阶段持续振荡。而采用固定最小学习率的策略(红色),则可以解耦初始学习率与终值,在训练后期始终保持合适的更新步长。这个思路与业界的训练经验也是一致的。
2. 训练损失与验证损失的最优超参一致性
根据DeepMind Chinchilla的研究,平滑训练损失可以作为验证损失的无偏估计。我们也沿用了这个设定,并通过实验做了补充验证。
图八:平滑训练损失的超参-损失等高线图(左)和验证损失的超参-损失等高线图(右)
(图片保留)
在429M参数模型上训练40B Token的验证结果显示,当平滑训练损失达到最优时,学习率为1.95×10^-3,批量大小为393,216——这个点恰好与验证损失最优时的超参数完全重合。而且,两种损失在不同学习率和批量大小下的偏离趋势也高度一致。这表明,平滑训练损失能够为学习率和批量大小的选择提供稳定可靠的优化指导,而且可以省去大量在验证集上推理的算力。
3. 最优超参的Scaling Law拟合
图九:(a)散点图表示模型规模为N时,经验最优学习率与批量大小的关系;(b)散点图表示数据集规模为D时,经验最优学习率与批量大小的关系
(图片保留)
在双对数坐标下,可以清晰地看到:
- 最优学习率:随模型规模增大而减小,随数据规模增大而增大。
- 最优批量大小:随数据规模增大而增大,与模型规模弱相关。
因此,我们可以把这两个关系建模成幂律公式。通过对数变换和Bootstrap采样,我们最终给出了一套精确的预测公式。这套公式在研发过程中已经被我们在大于1B的模型和非极端的D/N设置下广泛使用。
讨论与未来工作
面对海量的实验结果,我们也坦诚地说,目前的分析仍然是不足的。我们非常期待更多社区的研究员能参与进来,发掘这些海量实验中的宝藏,揭示出更多的规律,并给出理论解释。我们会陆续将实验的各个细节整理并开源出来,但由于实验量巨大,还需要些时间来整理。
抛砖引玉,我们觉得以下几个话题非常值得进一步探讨:
- 在给定模型、训练Token数的情况下,(Loss, BS, LR)这个三维空间是否真的是一个“凸碗”?
- 有没有更好的方法去拟合最优BS和LR,能兼容它们之间内在的耦合关系?
- 虽然Step Law在不同Model Shape、不同稀疏度的MoE模型上是鲁棒的,但次优区域在不同配置下是变化的——能否给出更清晰的解释?
- 基于海量Grid Search数据驱动的结论,其背后更深层次的理论解释是什么?
- 不同的超参、不同Model Size、Model Shape、Model Sparsity下的Training Dynamic研究。
我们团队的「Predictable Scale」是一个论文系列,很多实验已经完成。后续将可能讨论超大模型性能预测、Code & Math Scaling性质、不同Attention类型的Scaling性质等。敬请期待!
