本文转载来源:新智元公众号,不代表官方立场
测试时计算之后,谷歌三大团队集众人之力,发现了全新的Scaling Law!
谷歌研究员Zachary Charles宣布:「在越来越大的模型上,分布式训练取得重大突破」。
这个核心算法,便是——DiLoCo的Scaling Law。
新的训练方法无惧模型规模,未来,在「多个数据中心」训练大模型的规模不再是问题。
论文得出四大发现,DiLoCo训练方法的Scaling law,效果远超「数据并行」:
更稳健(Harder)
:在不同模型规模下,DiLoCo的超参数依然保持稳定且可预测。
更优越(Better)
:随着模型规模扩大,DiLoCo相较于数据并行训练的优势进一步提升。
更高效(Faster)
:DiLoCo所需的带宽比数据并行训练少几个数量级。
更强大(Stronger)
:DiLoCo能够容忍比数据并行训练大得多的批大小。
值得一提的是,这篇巨作集结了谷歌三大团队:谷歌Research、谷歌Search、谷歌DeepMind。
论文地址:https://arxiv.org/pdf/2503.09799
在固定计算预算下,研究人员探讨了DiLoCo在训练大模型时的Scaling law。
论文中,重点分析了算法因素(如模型副本数量、超参数设置、token预算)如何影响训练过程,并证明这些影响可通过Scaling law准确预测。
结果表明,DiLoCo在模型规模增长时,表现出稳定且可预测的扩展性。论文合著者Arthur Douillard再次强调:DiLoCo生效了!
智能的未来将是分布式的,而DiLoCo可能正是那个关键的要素
在合理调优的情况下,DiLoCo比数据并行训练更具扩展优势,即使在小规模模型上也可能优于数据并行训练。
这些发现,揭示了DiLoCo的强大优势:不仅解决了通信瓶颈,还为大规模模型训练开辟了全新的可能。
有网友惊叹地表示,「DiLoCo可能会重新定义LLM Scaling的方式!更少的带宽需求,更高的效率」。
数据并行训练在大模型上表现出色,前提是在计算资源集中分散的情况下,才能实现。
如果计算分布较广,通信就可能成为巨大的瓶颈,尤其是当模型规模增长时,问题会更加严重!
机器学习采用的解决方案,例如在联邦学习和数据中心训练中,就是让多个独立模型进行训练,并定期同步。
随着机器学习模型规模的扩大,数据并行方法固有的频繁同步需求会导致显著的性能下降,这对进一步扩展模型构成了关键挑战。
那么,如何在保持模型质量的同时,降低同步需求,以突破这一瓶颈呢?
答案或许就在,DiLoCo(Distributed Low-Communication)这一创新方法中。
论文链接:https://arxiv.org/abs/2311.08105
每个DiLoCo模型副本都会独立训练H个内部优化(inner optimization)步骤。
这些模型通过外部优化(outer optimization)步骤进行同步,通常在外部优化步骤之间引入动量机制。
在下图中,示例中共有M=4个模型副本。
DiLoCo的成功已经被反复验证。它的运作方式与联邦学习的FedOpt方法类似。
此外,研究人员也多次证明DiLoCo在大模型(LLM)训练中的卓越表现。
那么DiLoCo有什么问题?简单来说——规模
。
DiLoCo与数据并行训练不同,它引入了额外的「外部」超参数,并且实际上的表现和理论上明显不同。
这正是研究scaling laws的目的!
这次研究从零开始构建了DiLoCo和数据并行训练的Scaling law,用于预测它们在大规模模型上的表现对比。
在数据并行训练中,每个训练步长都会处理一个大小为B的数据批。
在本研究中,批大小指的是批中的token数量(而不是序列数量)。
计算批梯度,并使用学习率γ进行优化。
在DiLoCo训练过程中,每个时间步t处理一个全局批大小为B的数据,并在序列级别将其均匀分配到M个DiLoCo副本中。
因此,全局批大小仍然是B,而每个DiLoCo副本的本地批大小为B/M。与数据并行训练类似,每个副本都会计算批梯度,并使用学习率γ执行一次内部优化(inner optimization)。
但与数据并行不同的是,DiLoCo每H步会执行一次「外部优化」(outer optimization),基于参数空间计算的外部梯度(outer-gradients),并使用学习率η进行更新。
一个重要的对比是数据并行vs.DiLoCo(M=1)。
虽然它们相似,但并不完全相同。
DiLoCo在M=1的情况下,仍然包含一个外部优化器(OuterOpt)步骤,因此它可以被视为Lookahead优化器的变体。
而在DiLoCo中,OuterOpt通常使用带有Nesterov动量的GD,这意味着DiLoCo(M=1)实际上是数据并行训练的一个变体,但动量操作仅每H步执行一次。
还进行了大量实验,涵盖训练过程的各个方面,全面分析了它们的扩展行为。
大部分实验里,研究团队使用C4数据集的训练集来训练模型,评估指标用C4的验证集。
另外,还在三个下游任务上算了零样本评估指标:HellaSwag、Piqa和Arc-Easy。
研究团队用的是一个类似「Chinchilla」的纯解码器Transformer架构,加入了QK-LayerNorm,还使用了z-loss正则化来让训练更稳定。
他们把多个序列打包到每个批里,最大序列长度全程固定为2,048。
所有模型都是从零开始训练的,因为这次主要想研究预训练阶段的规模规律。
研究团队训练了一堆模型,调整了Transformer层数、注意力头的数量、QKV维度和前馈层的隐藏维度。
除非特别说明,他们都用Chinchilla的token预算,并且对除了最大的两个模型(4B和10B参数)外,其他模型都做了大量的超参数调整。
研究团队用AdamW作为数据并行(Data-Parallel)的优化器,也是DiLoCo的内层优化器。两个算法的β1设为0.9,β2设为0.99。
训练开始有1000步的预热,然后用余弦学习率衰减。权重衰减参数λ设为T⁻¹,其中T是总训练步数(取决于批大小和token预算)。到训练结束时,学习率衰减到峰值的5%。
为了训练稳定,他们把(内层)梯度的全局ℓ2范数剪裁到1,外层梯度不剪裁。
对于DiLoCo,他们用带Nesterov动量的SGD作为外层优化器。动量设为0.9,外层学习率保持不变。
DiLoCo的评估损失随着N的增加,相对于数据并行(Data-Parallel)有所改善。
Scaling law预测,当M=2时,DiLoCo在参数达到几十亿以上时,损失会比数据并行更低。这一现象在研究调优的最大模型以及4B和10B模型的训练中都得到了验证。
下图2展示了DiLoCo和Data-Parallel两种算法在不同模型规模(N)下的表现对比。
图(a)显示,随着模型规模从2^25到2^31逐渐增大,DiLoCo(分别在M=1、2、4、8时)和Data-Parallel的评估损失(EvalLoss)都在下降,但DiLoCo的损失下降得更明显,尤其是在M值较大时。
图(b)进一步展示了DiLoCo相对于Data-Parallel的评估损失的百分比差异,可以看出,随着模型规模增加,DiLoCo的损失比Data-Parallel低得越来越多,说明DiLoCo在模型规模扩大时表现更优越。
这个发现有两个独立但相关的部分:
-
DiLoCo(M=1)表现更好:就像上面提到的,DiLoCo在M=1时,所有模型规模的评估损失都比Data-Parallel低。而且随着模型参数规模N增加,Data-Parallel和DiLoCo(M=1)之间的差距越来越大。
-
DiLoCo(M≥2)的表现:在大多数模型规模下,DiLoCo在M≥2时评估损失会更高。不过,如果看DiLoCo和Data-Parallel之间的百分比差异(带正负号),会发现随着N增大,DiLoCo相对Data-Parallel的表现越来越好,甚至在M=2、N=2.4亿参数时超过了Data-Parallel。
比如,研究团队在下表4中列出了Data-Parallel和DiLoCo在不同模型规模N下的评估损失。
可以看出,不管M是多少,百分比差异都随着N增加严格减小。
这个趋势在图2中也有展示:随着N增加,DiLoCo的相对评估损失逐渐降低。