谷歌团队再次证明Scaling Law的准确性，仍能超越数据并行训练

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 也能胜过数据并行训练方法。

研究人员：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）向量的维度。这些向量是自注意力机制的核心组成部分，用于计算输入序列中不同位置之间的关系。

在大多数实验中，研究人员全程使用 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 亿参数量的模型训练得拥有更低的评估损失。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

参考资料：

https://arxiv.org/pdf/2503.09799

排版：初嘉实