Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo

文章标题: 通信高效的语言模型训练可实现可靠且稳健的扩展：DiLoCo的缩放定律
作者: Zachary Charles, Gabriel Teston, Lucio Dery, Keith Rush, Nova Fallen, Zachary Garrett, Arthur Szlam, Arthur Douillard
机构: 1Google Research, 2Google Search, 3Google DeepMind

A1 主要贡献

随着我们扩展到更庞大的机器学习模型，数据并行方法固有的频繁同步需求造成了严重的减速，对进一步扩展构成了关键挑战。最近的研究【索引10，DiLoCo: Distributed low-communication training of language models，2023，arXiv，https://arxiv.org/abs/2311.08105】、【索引12 ，Streaming DiLoCo with overlapping communication: Towards a distributed free lunch，2025，arXiv，https://arxiv.org/abs/2501.18512】开发了一种方法（DiLoCo），该方法放宽了同步要求而不损害模型质量。然而，这些工作没有仔细分析DiLoCo的行为如何随模型大小变化 。

核心问题与研究目标：本文研究了在固定计算预算下训练大型语言模型（LLMs）时DiLoCo的缩放定律行为。我们关注算法因素（包括模型副本数量、超参数和token预算）如何以可通过缩放定律准确预测的方式影响训练。具体而言，本文旨在解决两个核心问题：(1) 如何预测作为模型大小函数的评估损失；(2) 如何预测给定模型大小的最佳超参数选择，从而避免昂贵的超参数调整。

主要贡献与创新点：

1. 为DiLoCo和数据并行训练建立缩放定律：本文基于广泛的实验数据，为数据并行训练和DiLoCo的评估损失及最优超参数建立了缩放定律。我们使用从3500万到24亿参数的9种不同模型尺寸，以及不同的DiLoCo副本数 M，对学习率、批大小和DiLoCo的“外部”学习率等超参数进行了详尽的实验。
2. 验证缩放定律的准确性：我们证明了这些缩放定律在推断到更大模型尺寸时，能很好地估计损失和最优超参数。利用这些定律预测的超参数，我们成功训练了40亿和100亿参数的模型，并验证了预测的准确性：DiLoCo的表现优于数据并行训练，同时将总通信量减少了超过100倍。

3. 揭示DiLoCo的普适性优势：我们的实验数据揭示了DiLoCo比以往文献记录的更广泛的优势：

   • 更大的最优批量大小：DiLoCo能够容忍比数据并行训练大得多的批量大小，从而提高水平可扩展性。
   • 随规模提升的泛化能力：随着模型规模的增大，DiLoCo在下游任务上的泛化能力相对于数据并行训练有更显著的提升。
   • 固定token预算下的更低评估损失：在相同的token预算下，DiLoCo能达到比数据并行训练更低的评估损失。

4. 即使在无通信瓶颈时DiLoCo也表现更优：一个令人惊讶的发现是，即使在通信不成问题的场景下，DiLoCo也表现出色。特别地，当副本数 M=1 时，DiLoCo（实际上是Lookahead优化器【索引60，Lookahead optimizer: k steps forward, 1 step back，2019，NeurIPS】的一个增强版本）在所有模型规模上都比数据并行训练取得了更低的评估损失和更好的训练时间，并且对更大的批量大小更具鲁棒性。

摘要中的核心观点：
* 更稳健 (Harder): DiLoCo的超参数在不同模型规模下是稳健且可预测的。
* 更优越 (Better): 随着模型规模的增加，DiLoCo相对于数据并行训练的优势进一步扩大。
* 更快速 (Faster): DiLoCo使用的带宽比数据并行训练少几个数量级。
* 更强大 (Stronger): DiLoCo能容忍比数据并行训练大得多的批量大小。

A3 背景知识

在整篇文章中，我们用 $θ$ 表示模型参数。我们用 $θ^{(t)}$ 表示在步骤 $t$ 时的模型。由于DiLoCo在 $M$ 个并行的模型参数上操作，我们将使用下标符号 $θ_m$ 来表示第 $m$ 个模型。当没有下标时，假定参数在所有DiLoCo副本之间是复制的。对于一批数据 $x$，我们用 $f(θ, x)$ 表示 $θ$ 在该批数据上的损失。

2.1 DiLoCo: 分布式低通信

DiLoCo的设计目标。DiLoCo【索引10，DiLoCo: Distributed low-communication training of language models，2023，arXiv，https://arxiv.org/abs/2311.08105】是一项专为存在通信约束下训练模型而设计的技术。它尤其适用于跨设备训练大型模型，而这些设备并非都通过低延迟带宽连接。为了避免产生延迟成本，DiLoCo并行训 练 $M$ 个模型（理想情况下，每个模型都使用通过低延迟带宽连接的共置计算资源进行训练），并仅每 $H$ 步同步一次模型。这类似于联邦学习中使用的FedOpt算法【索引48，Adaptive federated optimization，2021，ICLR】，但有一个重要的区别，即副本在各个轮次之间保持其内部优化器状态。

DiLoCo的双层优化框架。DiLoCo在多个模型上应用了一个双层优化框架：每个DiLoCo副本都有自己的模型 $θ_m^{(t)}$，并且有一个全局模型 $θ^{(t)}$。在每一步，每个副本都执行一个内部优化步骤（InnerOpt）。每隔 $H$ 步，每个副本计算 $∆_m^{(t)} = θ^{(t−H)} − θ_m^{(t)}$，即副本当前模型与最近更新的全局模型在参数空间上的差异。我们在所有副本上对这些增量（delta）进行平均，得到一个向量 $∆^{(t)}$，我们称之为外部梯度。我们将其视为外部模型的梯度估计，并对外部模型 $θ^{(t−H)}$ 应用一个外部优化步骤（OuterOpt）。这会产生一个更新后的外部模型 $θ^{(t)}$，该模型被广播到所有副本，并设置为它们当前的内部模型。我们在算法1中给出了DiLoCo的完整伪代码。

Table 1: 通用符号

Table 2: 算法特定符号

Figure 1: DiLoCo。每个DiLoCo模型副本独立训练H个内部优化步骤。这些模型通过一个外部优化步骤进行同步，通常涉及跨外部优化步骤的动量。在此图中，有M=4个副本。

