在人工智能的激烈竞争中,一场耗资百万美元的大规模实验正悄然改变着大语言模型的训练方式。阶跃星辰研究团队日前发布重磅研究成果,他们通过耗费近100万NVIDIA H800GPU小时的算力,从零开始训练了3,700个不同规模的模型,累计训练了惊人的100万亿个token,揭示出一条被称为"Step Law"的普适性缩放规律,为大语言模型的高效训练提供了全新指南。
这项研究不仅仅是对超参数优化的探索,更是第一个全面考察模型最优超参在不同形状、稀疏度和数据分布下稳定性的工作。研究结果表明,无论模型采用何种架构设计,无论训练数据来自何种语言或领域,Step Law都表现出令人惊叹的鲁棒性,这大大增强了该工具在实际应用中的价值。
团队训练的3,700个模型涵盖了不同规模、不同超参数组合、不同形状、不同数据配比以及不同稀疏度的配置,包括MoE和Dense两种架构。通过这些海量实验,他们发现最优学习率随模型参数规模与数据规模呈现幂律变化,而最优批量大小主要与数据规模相关。这一发现颠覆了业界对超参数设置的传统认知。
图源备注:图片由AI生成,图片授权服务商Midjourney
实验数据显示,在固定模型大小与数据规模的条件下,超参数优化的Landscape呈现明显的凸性特征,这意味着存在一个稳定且容易寻找的最优超参数区域。为了验证这一点,研究团队构建了三维可视化空间,直观展示了学习率与批量大小对训练损失的影响。结果清晰地展现出"山谷"形态,凸性底端是一个相对平坦的区域,这为实践中的超参数调优提供了宝贵的理论依据。
为了让这一发现惠及整个AI社区,团队开发并推出了一款通用的最优超参数估算工具。该工具的预测结果与通过穷举搜索得到的全局最优超参数相比,性能差距仅为0.09%。这意味着研究人员和工程师们可以不再依赖昂贵的网格搜索,而是直接通过这一工具获得接近最优的超参数配置。
更让人印象深刻的是Step Law的普适性。研究团队从三个不同角度验证了其适用范围:首先,无论模型形状如何变化——是偏向宽度、偏向深度,还是宽深平衡——Step Law都能准确预测最优超参数区域;其次,这一规律不仅适用于Dense模型,还能很好地扩展到不同稀疏度的MoE模型;最后,无论训练数据是英语主导、中英双语、代码与英语混合,还是以代码为主的分布,Step Law都表现出了惊人的稳定性。
研究还揭示了学习率调度策略的优化方向。与传统的学习率衰减策略不同,团队提出采用固定的最小学习率(1e-5),而非传统方法中将最小值设为最大值的十分之一。这一改变使得训练在后期能够维持更为合理的参数更新步长,有效避免了损失函数在收敛阶段的持续振荡。
此外,研究发现平滑训练损失与验证损失的最优超参数高度一致,这一发现为超参数选择提供了更为经济的方法——研究人员可以通过监控平滑训练损失来指导超参数调整,而无需频繁在验证集上评估模型性能。
尽管取得了显著成果,阶跃星辰研究团队坦言这仅仅是一个开始。他们计划陆续开源实验的各个细节,包括近4000个模型的最终检查点,以供整个社区进行更深入的分析和理论解释。未来的研究方向包括探索Loss-BS-LR三维空间的凸性、改进最优超参数的拟合方法、解释不同配置下次优区域的变化,以及深入研究不同设置下的训练动态。
Predictable Scale系列的后续工作可能将进一步讨论超大模型性能预测、Code&Math的缩放性质,以及不同Attention类型的缩放特性。可以预见,这一系列研究将为大语言模型的高效训练提供更全面的理论指导和实践工具,推动AI技术向更高效、更可控的方向发展。