Scaling Law 由 OpenAI 团队于 2020 年正式提出,并在其论文《神经语言模型的扩展定律》(Scaling Laws for Neural Language Models)中进行了详细阐述。Scaling Law 揭示了模型性能与模型规模、数据量和计算资源之间的数学关系。但在这两年,由于实际应用局限性、理论解释不完善等原因,Scaling Law 遭到了一些质疑。而 Google 旗下两支团队的最新研究,似乎又给 Scaling Law 投上了一张支持票。

在大模型领域,数据并行训练(Data-Parallel)方法占据着核心地位,它是训练大模型的基础技术之一。然而,随着模型规模的不断增大,数据并行训练方法的通信开销和内存限制等缺点日益凸显。

为此,Google DeepMind 此前研发出一种名为 DiLoCo(分布式低通信优化,Distributed Low-Communication Optimization)的方法,其具备减少通信开销和提高扩展性等优势。

近日,来自 Google Research 和 Google DeepMind 的研究人员发现 DiLoCo 竟然具备更多优势。

研究中,他们研究了在固定计算预算之下训练大模型时 DiLoCo 的 Scaling Law 行为,借此揭示了 DiLoCo 一系列更通用的优点:包括其能够增加最佳 batch 大小、能够提高下游规模的泛化能力、能够在固定 token 预算之下改善评估损失。(注:评估损失,是指模型在验证集或测试集上计算出的损失值。)

研究人员发现随着模型尺寸的增大,DiLoCo 会呈现出可预测的稳健扩展。如果调整得当,DiLoCo 的模型规模扩展性优于数据并行训练方法,即使在小模型规模下 DiLoCo 也能胜过数据并行训练方法。


(来源:arXiv)



研究人员:Scaling Law 是准确的

而本次研究中,基于相关数据研究人员为数据并行训练方法和 DiLoCo 分别建立了评估损失和最优超参数的 Scaling Law。

借此证明,当 Scaling Law 扩展到更大的模型规模时,其具备更好的评估损失和最优超参数。

令人十分惊讶的是:许多情况下对于相同的 token 预算,在通信效率更高的 DiLoCo 上,这些 Scaling Law 预测会比数据并行训练方法产生更低的评估损失。

为了利用研究人员的 Scaling Law 来预测 DiLoCo 的超参数,他们分别在训练具有 40 亿参数和 100 亿参数的模型时加以测试。

事实证明,Scaling Law 是准确的。即使总通信量减少 99% 以上,DiLoCo 的表现也能优于数据并行训练。

对于研究中的每个实验,在不同带宽和不同延迟的网络下,研究人员都能实现理想的端到端 wall-clock 训练时间。(注:wall-clock time,指的是从任务开始到结束所经历的实际时间,包括所有等待时间和计算时间。)

研究人员的实验包括 M = 1 的 DiLoCo(即仅使用单个副本)。而在这些数据中一个令人惊讶的现象是:即使没有通信瓶颈,DiLoCo 也能改进训练效果。

由于 DiLoCo 实际上是另一种优化算法 Lookahead optimizer 的增强版本,因此它不会导致任何通信的减少。

然而,通过使用低频动量操作,无论在怎样的模型规模下,DiLoCo 在评估损失和对更大 batch 的容忍度上,都要优于数据并行训练方法。

值得注意的是,M=1 的 DiLoCo 在评估损失和训练时间方面都优于数据并行训练方法。

研究人员发现,M=1 的 DiLoCo 在所有模型尺度上都能实现较低的评估损失,并且能对更大的 batch 产生更好的鲁棒性,从而能够大大减少 wall-clock 训练时间。

研究中,研究人员分别通过数据并行训练方法和 DiLoCo 来开展模型训练。

在数据并行训练方法中,在每一步研究人员都会获取一个 batch 大小为 B 的数据。

然后,研究人员计算批次梯度(batch gradient),并使用学习率 γ 进行优化。

在比较数据并行方法和 DiLoCo 时,研究人员始终确保模型大小 N 和总 token 预算 D 保持相同。

为了计算某些留出集的评估损失 L,针对数据并行方法研究人员使用当前模型,针对 DiLoCo 研究人员使用最新的全局模型。(注:留出集,是指从原始数据集中特意保留的一部分数据,用于评估模型的泛化性能。)



在谷歌张量处理单元 TPUv5e 和 TPUv6e 上进行实验

在具体的实验方法上,研究人员使用 Chinchilla-style 的仅解码器 Transformer 架构。

同时,使用 QKLayerNorm 来降低对于学习率的敏感性。需要说明的是,QKLayerNorm 是一种改进的层归一化技术,主要用于 Transformer 架构中的自注意力机制。另外,研究人员还使用 z 损失正则化来提高训练稳定性。

本次研究使用的词汇量为 32,768:其中有 32,000 个词汇表内的单词,另外还有一些额外的标记用于表示句子开头和词汇表外的内容。

实验中,研究人员将多个序列打包到每个 batch 中,整个 batch 的最大序列长度为 2048。

这时,研究人员开始针对一系列的模型进行训练,这些模型的 transformer 层数、注意力头数量、QKV 维度和前馈层隐藏维度各不相同。

其中,QKV 维度是指查询(Query)、键(Key)和值(Value)向量的维度。这些向量是自注意力机制的核心组成部分,用于计算输入序列中不同位置之间的关系。