2.2 数据并行与DiLoCo对比

两种训练方式的定义。在我们的工作中，我们将通过分布式数据并行训练（为简洁起见，我们称之为Data-Parallel）和DiLoCo进行模型训练。我们讨论它们的不同之处以及我们将在整个工作中使用的操作细节。在Data-Parallel中，我们在每一步取一个大小为 $B$ 的数据批次。在我们的工作中，与Kaplan等人【索引26，Scaling laws for neural language models，2020，arXiv，https://arxiv.org/abs/2001.08361】、Hoffmann等人【索引19 ，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】的工作类似，批大小指的是一个批次中的token数量（而不是序列数量）。然后我们计算一个批梯度，并以学习 率 $γ$ 应用优化。

DiLoCo的批次处理与优化步骤。在DiLoCo中，在每个步骤 $t$，我们取一个大小为 $B$ 的全局数据批次，并在序列级别上将其均匀地划分给 $M$ 个DiLoCo副本。因此，全局批大小为 $B$，但每个DiLoCo副本使用大小为 $B/M$ 的本地批次。与Data-Parallel类似，每个副本然后计算一个批梯度，并以学习率 $γ$ 应用一个内部优化步骤。与Data-Parallel不同，DiLoCo每 $H$ 步以学习率 $η$ 进行外部优化（对在参数空间中计算的外部梯度进行操作）。

M=1的DiLoCo与数据并行的区别。一个重要的比较是Data-Parallel与DiLoCo在 $M=1$ 时的对比。虽然相似，但它们并不完全相同。$M=1$ 的DiLoCo仍然有一个使用OuterOpt的外部优化器步骤，因此是Lookahead优化器【索引60，Lookahead optimizer: k steps forward, 1 step back，2019，NeurIPS】的一个变体。在我们的符号体系中，【索引60】对应于OuterOpt是学习率为 $η$ 的SGD的情况。在DiLoCo中，OuterOpt通常设置为带有Nesterov动量的SGD【索引10，DiLoCo: Distributed low-communication training of language models，2023，arXiv，https://arxiv.org/abs/2311.08105】，在这种情况下 ，$M=1$ 的DiLoCo成为Data-Parallel的一个变体，其中动量操作仅每 $H$ 步应用一次。

评估方法的一致性。在Data-Parallel和DiLoCo之间进行比较时，我们确保模型大小 $N$ 和总token预算 $D$ 是相同的。为了在某个留出集上计算评估损失 $L$，对于Data-Parallel，我们使用当前模型；对于DiLoCo，我们使用最新的全局模型。我们在表1中总结了与算法无关的符号，在表2中总结了特定于算法的符号。

A2 方法细节

算法1 DiLoCo

6.1 独立缩放定律

损失的缩放定律。我们首先拟合数据并行训练（Data-Parallel）所获得损失的缩放定律。我们通过幂律近似 $L(N) \approx AN^\alpha$ 来拟合数据并行的评估损失作为 $N$ 的函数。请注意，这可以通过对 $log(L)$ 应用线性拟合技术轻松完成，并且对 $A, \alpha$ 的初始值不敏感。得到的幂律在表7的第一行。

DiLoCo独立拟合的损失缩放定律。在进行独立拟合时，我们对DiLoCo也采用了上述方法。对于每个 $N, M$ 的值，我们记录所有超参数中最低的损失值。然后，我们可以为每个 $M$ 拟合一个幂律 $L_M(N) := L(N, M) \approx AN^\alpha$。结果在表7中给出。

Table 7: 损失的幂律近似 $L(N) \approx AN^\alpha$.

幂律拟合结果分析。结果显示，随着 $N$ 的增加，数据并行和DiLoCo在损失上的预测减少情况相似。值得注意的是，拟合的参数表明，在除了绝对最小的模型规模外，DiLoCo, M=1 的性能几乎都优于数据并行。这与第4节中讨论的结果相符。

超参数的缩放定律。对于数据并行，我们为学习率 $γ$ 和批大小 $B$ 拟合缩放定律。对于DiLoCo，我们为内部学习率 $γ$ 和全局批大小 $B$ 拟合缩放定律。鉴于它们在算法中的类似作用，我们以相同的方式拟合它们。对于（内部）学习率，我们使用与拟合损失缩放定律相同的方法：对于每个 $N$（以及DiLoCo的 $M$），我们选择最佳超参数，并拟合一个幂律。结果见表8。

批大小缩放定律的拟合方法。对于（全局）批大小，我们稍微改变了方法。如第3节所述，我们的扫描使用2的幂次方作为批大小，以饱和计算资源。然而，最优批大小可能位于这些值之间。为了解决这个问题，我们首先对批大小进行二次近似。具体来说，对于每个 $N$ 值，我们观察损失作为 $log_2(B)$ 的函数（当使用该 $B$ 的最佳学习率时），并对该函数进行二次拟合。我们选择这些二次函数的最小值，并对它们拟合一个关于 $N$ 的幂律。结果见表9。

外部学习率的特殊性。DiLoCo有第三个可以拟合缩放定律的超参数：外部学习率。然而，如第4节所示，最优的外部学习率（对于足够大的模型）似乎是恒定的。因此，一个缩放定律似乎不会比简单地为每个 $M$ 使用最佳外部学习率（见图7）提供任何改进的预测性能。

6.2 联合缩放定律

联合幂律拟合。作为另一种方法，我们可以为DiLoCo的各个方面拟合联合幂律，使用一个双变量幂律 $f(N, M) \approx AN^\alpha M^\beta$。我们对损失 $L$、内部学习率 $γ$ 和全局批大小 $B$ 都这样做。对于前两者，我们为每个 $N, M$ 的值选择最佳的学习率和损失。对于批大小，我们做同样的事情，但使用上一节的二次近似。然后，我们可以通过标准的线性回归技术拟合一个联合幂律。得到的幂律在表10中。与独立拟合一样，我们不尝试对外部学习率建模，因为最优值与 $N$ 无关。

Table 8: （内部）学习率的幂律近似 $γ(N) \approx AN^\alpha$.

Table 9: （全局）批大小的幂律近似 $B(N) \approx AN^\alpha$.

Table 10: DiLoCo的损失L、内部学习率γ和批大小B的联合幂律近似 $f(N, M) = AN^\alpha M^\beta$.

