主要观点总结
本文介绍了Google旗下团队对DiLoCo方法的研究，该方法是一种针对大模型的分布式训练优化方法。研究发现，DiLoCo在模型性能、评估损失、最优超参数等方面表现出良好的扩展性，优于数据并行训练方法。DiLoCo通过减少通信开销和提高扩展性，能够实现更自然的水平扩展。此外，该研究还涉及模型训练的实验细节和DiLoCo的优势。
关键观点总结
关键观点1: DiLoCo是一种针对大模型的分布式训练优化方法，具备减少通信开销和提高扩展性等优点。
DiLoCo能够增加最佳 batch 大小、提高下游规模的泛化能力，并在固定 token 预算下改善评估损失。随着模型尺寸的增大，DiLoCo呈现出可预测的稳健扩展。
关键观点2: DiLoCo在模型训练方面的表现优于数据并行训练方法。
在相同条件下，DiLoCo的模型规模扩展性优于数据并行训练方法。即使在通信效率更高的DiLoCo上，Scaling Law预测会比数据并行训练方法产生更低的评估损失。
关键观点3: DiLoCo具有更自然的水平扩展性。
DiLoCo能够实现更快速的训练，通过降低通信成本和利用更大的batch来加速过度训练，需要的串行训练步骤更少。
正文
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 也能胜过数据并行训练方法。
而本次研究中，基于相关数据研究人员为数据并行训练方法和 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 时间。
实验中，研究人员假设其模型正在跨多个数据中心开展训练。当在数据中心之内的时候，能够拥有一个高带宽网络。当跨数据中心的时候，则分别拥有高带宽、中带宽或低带宽的网络。
当通信时间处于理想状态的时候，研究人员始终将高带宽网络用于数据中心内的网络，以及将三种网络中的任意一种用于跨数据中心网络。
利用相关数据，研究人员推导出了 Scaling Law，从而能够预测较大模型的评估损失和最优超参数。