图 | 每个模型比例的架构细节(来源:arXiv)

在大多数实验中,研究人员全程使用 C4 数据集的训练集来训练模型。

此外,还计算了 3 项任务的下游零样本评估指标:HellaSwag、Piqa 和 Arc-Easy。在进行过度训练消融时,使用了 Dolma 数据集。

其还使用了一个改进版的 NanoDO,它利用 DrJAX 在副本之间并行化内部训练步骤,并将模型副本轴暴露出来以便进行显式编程。这对于 JAX 中更好的扩展性能至关重要,因为 DrJAX 提供了 jax.vmap 的丰富版本,能够提供有关 DiLoCo 副本的更明确的分片信息。

与此同时,研究人员使用 all-reduce 进行外部优化。all-reduce 是一种分布式计算中的通信操作,用于在多个 GPU 或多个 TPU 之间高效地聚合数据。

研究中,其始终使用 bfloat16 来表示模型权重和梯度。bfloat16 是一种 16 位浮点数格式,它由 Google 团队开发,它在保留 32 位浮点数动态范围的同时,能够减少存储和计算开销。

研究人员在谷歌张量处理单元 TPUv5e 和 TPUv6e 上进行了大部分实验,并在 TPUv-5 上进行了最大规模的实验。

对于每个实验,其还计算了一个理想化的 wall-clock 训练时间,并在这一时间中考虑了理想化的计算时间和通信时间。同时,其还专门测量了端到端的 wall-clock 时间。

其发现,更大的水平并行化例如通过将 batch 大小加倍,将能够减少 wall-clock 时间。

实验中,研究人员假设其模型正在跨多个数据中心开展训练。当在数据中心之内的时候,能够拥有一个高带宽网络。当跨数据中心的时候,则分别拥有高带宽、中带宽或低带宽的网络。

当通信时间处于理想状态的时候,研究人员始终将高带宽网络用于数据中心内的网络,以及将三种网络中的任意一种用于跨数据中心网络。



DiLoCo:能实现更自然的水平扩展

利用相关数据,研究人员推导出了 Scaling Law,从而能够预测较大模型的评估损失和最优超参数。

根据 Chinchilla Scaling Law,研究人员假设最优 token 预算由 D=20N 给出。

这意味着在模型大小固定的情况下,如果将 batch 大小 B 加倍,则训练步骤的数量将会减半。

实验中,研究人员通过使用 Scaling Law 设置的超参数来训练 4B 和 10B 模型来验证上述观点。

尽管下图展示的是“插值”区域的情况,同时这也是大量扫描的结果,但是这些发现从定性角度来看,同样适用于“外推”区域。这使得研究人员在 M = 1、2 时,能够使用 DiLoCo 将 40 亿参数和 100 亿参数量的模型训练得拥有更低的评估损失。


(来源:arXiv)

结果显示,DiLoCo 能够实现更自然的水平扩展性。无论在任何情况,token 预算 D 仅仅是 N 的函数。

这意味着当使用 4 倍大的 batch 时,训练步骤能够减少到原来的 1/4。

对于 DiLoCo 来说,这会产生相当好的性能,并且可以同时使用更多资源,从而减少模型总训练时间。

相比之下,数据并行训练方法似乎需要更多的串行训练。因为,训练时间的减少与通信的减少是相辅相成的。

为了展示这些效果,研究人员在下图中绘制了在不同带宽的网络下训练时的理想化 wall-clock 时间。


(来源:arXiv)

借此发现,DiLoCo 对更大 batch 的容忍度使其能够更快地实现与数据并行方法大致相当的损耗。

在下图中,研究人员绘制了不同过度训练量下数据并行方法和 DiLoCo 的理想训练时间(M = 2)。


(来源:arXiv)

其发现,DiLoCo 能够通过降低通信成本和利用更大的 batch 来加速过度训练,因此需要的串行训练步骤更少。

这表明 DiLoCo 对于过训练来说是一个巨大的福祉,因为这意味着可以通过横向可扩展性来摊销计算时间。(注:过训练,是指模型在训练过程中过度拟合训练数据,导致其在未见数据上的性能下降。)

与此同时,研究人员的结果表明,DiLoCo 和数据并行训练这两种方法通常都是有效的,尽管没有明确的赢家,但是 M 之间的残差存在显著差异。

尽管如此,研究人员发现平均而言,虽然独立拟合在预测损失和全局批量大小方面略优于联合拟合,但独立拟合在预测内部学习率方面要明显好得多。

总的来说,本次结果表明与数据并行方法一样,DiLoCo 可以随着模型大小的变化实现可预测的扩展,这使得在极大的尺度上调整超参数和训练模型变得更加简单。

此外,DiLoCo 相比数据并行方法可以带来显著优势,包括:使用单个模型副本时具备更优的评估损失,在任何模型副本数量下最优 batch 大小都能增大。对于模型规模、过训练和同步频率来说,这将产生更好的鲁棒性。

不过,眼下显然需要可用于大规模部署 DiLoCo 等方法的系统和软件,并能在真实的超大模型环境中实现其通信效率优势。

参考资料:

https://arxiv.org/pdf/2503.09799

排版:初嘉实

ad1 webp
ad2 webp
ad1 webp
ad2 webp