6.3 拟合优度度量

拟合方法比较。既然我们有两种不同的方法来为DiLoCo建立缩放定律，我们可以尝试回答哪一种能提供更好的预测。首先，我们通过留一法交叉验证来做到这一点。具体来说，我们使用与上述相同的方法来拟合 $L, γ, B$ 的缩放定律，但只使用到 $N=13$亿参数的数据，排除了我们关于 $N=24$亿参数的数据。然后，我们使用这个缩放定律来预测在 $N=24$亿参数时 $L, γ, B$ 的最优值，针对不同的 $M$ 值。

残差计算方法。给定一个预测值 $\tilde{y}$ 和一个实际最优值 $y$，我们将预测的残差计算为对数的平均绝对误差：$res(y, \tilde{y}) = |log(y) − log(\tilde{y})|$。我们使用这个度量，因为它能同时很好地适用于这三个变量，尽管它们的尺度差异很大。对于每个 $M \in \{1, 2, 4, 8\}$，我们计算上述三个参数的预测值，并测量相对于 $N=24$亿时实际最优值的残差。我们在表11中报告了这些值以及它们在 $M$ 上的平均值。

拟合结果评估。我们的结果通常表明，两种方法大体上都是有效的（因为没有明显的赢家），但也表明不同 $M$ 之间的残差存在显著差异。尽管如此，我们看到平均而言，虽然独立拟合在预测损失和全局批大小方面略好一些，但独立拟合在预测内部学习率方面明显更好。

6.4 外推至更大模型

超参数预测与实验。我们使用独立和联合拟合来预测数据并行和DiLoCo在 $M \in \{1, 2, 4\}$ 时，模型规模为40亿和100亿的最佳超参数。请注意，对于数据并行，我们只能使用独立拟合。我们使用这些超参数和Chinchilla最优的token预算在这些模型上进行训练，并比较结果。

外推结果分析。我们观察到两个重要的方面。首先，与之前的结果不同，在40亿和100亿的规模下，我们发现无论使用独立还是联合拟合方法，$M=2$ 的DiLoCo实际上都优于数据并行和 $M=1$ 的DiLoCo。其次，我们发现 $M=1$ 的DiLoCo需要联合拟合才能比数据并行做得更好。除了这种情况外，联合和独立拟合在整个过程中表现相当。总而言之，联合拟合方法在超参数外推方面似乎比独立拟合略有优势。结合其也能外推到更大的 $M$ 的能力，我们通常推荐对所有超参数使用联合拟合方法。

Table 11: 联合拟合缩放定律与独立拟合相当或更优。此处我们给出了在N=2.4B和不同M值下缩放定律预测的残差。我们比较了独立和联合拟合策略在预测损失L、内部学习率γ和全局批大小B方面的残差。对于平均残差，我们高亮显示了独立或联合拟合中残差较低者。我们看到联合拟合在L和B上与独立拟合相当，但在预测γ方面做得更好。

损失预测验证。我们现在使用这些损失值来比较它们与上面拟合的缩放定律。我们通常发现损失值被预测得非常好，与缩放定律预测的损失相差在几个百分点之内。我们在图13中展示了拟合的缩放定律和外推的损失值。

Figure 13: DiLoCo缩放定律能很好地外推到更大的模型。我们展示了数据并行和DiLoCo的损失缩放定律。图中既有用于形成缩放定律的损失值（通过训练高达2.4B的模型），也有在更大模型（4B和10B）上获得的损失值。为简单起见，我们展示了独立拟合的缩放定律，但联合拟合也能很好地预测损失。

Table 12: 联合拟合超参数能很好地外推到更大的模型。此处我们展示了在4B和10B模型上的评估结果，使用了由独立和联合缩放定律预测的超参数。我们高亮显示了比数据并行更好（即更低）的DiLoCo评估结果。我们看到，虽然M=2的DiLoCo使用独立或联合拟合规则都比数据并行做得好，但DiLoCo M=1仅在使用联合拟合时才表现更好。

6.5 参数函数拟合

拟合更复杂函数的动机。虽然幂律具有有用的特性，但关于缩放定律的各种工作常常发现拟合更复杂的函数到数据上是有益的【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】。特别地，一个 像 $AN^\alpha$ 这样的幂律在 $N \to \infty$ 时只会趋向于0、1或发散到 $\infty$。在许多情况下，拟合更复杂的函数是有意义的。

联合缩放定律的参数形式探索。我们对联合缩放定律的参数形式特别感兴趣。虽然Hoffmann等人【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】使用风险分解论证将损失分解 为 $N$ 和 $D$ 的函数，但如何分解 $L(N, M)$ 并不立即清楚。为此，我们采用一种经验方法。我们开发候选函数，并确定哪种函数在外推意义上表现最好。我们使用以下函数形式：
1. $L(N, M) \approx AN^\alpha M^\beta$
2. $L(N, M) \approx AN^\alpha M^\beta + C$
3. $L(N, M) \approx AN^{\alpha+\beta M} + C$
4. $L(N, M) \approx AN^\alpha + BM^\beta + C$

参数拟合方法。第一个被包含进来是为了比较，因为它恢复了上面使用的幂律缩放定律。拟合更复杂的函数可能对参数的初始值更为敏感，也对异常数据敏感。因此，当将这些函数拟合到我们上面的损失值时，我们使用Hoffmann等人【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】提出的通用策略 。

详细拟合过程。具体来说，让 $f_Q(N, M)$ 表示上述函数形式之一，其中 $Q$ 代表要拟合的参数集（例如，对于第一个函数，$Q = \{A, \alpha, \beta\}$）。让 Huber$_\delta$ 表示参数为 $\delta$ 的Huber损失。让 $\mathcal{N}, \mathcal{M}$ 表示所考虑的 $N$ 和 $M$ 的值集。对于每个 $N, M$，我们有一个经验损失 $L(N, M)$，和一些损失的估计 $f_Q(N, M)$。然后我们解决以下最小化问题：

我们通过L-BFGS最小化这个问题，使用一些参数的初始化 $Q_0$。我们对256个随机的参数初始化 $Q_0$ 重复这个过程。我们保留所有在 $N=2.4B$ 规模下的损失值，并选择最能拟合保留数据的参数 $Q$，以所有 $M$ 上的平均残差 $|log f_Q(N, M) − log L(N, m)|$ 来衡量。

参数函数拟合结果。我们发现幂律（第1行）和加性分解（第4行）在外推损失值方面明显比更细致的参数形式差。我们注意到，加性分解类似于Hoffmann等人【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】使用的将损失分解为模型大 小 $N$ 和token预算 $D$ 的函数，但似乎不能反映 $M$ 如何影响DiLoCo的损失。我们将其作为一个开放问题，以确定何种参数形式能更好地预测损失，并能通过对通信高效训练的理论理解来解释。

Table 13: 参数函数拟合改善了联合缩放定律。我们展示了经验损失函数 $L(N, M)$ 的各种参数近似，以及它们在 $N = 2.4B$ 规模的留出损失数据上的验证误差。我们看到，联合幂律（第一个）和经典损失分解（最后一个）在预测留出数据上的损失方面明显更差。

A4 实验环境

模型架构：
* 使用Chinchilla风格的仅解码器transformer架构【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】 。
* 为降低对学习率的敏感性，使用了QK-LayerNorm【索引58，Small-scale proxies for large-scale transformer training instabilities，2023，arXiv，https://arxiv.org/abs/2309.14322】 。
* 为增加训练稳定性，使用了z-loss正则化【索引8，PaLM: Scaling language modeling with pathways，2023，Journal of Machine Learning Research】。
* 词汇表大小为32,768。最大序列长度为2,048。
* 所有模型均从头开始预训练。
* 模型规模从3500万到100亿参数不等，具体架构参数见下表3。

Table 3: 模型详情。我们展示了每个训练模型的规模、层数、层维度和token预算。对于较大的模型（4B和10B），我们使用缩放定律来预测最佳超参数，而不是进行广泛的超参数调整。

数据集：
* 预训练：主要使用C4数据集【索引45，Exploring the limits of transfer learning with a unified text-to-text transformer，2020，Journal of machine learning research】的训练集，并在其留出的验证集上报告评估指标。
* 下游任务评估：在HellaSwag【索引59，Hellaswag: Can a machine really finish your sentence?，2019，arXiv】，Piqa【索引3，Piqa: Reasoning about physical commonsense in natural language，2020，AAAI】，和Arc-Easy【索引9，Think you have solved question answering? try arc, the AI2 reasoning challenge，2018，arXiv】上进行零样本评估。
* 过训练实验：使用Dolma数据集【索引52，Dolma: An open corpus of three trillion tokens for language model pretraining research，2024，arXiv】。

硬件配置：
* 主要在TPUv5e和TPUv6e上进行实验。
* 最大规模的实验在TPUv5上进行。

软件配置：
* 实现：使用修改版的NanoDO【索引31，NanoDO: A minimal transformer decoder-only language model implementation in JAX.，2024c，http://github.com/google-deepmind/nanodo】，该版本使用DrJAX【索引50 ，DrJAX: Scalable and differentiable mapreduce primitives in JAX，2024，WANT@ICML】在JAX【索引4，JAX: composable transformations of Python+NumPy programs，2018，http://github.com/jax-ml/jax】中并行化内部训练步骤 。
* 优化器：
* Data-Parallel和DiLoCo的内部优化器使用AdamW【索引32，Decoupled weight decay regularization，2017，arXiv】，其中 $\beta_1=0.9, \beta_2=0.99$。
* DiLoCo的外部优化器使用带Nesterov动量的SGD【索引54，On the importance of initialization and momentum in deep learning，2013，ICML】，动量项为0.9，采用恒定外部学习率。

• 训练设置：

  • 1000步的warmup，然后是cosine学习率衰减（衰减至峰值的5%）。
  • 权重衰减参数 $\lambda$ 设置为 $T^{-1}$，其中 $T$ 是总训练步数。
  • 内部梯度裁剪到全局 $\ell_2$ 范数为1，外部梯度不裁剪。
  • DiLoCo的同步周期 $H$ 默认为30。
  • 模型权重和梯度使用bfloat16表示。

• 缩放定律实验：

  • 对35M至2.4B参数的模型进行了全面的超参数扫描（学习率 $\gamma$，批大小 $B$，外部学习率 $\eta$）。
  • DiLoCo的副本数 $M$ 取值为{1, 2, 4, 8}。
  • 遵循Chinchilla缩放定律【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】，token预算设 为 $D=20N$。

A4 实验结果

核心发现

发现1：DiLoCo随模型规模表现更优

• 实验内容：比较了数据并行（Data-Parallel）与不同副本数（M）的DiLoCo在不同模型规模（N）下的评估损失。
• 实验结果：如图2和表4所示，随着模型规模N的增加，DiLoCo相对于数据并行的评估损失持续改善。对于所有M值，DiLoCo与数据并行的损失百分比差异随N增大而单调递减。例如，当M=4时，在35M模型上DiLoCo损失高2.0%，而在2.4B模型上仅高0.26%。特别地，在N=2.4B时，M=2的DiLoCo已取得了比数据并行更低的损失。
• 分析结论：这一趋势通过在4B和10B模型上的外推实验得到验证（表5），其中M=1和M=2的DiLoCo均取得了更低的损失，证实了DiLoCo的优势随模型规模扩大而增强。

(a) 不同算法的评估损失随N变化的函数。

(b) 相对于数据并行的评估损失百分比差异。

Figure 2: DiLoCo随规模扩大表现更佳。我们比较了不同模型规模N下数据并行与DiLoCo。对于所有M，DiLoCo相对于数据并行的改善随N增加而单调增强。

Table 4: 不同模型规模下数据并行和DiLoCo的评估损失。我们改变模型规模N和副本数M，报告评估损失L以及相对于数据并行(DP)的百分比差异。我们用粗体标出DiLoCo取得比数据并行更低损失的设置。

Table 5: DiLoCo在更大规模上优于数据并行。此处我们展示了使用缩放定律预测的超参数在4B和10B模型上的评估结果。我们用粗体标出DiLoCo达到比数据并行更低损失的设置。

发现2：单副本DiLoCo（M=1）的优势
* 实验内容：专门比较了数据并行与M=1的DiLoCo在不同模型规模和全局批大小下的评估损失和下游任务准确率。
* 实验结果：如图3所示，在几乎所有设置下，M=1的DiLoCo都取得了比数据并行更低的评估损失和更高的零样本评估准确率。
* 分析结论：M=1的DiLoCo性能对批大小变化表现出更强的稳定性。数据并行的性能会因批大小加倍或四倍而大幅下降，而这对M=1的DiLoCo影响甚微。这表明即使没有通信减少的优势，DiLoCo的优化机制本身也更优越。

Figure 3: DiLoCo (M=1)比数据并行泛化能力更好。我们展示了数据并行和DiLoCo (M=1)在不同模型和全局批大小（以token计）下的评估损失和下游准确率。在所有设置中，DiLoCo (M=1)都比数据并行表现更好，且差距随批大小增加而增大。其他模型规模也有类似结果，为简洁起见省略。

发现3：DiLoCo提升了最优批大小
* 实验内容：研究了数据并行和不同M值的DiLoCo在不同全局批大小下的评估损失和下游任务准确率。
* 实验结果：如图4和图5所示，DiLoCo的最优全局批大小远大于数据并行，并且最优批大小随M的增加而增加。数据并行的性能在小批次时较好，但随批次增大迅速恶化；相比之下，DiLoCo对批大小变化的性能更稳定。
* 分析结论：这提高了水平可扩展性。由于总token预算固定，使用更大的批大小意味着更少的训练步骤，从而可以利用更多计算资源并行处理，显著减少端到端的训练时间。如图6所示，在不同带宽的网络环境下，DiLoCo都因其通信效率和对大批次的容忍度而训练得更快。

Figure 4: DiLoCo增加了最优批大小，第一部分。数据并行和DiLoCo的评估损失随全局批大小（以token计）变化的函数。对于所有M，DiLoCo都展现出比数据并行更大的最优批大小。此外，最优批大小随M增加而增加。其他模型规模也有类似结果，为简洁起见省略。

Figure 5: DiLoCo增加了最优批大小，第二部分。数据并行和DiLoCo在HellaSwag上的零样本准确率随全局批大小（以token计）变化的函数。即使在较小的模型规模下，M=2的DiLoCo在较大的全局批大小时也能获得更高的准确率。其他模型规模也有类似结果，为简洁起见省略。

(a) 带宽为10 Gbit/s，延迟为10^-2秒的网络（低带宽）。

(b) 带宽为100 Gbit/s，延迟为10^-3秒的网络（中带宽）。

(c) 带宽为400 Gbit/s，延迟为10^-4秒的网络（高带宽）。

Figure 6: 无论有无通信瓶颈，DiLoCo训练都更快。我们绘制了数据并行和DiLoCo在不同模型规模下，跨高、中、低带宽网络连接的计算节点进行训练的理想化端到端训练时间（见附录A）。对于高达2.4B的模型，我们还改变了全局批大小。对于4B和10B模型，我们使用缩放定律预测的批大小，相关讨论在第6节。由于其减少的通信量和对更大批大小的容忍度，DiLoCo在几乎所有设置中都更快。即使在高带宽网络中，更大的批大小也减少了端到端训练时间。M≥4和更小模型有类似结果，为清晰起见省略。

发现4：外部学习率的规律
* 实验内容：分析了DiLoCo在不同模型规模N和副本数M下的最优外部学习率 $\eta$。
* 实验结果：如图7所示，对于足够大的模型（N≥335M），最优的外部学习率 $\eta$ 与模型大小N无关，仅依赖于副本数M。并且，更大的M值似乎需要更大的 $\eta$。
* 分析结论：这个发现简化了超参数调整，因为可以在较小的模型规模上确定最优的 $\eta$ 并将其推广到更大的模型。这与联邦学习中外部学习率应随客户端数量增加而增加的发现一致【索引6，On large-cohort training for federated learning，2021，NeurIPS】。

Figure 7: 最优外部学习率与模型大小无关。我们展示了DiLoCo在不同M和模型规模N下的最佳外部学习率。我们在{0.2, 0.4, 0.6, 0.8, 1.0}中选择最佳外部学习率，同时优化内部学习率γ和全局批大小B。对于足够大的模型，最佳外部学习率显然是恒定的。

消融实验

同步周期的影响 (H)
* 实验内容：在固定内部学习率和批大小的情况下，对同步周期 $H \in \{1, 5, 10, 30, 100, 300\}$ 和外部学习率 $\eta$ 进行扫描。
* 实验结果：
1. 最优外部学习率 $\eta$ 随 $H$ 和 $M$ 单调增加，但与模型规模 $N$ 基本无关（图8）。
2. $H=1$（每步同步）性能最差。对于 $H>1$ 的情况，评估损失随 $H$ 的增加而增加，但这种增加对于 $M=1$ 和更大的模型来说不那么明显（图9）。

• 分析结论：对于大型模型，可以采用非常稀疏的同步（更大的H）而几乎不影响性能，这对于在广域网上分布式训练至关重要。

计算利用率
* 实验内容：模拟了在不同网络带宽下，不同同步周期 $H$ 的DiLoCo和数据并行训练10B、405B和671B参数模型时的计算利用率。
* 实验结果：如图10和表6所示，DiLoCo显著提高了计算利用率。通过增加同步周期 $H$，即使在低带宽网络下，也可以达到很高的计算利用率（例如80%），而数据并行则需要高得多的带宽。
* 分析结论：DiLoCo能将带宽需求降低10到100倍，极大地提升了在有通信瓶颈环境下的训练效率。

过训练（Overtraining）的鲁棒性
* 实验内容：在使用Chinchilla最优token预算的1倍、4倍、16倍进行过训练时，比较了数据并行和DiLoCo的性能。
* 实验结果：如图11所示，在不同的过训练倍率下，各种算法的损失随FLOPs变化的曲线在对数空间中基本保持平行。DiLoCo（特别是M=1）在所有过训练量和模型规模下都略优于数据并行。
* 分析结论：DiLoCo的优势对于过训练是鲁棒的，并且由于其对大批次的容忍度，能够通过水平扩展显著加速长时间的过训练过程（图12）。

A7 补充细节

5.1 同步频率

实验设置。我们之前的实验都使用了 $H=30$ 的同步频率（即每进行 $H=30$ 次内部优化步骤，DiLoCo执行一次外部优化步骤）。我们现在进行一项消融研究，以确保我们的结果与其他 $H$ 值一致。为此，我们在不同的 $M$ 和各种模型规模 $N$ 上应用DiLoCo。对于每种设置，我们采用上述最佳的内部学习率和全局批大小，但对 $H \in \{1, 5, 10, 30, 100, 300\}$ 和 $\eta \in \{0.05, 0.1, 0.2, 0.4, 0.6, 0.8, 1.0\}$ 进行扫描。我们首先研究第4节中的观察，即 $\eta$ 应独立于 $N$ 进行调整，是否对其他 $H$ 值也成立。

(a) 每个同步频率的最佳外部学习率。阴影区域表示模型规模间的方差。

(b) 每个模型规模的最佳外部学习率。阴影区域表示同步频率间的方差。

Figure 8: 外部学习率随M和H扩展，而非N。最佳外部学习率η是同步频率H和M的单调递增函数（左图），但基本上与模型规模N无关（右图）。与图7一样，这意味着我们可以在较小的模型规模上调整外部学习率。

外部学习率的规律。我们在图8中给出了肯定的结果。我们看到，在各种模型规模下，最佳学习率基本上只是副本数 $M$ 和同步频率 $H$ 的函数，而与模型规模 $N$ 基本无关。虽然存在一些轻微的变化，但这可能是由于没有重新调整内部学习率所致。此外，我们的结果实际上揭示了一个可能令人惊讶的现象：最佳外部学习率随 $H$ 增加而增加。这可能与直觉相反；因为随着 $H$ 增加，DiLoCo副本可能会发散得更多，所以人们可能期望需要一个更保守的学习率。但事实并非如此。

同步频率对损失的影响。接下来，我们在图9中分析了 $H$ 如何影响DiLoCo的评估损失。对于所有模型，每步同步（$H=1$）表现最差，但在此之后，所有 $H$ 值的表现都相当。对于固定的 $N$ 和 $M$，评估损失随 $H$ 的增加而增加。然而，这种增加对于 $M=1$ 和更大的模型来说不太明显。这得出了一个重要发现：随着模型规模 $N$ 的增加，我们可以实际上更不频繁地在DiLoCo副本之间进行同步，同时几乎保持评估性能。

Figure 9: 不频繁同步对大模型效果更好。除了表现最差的H=1之外，评估损失随H增加而增加。然而，对于M=1的DiLoCo和更大的模型，增加率不那么明显，这表明大型模型可以非常不频繁地进行同步。

计算利用率。DiLoCo的同步频率对于在全球范围内分布的大规模模型训练至关重要。确实，较不频繁的同步（更大的 $H$）降低了训练的带宽要求。我们参照Douillard等人【索引12，Streaming DiLoCo with overlapping communication: Towards a distributed free lunch，2025，arXiv，https://arxiv.org/abs/2501.18512】的方法，模拟了要达到尽可能高的计算利用率（计算时间/总时间）所需带宽量，针对三种类型的LLMs：一个10B的Chinchilla风格transformer【索引19 ，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】（图10a），一个405B的Llama3模型【索引13 ，The llama 3 herd of models，2024，arXiv】（图10b），以及一个671B的DeepSeek-v3 MoE【索引29，Deepseek-v3 technical report，2024a，arXiv】（图10c）。我们也在表6中报告了原始数据。

Figure 10: DiLoCo极大地提高了计算利用率。这里我们展示了DiLoCo和数据并行在一系列带宽和同步频率H下的模拟计算利用率。计算利用率为0.8意味着80%的时间用于计算，20%用于通信。数值越高越好。我们在其他模型规模上也看到了类似的结果，但为了视觉清晰而省略。

Table 6: 模拟的计算利用率。我们根据所需的flops，使用Kaplan等人[2020]提出的规则和60%的最大flop利用率来估计步长时间。我们使用[Douillard等人, 2025]的模拟器来估计达到一定计算利用率水平所需的带宽（以Gbit/s为单位）。我们用浅蓝色突出显示带宽减少10倍，深蓝色突出显示减少100倍。

5.2 过训练（Overtraining）

过训练的定义。在上述所有实验中，我们对每个模型规模都使用了Chinchilla最优的token数量【索引19，Training compute-optimal large language models，2022，arXiv，https://arxiv.org/abs/2203.15556】。然而，在许多情况下（例如，考虑到推理成本训练模型时），进行过训练是有益的，即我们使用更多的token【索引16 ，Language models scale reliably with over-training and on downstream tasks，2024，arXiv】。我们希望确保我们的结果对过训练量是稳健的。为此，我们对各种过训练乘数进行了一系列的消融研究。给定一个过训练乘数 $\lambda \ge 1$，我们在 $D = 20N\lambda$ 个token上进行训练，因此 $\lambda=1$ 对应于Chinchilla最优的token数量。

过训练实验设置。在我们的实验中，对于不同的模型规模 $N$ 和算法，我们采用了我们上述结果中表现最好的超参数（即（全局）批大小 $B$、（内部）学习率 $\gamma$ 和使用DiLoCo时的外部学习率 $\eta$）。然后，我们在Dolma数据集【索引52，Dolma: An open corpus of three trillion tokens for language model pretraining research，2024，arXiv】上训练 $D = 20N\lambda$ 个token，因为它比C4有更多的token。由于我们使用QK-layernorm，通常不需要重新调整学习率【索引16，Language models scale reliably with over-training and on downstream tasks，2024，arXiv】。我们获取得到的模型，并在与之前相同的评估集上进行评估，即C4的验证分割。我们将过训练乘数 $\lambda$ 在 $\{1, 4, 16\}$ 中变化。请注意，由于我们更改了底层数据集，因此需要对 $\lambda = 1$ 进行重新训练。我们使用Data-Parallel以及 $M \in \{1, 2, 4\}$ 的DiLoCo进行训练。

Figure 11: DiLoCo在过训练中表现稳定可靠。对于每个过训练乘数，损失相对于FLOPs的曲线在对数空间中基本上都是平行的线。我们看到，改变算法带来的损失变化远小于改变模型大小或过训练量。此外，在所有过训练量和模型大小下，M=1的DiLoCo都略优于包括数据并行在内的所有其他算法。

过训练结果分析。我们的结果呈现在图11中。我们看到，随着我们进行更多的过训练，缩放行为在性质上基本保持不变。我们强调，在这些实验中我们没有重新调整任何超参数。对于每个模型大小和算法，我们只是采用了我们在C4上进行的Chinchilla最优实验中表现最好的超参数。这意味着DiLoCo在我们进行过训练时的稳定性得以保持，尽管对于 $M > 1$ 的情况，我们使用的批大小比Data-Parallel大得多。

过训练的训练时间。为了说明这一点，我们在图12中绘制了Data-Parallel和DiLoCo（$M=2$）在不同过训练量下的理想化训练时间。DiLoCo通过减少通信成本和利用更大的批大小来加速过训练，因此需要更少的串行训练步骤。这表明DiLoCo在过训练中是一个巨大的福音，因为我们可以通过水平可扩展性来分摊计算时间（对于过训练来说，计算时间可能相当长）。

(a) 带宽为10 Gbit/s，延迟为10^-2秒的网络（低带宽）。

(b) 带宽为100 Gbit/s，延迟为10^-3秒的网络（中带宽）。

(c) 带宽为400 Gbit/s，延迟为10^-4秒的网络（高带宽）。

Figure 12: DiLoCo通过利用水平可扩展性加速过训练。对于数据并行和DiLoCo M=2，我们绘制了在连接高、中、低带宽网络的计算节点上进行训练的理想化端到端时间（见附录A）。对于每种算法和过训练量，显示的线条代表从335M到2.4B参数不等的模型规模。由于其减少的通信量和对更大批大小的容忍度，DiLoCo在所有设置中都更快。即使在高带宽设置下，更大的批大小也增加了水平可扩展性，减少了端到端训练时间。对于M≥4和更小的模型，我们看到了类似的结果，但为清晰起见省略了。

A5 结论

本文的研究结果表明，与数据并行（Data-Parallel）类似，DiLoCo能够以可预测的方式随模型规模扩展，这使得在极大尺度上调整超参数和训练模型变得更加简单。此外，DiLoCo相比于数据并行能提供显著的优势，包括：
1. 更优的评估损失：在使用单个模型副本（M=1）时，DiLoCo能达到比数据并行更低的评估损失。
2. 更大的最优批量大小：对于任何数量的模型副本，DiLoCo都支持更大的最优批量大小，从而增强了水平扩展能力。
3. 广泛的鲁棒性：这些优势对于模型规模、过训练量和同步频率的变化都具有鲁棒性。

未来工作展望：
1. 扩展缩放定律分析：可以将已用于数据并行缩放定律的各种方面，如对下游任务的分析【索引16，Language models scale reliably with over-training and on downstream tasks，2024，arXiv】、数据受限的缩放定律【索引36，Scaling data-constrained language models，2023，NeurIPS】以及推理成本【索引51，Beyond chinchilla-optimal: Accounting for inference in language model scaling laws，2023，arXiv】，来增强现有的分析。
2. 适应DiLoCo的改进方法：可以将缩放定律调整以涵盖对DiLoCo及相关方法的改进提议，包括异步更新【索引30，Asynchronous local-SGD training for language modeling，2024b，arXiv】、流式DiLoCo【索引12，Streaming DiLoCo with overlapping communication: Towards a distributed free lunch，2025，arXiv】以及与训练方法协同设计的模块化架构【索引11，DiPaCo: Distributed path composition，2024，arXiv】、【索引22，Training neural networks from scratch with parallel low-rank adapters，2024，arXiv】。
3. 系统与软件支持：明确需要能够用于在实际的超大规模环境中部署像DiLoCo这类方法并实现其通信效率优势的系统和软件【索引24，OpenDiLoCo: An open-source framework for globally distributed low-communication training，2024b，arXiv】。

A6 附录

A.1 计算时间

计算时间定义。这里，我们指的是模型训练中浮点运算所花费的时间，忽略了节点间的通信时间（我们将在下面的部分单独处理）。我们使用理想化模型，其中总FLOPs $C=6ND$。给定 $R$ 个芯片，每个芯片每秒可以执行 $Q$ 次浮点运算，总计算时间下界为 $C/RQ$。芯片数量 $R$ 纯粹是模型大小 $N$ 和全局批大小 $B$ 的函数。芯片数量不依赖于算法（数据并行或DiLoCo）或使用DiLoCo时的模型副本数。

A.2 通信时间

通信时间模型。网络连接由带宽 $W$ 和延迟 $ϵ$ 表征。当在 $R$ 个计算节点上对 $N$ 个参数执行all-reduce时，参与all-reduce的至少一个计算节点发送和接收的流量下界为 $2N(1-R^{-1})$【索引43，Bandwidth optimal all-reduce algorithms for clusters of workstations，2009，J. Parallel Distrib. Comput.】。这类算法被称为带宽最优算法。由于节点间的通信是同步但并行完成的，在一个每对节点间带宽为 $W$、延迟为 $ε$ 的网络中，完成all-reduce的时间至少为：

$$\frac{2 N}{W}\left(1-\frac{1}{R}\right)+\varepsilon.$$

跨数据中心通信模型。DiLoCo【索引10，DiLoCo: Distributed low-communication training of language models，2023，arXiv，https://arxiv.org/abs/2311.08105】是为模型太大无法容纳在单个数据中心，因此必须跨由低带宽连接的计算岛进行训练的场景而设计的。为了对此建模，我们假设我们 在 $R$ 个计算节点（通常是GPU或TPU）上进行训练。其中一些通过数据中心内部的网络连接，另一些则跨数据中心连接。我们让 $W_0, ε_0$ 表示数据中心内部网络的带宽和延迟，而 $W_1, ε_1$ 同样定义为跨数据中心网络。通常，$W_0 \ge W_1, ε_0 \le ε_1$。

数据并行（Data-Parallel）通信时间。在每个训练步骤 $T$，我们必须在所有 $R$ 个计算节点上执行一次all-reduce。由于一些节点是跨数据中心连接的，总通信时间至少为：

$$\left(\frac{2N}{W_1}\left(1-\frac{1}{R}\right)+\varepsilon_1\right)T$$

DiLoCo, M=1 通信时间。在每个内部步骤 $T$，我们像在数据并行训练中一样，在所有 $R$ 个设备上执行一次all-reduce。我们还每 $H$ 步为外部优化进行一次all-reduce。其中一些节点是跨数据中心连接的，因此每次all-reduce的通信时间至少为 $2N(1 - R^{-1})W_1^{-1} + ε_1$。因此，总通信时间至少为：

$$\left(\frac{2 N}{W_{1}}\left(1-\frac{1}{R}\right)+\varepsilon_{1}\right) T\left(1+\frac{1}{H}\right) .$$

DiLoCo, M≥2 通信时间。我们假设 $M$ 个模型副本中的每一个都在同一数据中心内连接的计算节点上训练。在每个内部步骤 $T$，每个模型副本由 $R/M$ 个设备训练，这些设备需要进行一次all-reduce。然而，数据中心之间没有通信发生，因此每个内部步骤的通信时间由 $2N(1 - MR^{-1})W_0^{-1} + ε_0$ 界定。

每个外部优化步骤都涉及在所有 $R$ 个跨数据中心连接的设备上进行all-reduce。这会产生至少 $2N(1 - R^{-1})W_1^{-1} + ε_1$ 的通信时间。由于它只每 $H$ 步发生一次，总通信时间的下界为：

$$\left(\frac{2 N}{W_{0}}\left(1-\frac{M}{R}\right)+\varepsilon_{0}\right) T+\left(\frac{2 N}{W_{1}}\left(1-\frac{1}{R}\right)+\varepsilon_{1}\right) \frac{T}{H}$$

请注意，这表明只要 $H \ge W_0/W_1$，外部通信步骤最多只产生总通信成本的一半。

流式DiLoCo (Streaming DiLoCo)。我们注意到，上面计算的成本也适用于流式DiLoCo【索引12，Streaming DiLoCo with overlapping communication: Towards a distributed free lunch，2025，arXiv，https://arxiv.org/abs/2501.18512】。虽然内部步骤保持不变，但外部步骤被平滑处理，使得每个片 段 $p \in \{1, \dots, P\}$ 每 $H$ 步进行一次。然而，片段通信被偏移，使得某个片段每 $H/P$ 步进行通信，导致通信分摊到相同的每步成本。这是预期的，因为流式DiLoCo减少了任何一步的峰值通信，但并未减少整个训练过程中的总通信量。

重叠通信。Douillard等人【索引12，Streaming DiLoCo with overlapping communication: Towards a distributed free lunch，2025，arXiv，https://arxiv.org/abs/2501.18512】的另一个贡献是能够将外部优化器所需的通信与计算重叠，通过在外部优化器中使用一个旧版本的片段，并将此外部优化的结果与本地优化的片段合并。例如，如果计算时间主导通信时间，这将允许在计算端到端时间时省略通信项。这种设置在算法角度上是不同的，因此其对扩展性的影响需要独立研究 。

A.3 总端到端时间

总时间计算。总端到端时间是上述计算时间和通信时间的总和。要测量通信时间，我们必须知道每个实验使用的芯片数量 $R$，每个芯片每秒的FLOPs数量 $Q$，以及数据中心内部和跨数据中心网络的带宽和延迟。对于计算成本，我们根据我们的实验使用一个稍微理想化的芯片数量 $R$，但确保将全局批大小加倍会使芯片数量加倍。我们基于TPU v5e和v6e芯片的FLOPs能力的公开信息来确定常数 $Q$，这两种芯片的峰值计算能力（bfloat16）分别为197 teraflops和918 teraflops。假设最大FLOPs使用率为50%，这些芯片的实际计算能力分别约为100和408 teraflops。在计算理想化计算时间时，我们设置 $Q = 300$ teraflops，介于两者之间。

网络原型。对于带宽和延迟，我们考虑三种网络原型：
1. 高带宽网络，带宽 $W_{high} = 400$ Gbit/s，延迟 $ε_{high} = 10^{-4}$ 秒。
2. 中带宽网络，带宽 $W_{med} = 100$ Gbit/s，延迟 $ε_{med} = 10^{-3}$ 秒。
3. 低带宽网络，带宽 $W_{low} = 10$ Gbit/s，延迟 $ε_{low} = 10^{-2}$ 秒。

网络假设。我们强调，这些并非基于任何实际系统，而仅仅是作为指导性的网络原型设计的。对于理想化的通信时间，我们总是对数据中心内部网络使用高带宽网络，对跨数据中心网络使用三种中的一种。

B 额外实验结果

补充图表说明。在本节中，我们给出了扩展第4节中内容的额外实验结果。在图14、15、16和17中，我们展示了评估损失和各种下游零样本任务的评估准确率，作为算法、模型大小和全局批大小的函数。结果一致表明，数据并行的评估性能随着批大小的增加而迅速下降。相比之下，DiLoCo的性能下降得更慢，甚至有所改善。我们注意到Arc-Easy作为一个评估任务相当嘈杂。

Figure 14: 数据并行和DiLoCo的评估损失随全局批大小（以token计）变化的函数。

Figure 15: 数据并行和DiLoCo在HellaSwag上的零样本评估准确率随全局批大小（以token计）变化的函数。

Figure 16: 数据并行和DiLoCo在Piqa上的零样本评估准确率随全局批大小（以token计）变化的函数。

Figure 17: 数据并行和DiLoCo在Arc-Easy上的零样本评估准确率随全局批大小（以token计）变化的函数。

DiLoCo M=1 对比。在图18和19中，我们比较了数据并行和DiLoCo（M=1）在评估损失和HellaSwag、Piqa和Arc-Easy上的零样本评估准确率方面的表现。如上所述，我们注意到DiLoCo（M=1）对更大批大小具有更好的容忍度。

Figure 18: 数据并行和DiLoCo (M=1) 在不同模型和全局批大小（以token计）下的评估损失和零样本准确率。在所有设置中，DiLoCo (M=1) 都比数据并行表现更好，且差距随批大小增加而增大。

Figure 19: 数据并行和DiLoCo (M=1) 在不同模型和全局批大小（以token计）下，在Piqa和Arc-Easy上的零样本评估准确率。在几乎所有设置中，DiLoCo (M=1) 都比数据并行表现更好，且差距通常随批大小增加而增大。