DreamDDP: Accelerating Data Parallel Distributed LLM Training with Layer-wise Scheduled Partial Synchronization

作者/机构: Zhenheng Tang†1, Zichen Tang†2, Junlin Huang2, Xinglin Pan2, Rudan Yan2, Yuxin Wang3, Amelie Chi Zhou3, Shaohuai Shi4, Xiaowen Chu∗2,1, and Bo Li∗1
1香港科技大学 2香港科技大学（广州） 3香港浸会大学 4哈尔滨工业大学（深圳）

A1 主要贡献

核心问题与研究目标

大规模语言模型（LLM）的增长给跨数据中心的多GPU分布式训练带来了加速挑战，同时数据隐私和数据耗尽问题也推动了对地理分布式数据中心的需求。在低带宽环境下，采用随机梯度下降（S-SGD）的地理分布式数据并行训练（DDP）的主要瓶颈是通信。本地SGD（Local SGD）通过降低同步频率来缓解通信开销，并已成功应用于LLM的地理分布式预训练。然而，本文发现其模型同步机制阻碍了通信与计算的重叠，从而失去了系统优化的机会。本文的核心目标是克服本地SGD的这一局限性，设计一种新的训练框架，能够在保证收敛性的前提下，通过重叠通信与计算来加速低带宽环境下的分布式训练，并且不增加额外的GPU内存开销。

创新与贡献

为了解决上述问题，本文提出了DreamDDP，一个旨在加速低带宽分布式训练的框架，其核心创新点如下：

1. 部分本地SGD（Partial Local SGD）：提出了一种新的同步范式，即在每次迭代中只同步模型的部分层，而不是在H次迭代后同步整个模型。这种按层解耦模型同步的方法为在本地SGD中重叠通信与计算创造了机会。本文从理论上证明了该方法（PLSGD）具有与标准S-SGD相当的收敛速度保证。

2. 无额外内存占用的通信计算重叠：DreamDDP通过在每层的反向传播（BP）之后立即启动该层参数的同步，实现了原地（in-place）的部分参数同步。这种设计巧妙地将参数同步与后续层的BP计算重叠起来，而无需占用额外的GPU内存，解决了LLM训练中常见的内存墙问题。

3. 基于分析和优化的调度策略：

   • 细粒度分析：DreamDDP首先通过分析器（profiler）获取模型每一层通信和计算的精确时间。
   • 优化问题建模：基于分析数据，将总训练时间（包括前向、反向和通信）建模为一个关于层划分和同步周期的优化问题。
   • 高效求解算法：为解决这个复杂的组合优化问题，本文识别并利用了三个关键属性（最优隐藏、延迟CO分配、至少一次分配），设计了一个带有显著剪枝的深度优先搜索（DFS）算法，从而高效地找到近似最优的调度方案。

4. 通信“气泡”时间填充：在优化调度后，DreamDDP会利用计算过程中可能出现的通信链路空闲时间（“气泡”），插入额外的模型参数同步消息。只要这些额外的通信能被计算完全隐藏，就不会增加总训练时间，反而能提高部分层的同步频率，从而加速模型收敛。

实证成果

本文在包含ResNet-18、ResNet-50、GPT-2和Llama-2等主流模型的32个GPU上进行了实验评估。结果表明，DreamDDP不仅改善了Local SGD（及Adam）的收敛特性，而且相较于领先的基线方法，实现了1.49倍至3.91倍的加速。

A3 数据并行分布式训练

2.1 小批量SGD

基本定义。给定一个由$w \in R^d$参数化的深度学习模型，其中$d$是模型大小。单个工作节点的训练目标是最小化目标函数$F(w) \triangleq E_{x\sim D} f (w; x)$，其中数据$x$从数据集$D$中采样。本文遵循【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review】的设定，假设目标函数是$\beta$-平滑和$\mu$-强凸的。在第$r$次迭代中，工作节点计算一个小批量样本的梯度$\nabla f (w_r; \xi_r)$，其中$\xi_r$代表当前迭代$r$的数据索引所带来的采样噪声。然后模型按$w_{r+1} = w_r - \eta_r\nabla f (w_r; \xi_r)$【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review】进行更新，其中$\eta_r$是学习率。

2.2 分布式同步SGD

S-SGD流程。我们考虑$K$个不同的工作节点使用S-SGD【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review；索引编号45，Exploiting simultaneous communications to accelerate data parallel distributed deep learning，2021，IEEE INFOCOM】进行分布式训练。在每次迭代中，工作节点从原始数据集$D$中采样不同的数据样本。然后，它们执行前向传播（FP）和反向传播（BP）来计算梯度$\nabla f (w_r^k ; \xi_r^k )$。接下来，工作节点通信并平均这些梯度以获得$\frac{1}{K} \sum_{k=1}^{K} \nabla f (w_r^k ; \xi_r^k )$。模型更新如下：

通信瓶颈。通信和平均梯度$\nabla f (w_r^k ; \xi_r^k )$需要额外的时间，这限制了分布式训练的可扩展性，尤其是在低带宽环境中【索引编号8，Consumer gpus vs datacenter gpus for cv；索引编号63，CocktailSGD: Fine-tuning foundation models over 500Mbps networks，2023，ICML；索引编号69，Galaxy: A resource-efficient collaborative edge ai system for in-situ transformer inference，2024，INFOCOM；索引编号70，Asteroid: Resource-efficient hybrid pipeline parallelism for collaborative dnn training on heterogeneous edge devices，2024，ACM MobiCom；索引编号71，Decentralized training of foundation models in heterogeneous environments，2022，NeurIPS】。带有动量的S-SGD【索引编号25，An improved analysis of stochastic gradient descent with momentum，2020，NeurIPS】和分布式Adam【索引编号22，Adam: A method for stochastic optimization，2015，ICLR】也遵循这些过程，但需要维护一些额外的变量。对于动量SGD【索引编号25，An improved analysis of stochastic gradient descent with momentum，2020，NeurIPS】和Adam【索引编号22，Adam: A method for stochastic optimization，2015，ICLR】，动量项和预处理器会使用在工作节点上平均后的梯度进行更新。

2.3 本地SGD

LSGD流程。在LSGD中，每个第$k$个工作节点计算本地梯度$\nabla f (w_r^k ; \xi_r^k )$并直接用它来更新本地模型。经过$H$次训练迭代后，工作节点通信并平均它们的模型参数，即$\frac{1}{K} \sum_{k=1}^{K} w_r^k$。这个过程可以公式化如下：

与S-SGD的区别。与S-SGD不同，使用LSGD的工作节点在本地训练期间拥有不同的模型参数，但在$H$次迭代后会进行同步。

分布式训练中的通信开销。图1显示，在深度学习模型训练和LLM训练中，与计算相比，通信可能占据大量时间。当带宽较低时，训练普通深度学习模型所需的通信时间比计算时间更长【索引编号13，Enabling communication-efficient federated learning via distributed compressed sensing，2023，INFOCOM；索引编号41，Swarm parallelism: training large models can be surprisingly communication-efficient，2023，ICML；索引编号42，Towards crowdsourced training of large neural networks using decentralized mixtureof-experts，2020，NeurIPS；索引编号54，Fusionai: Decentralized training and deploying llms with massive consumer-level gpus，2023，IJCAI Symposium on LLMs；索引编号63，CocktailSGD: Fine-tuning foundation models over 500Mbps networks，2023，ICML；索引编号71，Decentralized training of foundation models in heterogeneous environments，2022，NeurIPS】。对于参数巨大的LLM，即使在通信带宽较高的环境中，通信也占用了很大的比例【索引编号4，Efficient large scale language modeling with mixtures of experts，2021，CoRR；索引编号36，Language models are unsupervised multitask learners，2019，OpenAI blog；索引编号41，Swarm parallelism: training large models can be surprisingly communication-efficient，2023，ICML】。

时间复杂度。我们构建并比较S-SGD和LSGD的时间复杂度。考虑到前向和后向时间分别为$t_{FP}$和$t_{BP}$，整个模型的通信时间为$t_{COMM}$，S-SGD和LSGD所需的总训练时间为：

$$T_{SSGD}=R\times(t_{FP}+t_{BP}+t_{COMM}),$$ $$T_{LSGD}=R \times (t_{FP}+t_{BP})+R/H \times t_{COMM}$$

时间节省分析。LSGD节省的时间比例可以计算为$(T_{SSGD} - T_{LSGD})/T_{SSGD} = (1 - 1/H) \times t_{COMM}/(t_{FP} + t_{BP} + t_{COMM})$。更大的$H$意味着更多的时间减少。然而，大的$H$可能导致收敛变慢和最终模型性能变差。LSGD的加速主要取决于$(t_{FP} + t_{BP})/t_{COMM}$的比例。

LSGD的优势。LSGD的通信成本为$O(KRS(w)/H)$，其中$S(w)$表示模型大小，$R$是总训练迭代次数。与通信成本为$O(KRS(w))$的普通S-SGD相比，LSGD将其减少了近$H$倍。并且这种减少不受工作节点数量的影响【索引编号47，A distributed synchronous sgd algorithm with global top-k sparsification for low bandwidth networks，2019，IEEE 39th International Conference on Distributed Computing Systems (ICDCS)】。与梯度压缩不同，LSGD不需要额外的计算时间进行压缩【索引编号1，On the utility of gradient compression in distributed training systems，2022，Proceedings of Machine Learning and Systems】。此外，考虑到分布式训练中的异构GPU【索引编号14，Hybrid local sgd for federated learning with heterogeneous communications，2022，ICLR；索引编号20，A unified architecture for accelerating distributed {DNN} training in heterogeneous {GPU/CPU} clusters，2020，OSDI】，LSGD非常适合通过为不同工作节点分配不同的本地训练迭代次数（即不同的$H_k$）来实现负载均衡。图2显示了LSGD相比S-SGD减少的通信时间。通过设置本地训练迭代次数为$H = 5$，通信时间节省了5倍。因为当通信带宽较低时，通信时间主导了整个迭代时间。因此，LSGD也几乎将总训练时间减少了5倍。

2.4 改进LSGD的挑战

挑战1：如何移除硬同步点？ 如图3(c)所示，当前LSGD的设计将BP和通信安排在不同的非重叠阶段，其中更新操作是一个硬同步点。具体来说，更新必须在通过BP获得梯度后进行，而通信必须在更新后进行。因此，BP无法与通信过程重叠，GPU或通信链路在任何给定时间戳都将完全空闲。这丧失了大量的性能优化机会【索引编号24，Pytorch distributed: experiences on accelerating data parallel training，2020，Proc. VLDB Endow.；索引编号61，Adaptive communication strategies to achieve the best error-runtime trade-off in localupdate sgd，2019，MLsys】，并成为LSGD的一个严重的内在限制。

模型散度的影响。在LSGD中，本地训练产生的模型散度会影响收敛性【索引编号26，Communication-efficient learning of deep networks from decentralized data，2017，Artificial Intelligence and Statistics；索引编号28，Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally!，2022，ICML；索引编号48，Local SGD converges fast and communicates little，2019，ICLR；索引编号61，Adaptive communication strategies to achieve the best error-runtime trade-off in localupdate sgd，2019，MLsys】。形式上，根据公式2，模型散度定义为$\Gamma_r = \frac{1}{K} \sum_{k=1}^{K} ||\bar{w}_r - w_r^k ||^2$，其中$\bar{w}_r = \frac{1}{K} \sum_{k=1}^{K} w_r^k$是模型参数的平均值。在S-SGD中，由于梯度的强同步，每个工作节点$w_k$的模型都是$\bar{w}_r$。因此，我们可以说$\Gamma_r$代表了LSGD对S-SGD的近似误差。过大的$\bar{w}_r$会影响梯度估计。因此，先前的观察提出LSGD应保证工作节点从一个同步点开始，以避免过大的模型散度【索引编号11，Advances and open problems in federated learning，2021；索引编号26，Communication-efficient learning of deep networks from decentralized data，2017，Artificial Intelligence and Statistics；索引编号48，Local SGD converges fast and communicates little，2019，ICLR】。更频繁的模型同步（较小的$H$）使得LSGD更接近S-SGD，从而有更好的收敛速度【索引编号48，Local SGD converges fast and communicates little，2019，ICLR；索引编号61，Adaptive communication strategies to achieve the best error-runtime trade-off in localupdate sgd，2019，MLsys；索引编号65，Minibatch vs local sgd for heterogeneous distributed learning，2020，NeurIPS】。

挑战2：如何在不增加额外GPU内存且不影响收敛的情况下重叠参数同步与计算？ 为了重叠通信和计算，我们需要考虑何时以及通信什么。与S-SGD中的重叠（如WFBP【索引编号45，Exploiting simultaneous communications to accelerate data parallel distributed deep learning，2021，IEEE INFOCOM；索引编号73，Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters，2017，USENIX ATC】）中工作节点同步梯度不同，在LSGD中，工作节点同步模型参数。考虑到LLM参数巨大，且在现有GPU中面临内存墙问题【索引编号38，Zero: Memory optimizations toward training trillion parameter models，2019，SC20: International Conference for High Performance Computing, Networking, Storage and Analysis；索引编号40，{Zero-offload}: Democratizing {billion-scale} model training，2021，2021 USENIX Annual Technical Conference (USENIX ATC 21)】，新方法不应增加GPU内存成本。因此，模型参数最好进行原地（in-place）同步。然而，原地通信和平均模型参数会直接修改模型参数。FP和BP过程都依赖于模型参数。因此，原地同步可能会影响FP和BP，导致不正确的特征和梯度计算。

挑战3：如何调度通信和计算？ 在S-SGD中，如WFBP【索引编号45，Exploiting simultaneous communications to accelerate data parallel distributed deep learning，2021，IEEE INFOCOM；索引编号73，Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters，2017，USENIX ATC】那样的重叠调度是直接的。S-SGD中的所有梯度都需要在每次训练迭代中同步。因此，每层的梯度通信可以在其BP后立即启动（不考虑合并小张量【索引编号46，Communication-efficient distributed deep learning with merged gradient sparsification on gpus，2020，IEEE INFOCOM】）。然而，在LSGD中，模型参数在$H$次迭代后同步。为了实现重叠，给定$L$个层，我们需要在不同的迭代中分配通信不同的层。

搜索空间巨大。这个问题的搜索空间极其巨大，为$H^L$。如此大的搜索空间使得寻找最优解变得困难。

A2 方法细节

DreamDDP的整体设计。DreamDDP旨在解决上述挑战，以提高LSGD的训练效率。DreamDDP的整体设计如图4所示。

首先，提出部分同步LSGD。我们提出基于部分同步的LSGD来进行DDP训练，它在不同的迭代中同步模型的不同部分，为LSGD的通信和计算重叠提供了更多机会（第3.1节）。

其次，实现原地部分参数同步。DreamDDP将神经网络划分为$L$个细粒度的层。基于部分同步，我们分析了LSGD的BP、FP和通信的依赖关系，以及如何通过原地部分参数同步实现重叠。我们发现参数同步必须在其对应层的BP之后启动，否则原地同步会影响FP和BP的计算。因此，DreamDDP主要将一层的通信与其后续层的BP重叠（第3.2节）。

第三，设计调度优化算法。DreamDDP集成了一个分析器，根据给定的模块、GPU和带宽配置，来分析每层的通信和计算时间。然后，我们建立了一个关于部分同步的$L$个层和$H$次迭代的FP、BP和通信过程的整体时间成本模型，并将其形式化为一个优化问题。我们确定了三个可以被利用的属性，包括最优隐藏、延迟CO分配、至少一次分配，来设计一个具有显著剪枝搜索空间的DFS算法来解决这个优化问题（第3.3节）。

最后，填充气泡时间以加速收敛。调度后，我们发现通信带宽可能存在一些空闲的气泡。因此，我们利用这些气泡来插入更多的通信消息（模型参数）来填充气泡时间，只要额外的通信能够被计算隐藏。这样，填充气泡时间的这些层的同步频率更高，从而加速了训练收敛（第3.4节）。

3.1 部分同步：软化模型同步点

全同步的局限性。我们将原始LSGD中所有参数的模型平均称为全同步，如图3(a)所示。LSGD的严格同步使得$H$次迭代的训练时间无法像WFBP那样用于重叠通信（图3(b)）。

部分同步的提出。为了在计算中隐藏更多LSGD的通信，我们提出了部分同步。具体来说，我们根据层索引集$\mathcal{L}$将整个模型划分为不同的段$w_l$，并均匀地将它们分配到不同的本地训练迭代中进行通信，如图3(d)所示。现在，与图3(c)相比，部分同步为在计算中重叠通信不同层提供了更多机会。

实现方式。现在，分离的层可以在计算时间内同步。通过在层的反向传播后同步它们，可以在保证参数一致性的情况下实现重叠。此外，没有额外的GPU内存占用。形式上，我们在算法1和以下定义中给出了带部分同步的LSGD过程。

定义1（部分同步）。假设模型的层索引$\mathcal{L} = \{1,..., L\}$被顺序地分割成$H$个不相交的子集$\mathcal{L}_{1:H} = \{\mathcal{L}_1, ..., \mathcal{L}_H\}$。给定同步周期$H$，带有部分同步的LSGD的更新形式为

其中$l \in \mathcal{L}$，并且$H_l = I(l)$是一个索引函数，指示第$l$层属于$\mathcal{L}_{1:H}$中的哪个子集。请注意，梯度是相对于整个模型$w_r^k$计算的。我们使用$\nabla_l f (w_r^k ; \xi_r^k )$来表示第$l$层的梯度。

示例1（等数分区）。给定$H = 2/L$，我们将$\mathcal{L}$划分为$\mathcal{L}_1 = \{1, 2\}, \mathcal{L}_2 = \{3, 4\}, ..., \mathcal{L}_H = \{L - 1, L\}$，这将$L$个层平均分配到$H$个集合中。为简单起见，我们写作$\mathcal{L}_{1:H} = \{\mathcal{L}_1, \mathcal{L}_2, ..., \mathcal{L}_H\}$。定义1和算法1显示，层集合的平均周期是不同的。与全同步类似，在$H$次迭代内，所有模型参数都被同步一次。换句话说，全同步中累积的散度$\Gamma_r$在$H$次迭代后被清除，而在部分同步中，层级的$\Gamma_r^l$在$H$次迭代后被清除。

模型散度的变化。请注意，整个模型散度$\Gamma_r$在整个训练过程中始终存在。这在某种程度上与要求完全清除模型散度的全同步【索引编号11，Advances and open problems in federated learning，2021；索引编号48，Local SGD converges fast and communicates little，2019，ICLR】相悖。然而，我们凭经验表明，带有部分同步的LSGD具有较小的模型散度$\Gamma_r$，这有利于收敛。图5显示了使用带有部分和全同步的LSGD在32个工作节点上训练ResNet18。图5(a)显示，全同步的收敛速度明显慢于部分同步。并且全同步的模型散度是周期性累积和清除的，而部分同步将散度维持在较低的水平。

减少散度放大效应。直观上，我们表明全同步中较大幅度的模型散度来自于层级散度的放大。考虑到客户端$k$上的中间特征$z_{k,l}(x) = f_{k,l}(w_{k,1:l}, x)$，其中$f_{k,l}$表示通过前$1, ..., l$层从$x$到$z$的映射，$w_{k,1:l}$之间较大的散度会引入$z_{k,l}$之间较大的散度。并且后续层的特征依赖于前面的中间特征，这意味着散度将逐层放大。因此，在全同步一个周期的后期阶段，总散度被放大到很大程度，如图5(b)所示。然而，部分同步频繁地消除了某些层的散度，从而减轻了层级散度放大效应。

理论收敛性证明。现在，我们从理论上证明带有部分同步的LSGD具有与S-SGD【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review】相同的收敛速度。

定理1。设$f$是$\beta$-平滑和$\mu$-强凸的，$E_k ||\nabla f (w_r^k ; \xi_r^k ) - \nabla f (w_r^k )||^2 \le \sigma^2$, $E_k ||\nabla f (w_r^k ; \xi_r^k )||^2 \le G^2$，对于$r = 0, 1, ..., R-1$，其中序列$\{w_r^k \}_{r=0}^R$（对于$k \in [K]$）根据(5)生成，且$H_l \le H$, $l = 1, 2, ..., L$，步长为$\eta_r = \frac{4}{\mu(a+r)}$，位移参数$a > \max\{16\kappa, H\}$，其中$\kappa = \frac{L}{\mu}$。那么

$$\begin{aligned} \begin{aligned} \mathbb{E} f\left(\hat{w}_R\right)-f^* \leq & \frac{\mu a^3}{2 S_R}\left\|w_0-w^*\right\|^2 \\ & +\frac{4 R(R+2 a)}{\mu K S_R} \sigma^2+\frac{256 R}{\mu^2 S_R} G^2 H^2 \beta \end{aligned} \end{aligned}$$

其中$\hat{w}_R = \frac{1}{S_R} \sum_{k=1}^{K} \sum_{r=0}^{R-1} p_r w_r^k$，对于$p_r = (a + r)^2$和$S_R = \sum_{r=0}^{R-1} p_r \ge \frac{1}{3} R^3$。

证明简述。我们遵循【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review】的假设，梯度采样方差有界，即对于任何工作节点$k \in [K]$和迭代$r \in [R]$，$E_k ||\nabla f (w_r^k ; \xi_r^k ) - \nabla f (w_r^k )||^2 \le \sigma^2$，这是收敛性分析中的一个通用假设。因此，多工作节点的分布式训练有助于减小梯度方差，即$E||g_r - \bar{g}_r||^2 \le \frac{\sigma^2}{K}$。并且梯度幅度有上界$E||\nabla f (w_r^k ; \xi_r^k )||^2 \le G^2$。然后，我们定义一个虚拟状态序列$\bar{w}_r = \frac{1}{K} \sum_{k=1}^{K} w_r^k$和梯度序列$g_r = \frac{1}{K} \sum_{k=1}^{K} \nabla f (w_r^k ; \xi_r^k )$，$g_r = \frac{1}{K} \sum_{k=1}^{K} \nabla f (w_r^k )$。详细证明在附录7中提供。

推论2。设$\hat{w}_R$如定理1中定义，参数$a = \max\{16\kappa,H\}$。那么

$$\mathbb{E} f\left(\hat{w}_{R}\right)-f^{*}=O\left(\frac{\kappa H^{2}}{\mu R^{2}}+\frac{\kappa^{3}+H^{3}}{\mu R^{3}}\right) G^{2}+O\left(\frac{1}{\mu K R}+\frac{\kappa+H}{\mu K R^{2}}\right) \sigma^{2}.$$

备注。定理1和推论2表明，带有部分同步的LSGD具有与S-SGD【索引编号5，Optimization methods for large-scale machine learning，2016，SIAM Review；索引编号65，Minibatch vs local sgd for heterogeneous distributed learning，2020，NeurIPS】相同的收敛速度，即$O(1/R)$。较小的$H$可以帮助加速其他包含$O(1/R^2)$或$O(1/R^3)$的项的收敛，这些项在关于训练迭代次数$R$的收敛界中不是主导项。

墙钟时间复杂度分析。遵循第2节中的定义，我们分析并比较了单次迭代重叠的LSGD和部分同步的时间复杂度。根据定义1，带有部分同步的LSGD的新训练时间为

$$T_{P-LSGD}=R\times t_{FP}+\frac{R}{H}\sum_{h=1}^{H}(t_{BP}^{\mathcal{L}_{1:h-1}}+t_{BP}^{h_0}+\max(t_{BP}^{\mathcal{L}_{h:H}}-t_{BP}^{h_0},t_{COMM}^{\mathcal{L}_h})),$$

其中，$h_0$代表集合$\mathcal{L}_h$中的第一层，$t_{BP}^{\mathcal{L}_{1:h-1}} = \sum_{i=1}^{h-1} (\sum_{l\in\mathcal{L}_i} t_{l}^{BP})$是$\mathcal{L}_1, ..., \mathcal{L}_h$中各层的后向传播时间，$t_{BP}^{\sum_{i=h}^{H}(\sum_{l\in\mathcal{L}_i} t_{l}^{BP})}$是$\mathcal{L}_{h+1}, ..., \mathcal{L}_H$中各层的后向传播时间，$t_{\mathcal{L}_i}^{COMM} = \tau_{h_{|\mathcal{L}_i|-1}}^{COMM} + t_{h_{|\mathcal{L}_i|-1}}^{COMM} - \tau_{h_0}$是$\mathcal{L}_i$中各层的通信时间，其中$\tau_{h_j}^{COMM} = \max(t_{h_{0:j}}^{BP}, \tau_{h_{j-1}}^{COMM} + t_{h_{j-1}}^{COMM})$定义了第$h_j$层通信时间的时刻。表达式$t_{h_{0:j}}^{BP}$表示第$h_j$层后向传播时间的完成时刻，$\tau_{h_{j-1}}^{COMM} + t_{h_{j-1}}^{COMM}$标记了第$h_{i-1}$层通信时间的完成时刻。项$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}$来自于集合$\mathcal{L}_h$中第一层的通信必须等待后向传播完成以避免同时访问参数。

3.2 重叠通信与计算

可行性分析。通过部分同步，定理1表明基于部分同步的LSGD（PLSGD）具有与LSGD相同的收敛速度。然而，第3.1节的分析和定理1没有提供如何重叠通信和计算的细节。在本节中，我们剖析了LSGD中通信和计算之间的依赖关系。在不引入额外GPU内存占用的原则下，我们建议将层的通信与其对应层的BP重叠。

依赖关系。图6 (d)展示了LSGD的FP、BP和参数同步。在LSGD中，工作节点同步和平均模型参数$W$，而不是像S-SGD中那样同步参数上的梯度$g_W$。每层的FP和BP计算依赖于层参数$W_l$和$g_{W_l}$。因此，与FP计算进行原地位的层通信是不可行的，因为这会导致不正确的FP和BP计算。因此，尽管将通信与FP计算重叠有助于隐藏更多的通信时间，如图6 (a)所示，但考虑到我们需要原地平均参数$W$以避免额外的内存占用，这是不可行的。此外，考虑到FP通常比BP占用更少的时间，通信时间不由FP时间主导。

调度策略的重要性。给定$L$个层需要分配到$H$次迭代中进行同步，我们需要考虑如何调度重叠以最大程度地减少时间成本。图6 (b) (c)表明，直接在第一次本地迭代中通信所有层可能不是最优的。当通信时间远大于BP时间时，如图6 (c)所示的更具负载均衡的重叠策略更好。在下一节中，我们正式将一个通信周期内PLSGD的总训练时间定义为优化目标，并确定了三个可用于减少解的搜索空间的属性。

3.3 调度与优化

优化问题。通过部分同步，减少训练时间的直观方法是通过调整层索引集$\mathcal{L}_{1:H}$的划分来最小化训练时间$T_{P-LSGD}$。根据$T_{P-LSGD}$的时间成本（公式7），时间$t_{FP}$、$t_{BP}$和$t_{COMM}$是固定的，$\mathcal{L}_{1:H}$中层索引的不同划分会影响$\sum_{h=1}^{H} \max(t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}, t_{\mathcal{L}_h}^{COMM})$。如果我们能很好地调整$\mathcal{L}_{1:H}$，$t_{\mathcal{L}_h}^{COMM}$可能被完全重叠。调度问题非常复杂。为了解决这个问题，我们将每个层划分视为一个区间，这使我们能够近似一个解。注意，最后一层的通信不能被重叠，因为通信必须等待后向传播。然后，我们提出以下目标函数$P_H^L$作为DreamDDP的目标。

$$\min_{\mathcal{L}_{1:H}^{L:1}} P_{H}^{L} = \sum_{h=1}^{H} (t_{BP}^{\mathcal{L}_{1:h-1}} + t_{BP}^{h_0} + \max(t_{BP}^{\mathcal{L}_{h:H}} - t_{BP}^{h_0}, t_{COMM}^{\mathcal{L}_h})),$$

这是一个组合优化问题，$\mathcal{L}_{L:1}^{1:H}$表示层$\{L, ..., 1\}$的可能划分$\mathcal{L}_{1:H}$。这个问题的暴力解法需要的时间复杂度为$O(C_{H+L}^L) = \frac{(H+L)!}{L!H!}$。

搜索过程和子结构问题。为了有效地解决这个问题，我们确定了问题8的最优子结构。现在，我们用一个高效的搜索算法来解决这个问题，其复杂度降低为$O(2^{\min(L-H,H)})$。该问题的搜索空间通过其最优子结构被显著减小。搜索过程是迭代地决定是否将第$l$层分配给第$h$个集合$\mathcal{L}_h$。从这个角度来看，给定已分配的层$\{L, ..., l + 1\}$在不相交的集合$\{\mathcal{L}_1, ..., \mathcal{L}_{h-1}, \mathcal{L}_h\}$中，我们定义子问题如下，

$$\min_{\mathcal{L}_{h:H}^{l:1}} S_{h}^{l}=\sum_{i=h}^{H}\left(t_{B P}^{\mathcal{L}_{1: i-1}}+t_{B P}^{i_{0}}+\max (t_{B P}^{\mathcal{L}_{i: H}}-t_{B P}^{i_{0}}, t_{C O M M}^{\mathcal{L}_{i}})\right)$$

其中$i$从$h$迭代到$H$，$H$和$L$在$S_h^l$中是固定的，在$\mathcal{L}_{l:1}^{h:H}$上进行优化意味着将层$\{1, ..., l\}$划分到集合$\{\mathcal{L}_h, \mathcal{L}_{h+1}, ..., \mathcal{L}_H\}$中。注意，$l$遵循降序，因为后向过程是从输出到输入，而$h$遵循升序，因为本地训练在同步周期$r + 1, ..., t + H$中遵循时间顺序。当$h = 1$且$l = L$时，$S_1^L$成为原始问题8，即$P_H^L = S_1^L$。

定义2（分配$A^l$）。我们将把一个层$l$分配到$\mathcal{L}_h$中称为在原始集合$\mathcal{L}_{1:H}$中用$\mathcal{L}_h \cup \{l\}$替换集合$\mathcal{L}_h$，我们有$A^l = h$。

子结构递归。优化$P_H^L$可以看作是从优化$S_1^L$到$S_1^H$的迭代搜索过程。在优化$S_h^l$的每次迭代中，我们有已分配在不相交集合$\{\mathcal{L}_1, ..., \mathcal{L}_{h-1}, \mathcal{L}_h\}$中的$\{L, ..., l + 1\}$。然后，根据是否将层$l$分配给集合$\mathcal{L}_h$，每个$S_h^l$具有如下子结构：

为简单起见，我们没有展开项$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} = \sum_{i=h}^{H}(\sum_{l\in\mathcal{L}_i} t_{l}^{BP}) - t_{h_0}^{BP}$和$t_{\mathcal{L}_h}^{COMM} = \sum_{l\in\mathcal{L}_h} t_{l}^{COMM}$。第一项意味着我们将$l$分配给$\mathcal{L}_h$，因此总时间成本$S_h^{l-1}$重新估计了第$h$轮训练中集合$\mathcal{L}_h$的时间成本。注意，无论$l \in \mathcal{L}_h$与否，在调度$S_h^l$时$t_{\mathcal{L}_{h:H}}^{BP}$实际上是固定的，因为所有剩余的层$\{\mathcal{L}_h, ..., \mathcal{L}_H\}$都将在第$h$次迭代中执行BP操作，这对$S_h^l$是固定的。我们将$\min_{\mathcal{L}_{l:1}^{h:H}} S_h^l$的最优解记为$S_h^{l, \star}$。

优化目标。问题8的最优解$P_H^L$和$S_H^L$应尽可能地减少通信时间，因为从1到$H$的每次迭代$h$的总BP时间是固定的$t_{L}^{BP}$。$P_H^L$的最小值应包含通信的最小值。现在，我们给出以下定义和属性，用于分析和减少搜索空间。

定义3（通信隐藏分配）。给定已分配在不相交集合$\{\mathcal{L}_1, ..., \mathcal{L}_h\}$中的层$\{L, ..., l + 1\}$，以及优化目标$S_h^l$，如果分配（定义2）$A^l = h$满足以下条件，我们称之为通信隐藏（CH）：

$$t_{BP}^{\mathcal{L}_{h:H}} - t_{BP}^{h_0} \ge t_{COMM}^{\mathcal{L}_h \cup \{l\}}.$$

否则，我们称分配$A^l = h$为通信溢出（CO）。

属性1（最优隐藏）。考虑$A^l = h$对于$S_h^l$是CH的情况，那么$\max(t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}, t_{\mathcal{L}_h \cup \{l\}}^{COMM}) = \max(t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}, t_{\mathcal{L}_h}^{COMM}) = t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}$，我们有$S_h^{l-1, \star} \le S_{h+1}^{l, \star} + t_L^{BP}$。直观的解释如图7所示。通过计算，未调度的$H - h$次迭代$\{h + 1,..., H\}$比$l$层更容易隐藏其他$l - 1$层$\{l - 1,..., 1\}$的通信。延迟第$l$层的通信可能导致后续通信重叠失败。

剪枝策略1。因此，基于最优隐藏（属性1），如果分配$A^l=h$能被计算完全隐藏，DreamDDP就会执行该分配（算法2第20行）。这种半贪心调度显著减小了搜索空间。

算法2：DreamDDP调度算法
输入: 层集合$\mathcal{L}$，同步频率$H$，反向传播时间$\{t_l^{BP}|l \in \mathcal{L} \}$，通信时间$\{t_l^{COMM} |l \in \mathcal{L} \}$，初始分配集$\Psi = \{\mathcal{L}_h = \emptyset |h \in \{1, ..., H\}\}$，解集$\Omega$。
输出: 最终训练好的模型$w_R$。
1: SOLVESUBPROBLEM($\Psi$, $L$, $\mathcal{L}$, 1, $H$) ▷ 调度
2: $\Psi^\star = \arg\min_{\Psi\in\Omega} P_H^L$; ▷ 找到一个最优解
3: for $\mathcal{L}_h \in \Psi^\star$ do
4: $l_h^- = \min_l t_{L...l}^{COMM}$ s.t. $l \notin \mathcal{L}_h$ and Eq. 12.
5: $\mathcal{L}_h = \mathcal{L}_h \cup \{L, ..., l\}$ ▷ 增加更多通信
6: Return $\Psi^\star$.
7:
8: procedure SOLVESUBPROBLEM($\Psi$, $l$, $\mathcal{L}$, $h$, $H$)
9: if $l == 0$ then
10: Add $\Psi$ to $\Omega$. ▷ 得到一个解
11: Return
12: if $h == H$ then
13: for $i = l, ..., 1$ do
14: Replace $\mathcal{L}_h \in \Psi$ as $\mathcal{L}_h = \mathcal{L}_h \cup \{i\}$
15: Add $\Psi$ to $\Omega$. ▷ 得到一个解
16: Return
17: if $\mathcal{L}_h$ is $\emptyset$ then ▷ 至少一次分配
18: Replace $\mathcal{L}_h \in \Psi$ as $\mathcal{L}_h = \mathcal{L}_h \cup \{l\}$
19: SOLVESUBPROBLEM($\Psi$, $l - 1$, $\mathcal{L}$, $h$, $H$)
20: else if $t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} \ge t_{\mathcal{L}_h\cup\{l\}}^{COMM}$ then ▷ 最优隐藏
21: Replace $\mathcal{L}_h \in \Psi$ as $\mathcal{L}_h = \mathcal{L}_h \cup \{l\}$
22: SOLVESUBPROBLEM($\Psi$, $l - 1$, $\mathcal{L}$, $h$, $H$)
23: else if $t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} < t_{\mathcal{L}_h}^{COMM}$ then ▷ 延迟COA
24: SOLVESUBPROBLEM($\Psi$, $l$, $\mathcal{L}$, $h + 1$, $H$)
25: else ▷ 搜索
26: $\Psi^- = \text{DeepCopy}(\Psi)$
27: Replace $\mathcal{L}_h \in \Psi^-$ as $\mathcal{L}_h = \mathcal{L}_h \cup \{l\}$ ▷ 分配 $A^l = h$
28: SOLVESUBPROBLEM($\Psi^+$, $l - 1$, $\mathcal{L}$, $h$, $H$)
29: $\Psi^+ = \text{DeepCopy}(\Psi)$
30: SOLVESUBPROBLEM($\Psi^+$, $l$, $\mathcal{L}$, $h - 1$, $H$)

属性2（延迟CO分配）。考虑集合$\mathcal{L}_h$的通信无法被BP计算隐藏的情况，即$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} \le t_{\mathcal{L}_h}^{COMM}$。在这种情况下，分配$A^l = h$也是CO，因为$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} \le t_{\mathcal{L}_h}^{COMM} \le t_{\mathcal{L}_h\cup\{l\}}^{COMM}$且满足定义3。现在，$A^l=h$将导致额外的通信时间$t_l^{COMM}$，这部分完全没有被隐藏，并增加了总训练时间，如图8(a)所示。然而，如果我们将层$l$的通信延迟到迭代$h+1$进行，通信成本可能会被计算完全或部分隐藏，如图8(b)所示。

剪枝策略2。请注意，延迟CO分配（属性2）可能不是最优的。然而，这种启发式见解可以帮助我们显著减小搜索空间。对于每个$h$，如果$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} \le t_{\mathcal{L}_h}^{COMM}$，第$l$层将被延迟以备将来分配（算法2第23行）。

属性3（至少一次分配(ALO)）。如果在调度第$l$层于第$h$次迭代时，$\mathcal{L}_h$为空集，那么分配$A^l=h$是最优的。第一种情况是$t_l^{COMM} \le t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}$，满足属性1。在另一种情况$t_l^{COMM} > t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}$下，分配$A^l=h$有助于将$t_l^{COMM}$与部分计算重叠。否则，将$A^l$延迟到$h+1$仍然会导致相同的溢出通信时间，并缩小了通信后续层的重叠机会。

剪枝策略3。鉴于ALO（属性3），在决定第$l$层和第$h$次迭代时，如果$\mathcal{L}_h$为空，DreamDDP会立即分配$A^l = h$（算法2第17行）。

深度优先搜索。在所有属性1、2和3都无法用来减小搜索空间的情况下，算法必须执行深度优先搜索（DFS）来收集解决方案。

算法2的复杂度分析。我们考虑算法复杂度的最坏情况，即DFS搜索（算法2第25行）在调度过程中频繁发生。每次DFS搜索都会导致两个分支，需要保存到解集中。当$\mathcal{L}_h$为空时，第$l$层将被加入$\mathcal{L}_h$而无需搜索，这不会产生更多的解。因此，$|\Omega|$的大小小于$2^{L-H}$。考虑到分支$A^l=h$（算法2第28行），下一层$l-1$必须延迟分配到迭代$h+1$，因为$t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP} < t_{\mathcal{L}_h\cup\{l\}}^{COMM}$。因此，$|\Omega|$小于$2^H$。总之，解集$\Omega$的大小上界为$O(2^{\min(L-H,H)})$。

3.4 填充气泡时间

动机。如3.1节和定理1所示，模型散度$\Gamma_r = \frac{1}{K} \sum_{k=1}^{K} ||\bar{w}_r - w_r^k ||^2$是影响收敛速度的一个主要因素。更高的通信频率（更低的$H$）导致训练期间的模型散度$\Gamma_r$更小，从而加速收敛【索引编号11，Advances and open problems in federated learning，2021；索引编号48，Local SGD converges fast and communicates little，2019，ICLR】。与增加同步频率$H$（即更频繁地同步整个模型）正交的另一种减少模型散度的方法是为不同的层分配不同的通信频率。

发现通信空闲。在调度重叠通信后，我们发现在后期本地训练阶段（同步周期中较大的本地$r$），后期层的通信时间可能存在一些空闲，如图9所示。这意味着可以更频繁地同步这些层以减少模型散度，同时不增加通信时间。

与模型收敛特性的契合。众所周知，深度学习模型的早期层比后期层收敛得更快【索引编号6，Which layer is learning faster? a systematic exploration of layer-wise convergence rate for deep neural networks，2023，ICLR；索引编号37，Svcca: singular vector canonical correlation analysis for deep learning dynamics and interpretability，2017，NeurIPS】，包括transformer模型【索引编号15，Pipetransformer: Automated elastic pipelining for distributed training of large-scale models，2021，ICML】。这是因为早期层在训练后期应趋于稳定，这样后期层才能在更稳定的中间特征基础上进行训练【索引编号6，Which layer is learning faster? a systematic exploration of layer-wise convergence rate for deep neural networks，2023，ICLR；索引编号37，Svcca: singular vector canonical correlation analysis for deep learning dynamics and interpretability，2017，NeurIPS】。幸运的是，PLSGD中的空闲通信时间与后期层相对应，如图9所示。因此，利用空闲通信时间更频繁地同步后期层，相比于更多地同步早期层，对减少模型散度恰好更有利。

具体实现。形式上，DreamDDP考虑为通信迭代$h$通信更多的层$L^- = \{L, ..., l\}$（算法2第5行）以加速收敛，同时确保通信时间不增加，满足以下等式：

$\max(t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}, t_{\mathcal{L}_h\cup\{L,...,l\}}^{COMM}) \le \max(t_{\mathcal{L}_{h:H}}^{BP} - t_{h_0}^{BP}, t_{\mathcal{L}_h}^{COMM})$，其中$\{L, ..., l\}$与$\mathcal{L}_h$不相交。

A4 实验环境

• 模型与数据集：

  • ResNet-18 在 CIFAR-10 上训练。
  • ResNet-50 在 CIFAR-100 上训练，使用带动量的SGD。
  • GPT-2 small 【索引编号36，Language models are unsupervised multitask learners，2019，OpenAI blog】和 一个175M参数的小型Llama-2【索引编号58，Llama 2: Open foundation and finetuned chat models，2023，arXiv preprint arXiv:2307.09288】在 WikiText-2【索引编号27，Pointer sentinel mixture models，2017，ICLR】上训练，使用Adam优化器【索引编号22，Adam: A method for stochastic optimization，2015，ICLR】。

• 硬件配置：

  • 集群1 (用于ResNet)：8台机器，每台配备4块NVIDIA 2080Ti（共32块GPU），通过1GB/s以太网连接。
  • 集群2 (用于LLM)：4台机器，每台配备8块NVIDIA A6000（共32块GPU），通过20GB/s以太网连接。

• 软件与模拟环境：

  • 为了模拟地理分布式训练，通过将机器间的带宽从10MB/s变化到1000GB/s来代表不同地理位置的多样化网络条件。
  • 基线方法：与朴素S-SGD、ASC-WFBP【索引编号45，Exploiting simultaneous communications to accelerate data parallel distributed deep learning，2021，IEEE INFOCOM】（一种先进的S-SGD重叠算法）、以及带全同步的LSGD（FLSGD）【索引编号32，The limits and potentials of local sgd for distributed heterogeneous learning with intermittent communication，2024，COLT；索引编号72，Aperiodic local sgd: Beyond local sgd，2023，ICPP】进行比较。
  • 消融研究：还实验了DreamDDP的基础版本，包括带部分同步的LSGD（PLSGD-ENP），采用等数分区（示例1）。
  • 说明：对于LLM，优化器是Adam，但由于其计算和通信依赖性与SGD相似，因此在并行化模式的讨论中，为简化起见，随机Adam和本地Adam也分别称为S-SGD和LSGD。

A5 实验结果

4.2 墙钟迭代时间

如表1所示，DreamDDP在所有模型和工作节点配置中都持续取得了最低的迭代时间，显示了其在消除额外通信成本方面的卓越效率。具体来说，与先进的DDP训练方法ASC-WFBP相比，DreamDDP实现了1.73倍至5.22倍的加速。同样，与FLSGD相比，DreamDDP实现了1.16倍至1.5倍的加速。这种性能提升在低带宽环境中（地理分布式场景中常见）可能更为显著。

表1：1000次运行迭代的平均迭代墙钟时间（秒）。S1和S2分别代表DreamDDP相对于ASC-WFBP和FLSGD的加速比。

4.3 相对于轮次的收敛性

尽管FLSGD的迭代时间比ASC-WFBP快，但其通信量减少了$H$倍，导致收敛性变差。我们在图10和图11中比较了不同的同步频率$H$。结果显示，对于FLSGD和DreamDDP，更大的$H$都会导致相对于轮次的收敛变慢。然而，DreamDDP中的补充通信有助于加速收敛。图11显示，增加工作节点数量并未影响DreamDDP的性能提升，证明了其良好的可扩展性。

4.4 墙钟收敛性

图12展示了不同算法相对于墙钟时间的收敛情况，表2则列出了达到目标性能所需的墙钟时间。两者都表明，所有LSGD方法的墙钟收敛速度都显著快于SGD。周期性模型同步方案有效减少了S-SGD和ASC-WFBP中频繁同步带来的高通信成本。具体而言，DreamDDP相比ASC-WFBP最高实现了3.91倍的加速。

此外，由于先进的调度算法消除了额外的通信开销，DreamDDP相比FLSGD最高实现了1.56倍的加速。再者，与采用类似通信-计算重叠策略的PLSGD相比，DreamDDP也展现了更快的收敛速度（如图12和表2所示）。这一改进证明了DreamDDP的补充通信机制的有效性，该机制在不产生额外训练时间的情况下，最大限度地利用了重叠机会。

不同同步频率的影响。我们在图13和图14中，将不同同步频率H的DreamDDP与H=10的FLSGD及S-SGD进行了比较。在8个和32个工作节点上的结果均显示，更大的H会使DreamDDP的按轮次收敛变慢，这与直觉相符。然而，DreamDDP的补充通信有助于加速收敛，使其达到了与更高频率的FLSGD相当的收敛速度。图14表明，增加工作节点数量不影响DreamDDP的性能提升，展示了其良好的可扩展性。

表2：不同训练方法达到目标模型性能的墙钟时间。对于深度学习模型，使用测试准确率；对于LLM，使用训练损失。S1和S2代表DreamDDP相对于ASC-WFBP和FLSGD的加速比。粗体表示最佳结果。

4.5 调度性能

不同层数的影响。为了探究DreamDDP能否最大程度地重叠通信时间，我们将其未重叠时间与暴力搜索的结果进行了比较。由于暴力搜索的时间复杂度过高（$\frac{(H+L)!}{L!H!}$），我们展示了在$H=5$的情况下，调度最多30个层的结果。图15比较了DreamDDP和暴力搜索调度后未与计算重叠的额外通信时间。图15(a)显示，随着层数的变化，在一个训练周期（5次迭代）内，DreamDDP在大多数情况下表现与暴力搜索的最优解相当，仅在少数场景中额外通信时间稍长。

不同带宽的影响。图15(b)显示，在不同带宽下，DreamDDP与暴力搜索得到的最优结果性能也相近。这些结果表明，DreamDDP在显著降低搜索复杂度的同时，有效地识别了近乎最优的调度策略。

4.6 分析与搜索开销

我们比较了包括DreamDDP和暴力搜索在内的调度算法在调度不同模型时的理论和实际分析运行时间。图16(a)展示了用$H=5$调度四个模型的理论时间复杂度。由于这些模型层数过多，用暴力搜索进行调度的时间复杂度是不可行的。因此，我们仅为这四个模型选择30个层来比较它们的实际运行时间。图16(b)显示，与暴力搜索方法相比，DreamDDP的调度方法将搜索复杂度降低了几个数量级。DreamDDP减小了搜索空间，有效避免了探索所有可能的层划分。

A7 补充细节

5 相关工作

重叠通信与计算。在采用S-SGD的分布式训练中，每层的梯度可以在其BP后立即进行通信，并与后续层的BP重叠，从而节省总训练时间【索引编号73，Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters，2017，USENIX ATC】。将通信与前向计算重叠【索引编号33，A generic communication scheduler for distributed dnn training acceleration，2019，SOSP】有助于进一步减少总训练时间。一些研究观察到，合并小张量进行通信【索引编号46，Communication-efficient distributed deep learning with merged gradient sparsification on gpus，2020，IEEE INFOCOM】可以提高通信效率。ASC-WFBP【索引编号45，Exploiting simultaneous communications to accelerate data parallel distributed deep learning，2021，IEEE INFOCOM】通过同时通信和调度进一步改进了重叠效果。与这些工作不同，我们考虑在LSGD中重叠通信和计算。

LSGD。LSGD【索引编号43，Communication acceleration of local gradient methods via an accelerated primal-dual algorithm with an inexact prox，2022，NeurIPS；索引编号48，Local SGD converges fast and communicates little，2019，ICLR；索引编号65，Minibatch vs local sgd for heterogeneous distributed learning，2020，NeurIPS；索引编号72，Aperiodic local sgd: Beyond local sgd，2023，ICPP】通过周期性地平均模型权重而非梯度来减少通信时间。同步频率对收敛和通信时间有很大影响。因此，一些工作通过改变同步周期来权衡收敛和通信【索引编号48，Local SGD converges fast and communicates little，2019，ICLR；索引编号61，Adaptive communication strategies to achieve the best error-runtime trade-off in localupdate sgd，2019，MLsys；索引编号72，Aperiodic local sgd: Beyond local sgd，2023，ICPP】。强同步严重限制了LSGD的潜在改进，使其无法进行通信重叠。为此，我们提出了部分同步，以解耦整个模型并在不同的训练迭代中同步各层。

数据分布的影响。不同工作节点中的数据分布也影响LSGD的收敛速度【索引编号65，Minibatch vs local sgd for heterogeneous distributed learning，2020，NeurIPS】。在这项工作中，我们考虑的是集中式分布式训练中的IID数据，而非联邦学习。尽管如此，我们的方法和收敛性分析可以很容易地扩展到联邦学习场景【索引编号11，Advances and open problems in federated learning，2021；索引编号56，Virtual homogeneity learning: Defending against data heterogeneity in federated learning，2022，ICML】。

降低通信成本。为了减少通信时间，一些工作提出在S-SGD中压缩通信的梯度【索引编号12，Efficient sparse collective communication and its application to accelerate distributed deep learning，2021，ACM SIGCOMM；索引编号52，Communication-efficient distributed deep learning: A comprehensive survey，2020，arXiv preprint arXiv:2003.06307】，或在LSGD中压缩模型参数【索引编号13，Enabling communication-efficient federated learning via distributed compressed sensing，2023，INFOCOM；索引编号53，Gossipfl: A decentralized federated learning framework with sparsified and adaptive communication，2022，IEEE Transactions on Parallel and Distributed Systems；索引编号62，Online distributed optimization with efficient communication via temporal similarity，2023，INFOCOM】。然而，压缩可能需要大量的潜在计算时间并减慢收敛速度，这阻碍了其部署【索引编号1，On the utility of gradient compression in distributed training systems，2022，Proceedings of Machine Learning and Systems】。【索引编号13，Enabling communication-efficient federated learning via distributed compressed sensing，2023，INFOCOM；索引编号62，Online distributed optimization with efficient communication via temporal similarity，2023，INFOCOM】提出了压缩联邦学习【索引编号53，Gossipfl: A decentralized federated learning framework with sparsified and adaptive communication，2022，IEEE Transactions on Parallel and Distributed Systems；索引编号56，Virtual homogeneity learning: Defending against data heterogeneity in federated learning，2022，ICML】中的本地结果，其中使用了LSGD。与该方向正交，我们的方法旨在修改LSGD以提高系统效率。

6 局限性

异构环境。现实世界的地理分布式设置可能涉及异构的硬件配置和带宽，导致不同集群之间的计算和通信时间各不相同。因此，我们可能需要考虑如何设计合适的算法和调度来解决这个问题。

动态层划分。我们的方法依赖于分析出的计算和通信时间来优化调度。然而，在带宽可能频繁变化的地理分布式训练场景中，静态分析的时间可能不再准确反映当前的网络状况，这可能会降低调度策略的有效性。

A6 结论

在地理分布式训练中，传统的S-SGD因聚合梯度的通信成本而受到影响，尤其是在低带宽环境和LLM中。LSGD降低了通信成本，但它需要严格的同步，并且收敛速度较慢。在这项工作中，我们首先提出了部分同步来放宽LSGD的严格同步。然后，我们设计了DreamDDP来调度通信与反向传播计算重叠，以减少LSGD的通信开销。DreamDDP还通过最大限度地利用重叠而无需额外训练时间来补充通信。实验结果表明，在包括ResNets、GPT和Llama-2在内的四种流行类型的深度学习模型上，DreamDDP在墙钟迭代时间和收敛速度方面都比ASC-WFBP和LSGD（及Adam）更快。

A8 附录

本附录提供了定理1的详细证明。

引理3. 平均权重的收敛性。设$\{w_r\}_{r\ge0}$和$\{\bar{w}_r\}_{r\ge0}$（对于$k \in [K]$）如方程(5)中所定义，并设$f$为$\beta$-平滑和$\mu$-强凸，且$\eta_r \le \frac{1}{4\beta}$。那么

$$\begin{aligned} \begin{aligned} \mathbb{E}\|\bar{w}_{r+1}-w^*\|^2 &\leq (1-\mu\eta_r)\mathbb{E}\|\bar{w}_r-w^*\|^2 + \eta_r^2\mathbb{E}\|g_r-\bar{g}_r\|^2 \\ &- \frac{1}{2}\eta_r\mathbb{E}(f(\bar{w}_r)-f^*) + 2\eta_r\frac{\beta}{K}\sum_{k=1}^K\mathbb{E}\|\bar{w}_r-w_r^k\|^2 \end{aligned} \end{aligned}$$

引理3的证明。使用更新公式5，我们有

根据$f$的L-平滑性和$\mu$-强凸性，我们可以推导出第一项的上界为

$$\begin{aligned} \begin{aligned} &\|\bar{w}_r - w^* - \eta_r \bar{g}_r\|^2 \le (1 - \mu\eta_r)\|\bar{w}_r - w^*\|^2 \\ &- \frac{1}{2}\eta_r(f(\bar{w}_r) - f^*) + \frac{2\eta_r L}{M} \sum_{k=1}^K \|\bar{w}_r - w_r^k\|^2. \end{aligned} \end{aligned}$$

将(14)代回(13)，通过取期望我们得到

$$\begin{aligned} \begin{aligned} \mathbb{E}\|\bar{w}_{r+1}-w^*\|^2 \leq & (1-\mu \eta_r) \mathbb{E}\|\bar{w}_r-w^*\|^2 + \eta_r^2 \mathbb{E}\|g_r-\bar{g}_r\|^2 \\ & -\frac{1}{2} \eta_r \mathbb{E}(f(\bar{w}_r)-f^*) + \frac{2L\eta_r}{M} \sum_{k=1}^K \mathbb{E}\|\bar{w}_r-w_r^k\|^2 \end{aligned} \end{aligned}$$

引理4. 有界模型散度。给定有界同步频率$H_l \le H$（对于$l \in [L]$）和递减的正学习率序列$\{\eta_r\}_{r\ge0}$满足$\eta_r \le 2\eta_{r+H}$对于所有$r \ge 0$，则我们有

$$\frac{1}{K} \sum_{k=1}^{K} \mathbb{E}\|\bar{w}_{r}-w_{r}^{k}\|^{2} \leq 4 H^{2} \eta_{r}^{2} G^{2}.$$

引理4的证明。由于$w_k$是$w_{k,l}^r$关于$l$的拼接，我们有$||\bar{w}_r - w_r^k ||^2 = \sum_{l=1}^{L} ||\bar{w}_l^r - w_{k,l}^r ||^2$，$\sum_{l=1}^{L} ||\nabla_l f (w_h^k; \xi_h^k)||^2 = ||\nabla f (w_h^k; \xi_h^k)||^2$。使用$E||X - EX ||^2 = E||X ||^2 - ||EX ||^2$和$||\sum_{i=1}^{H_l} a_i||^2 \le H_l \sum_{i=1}^{H_l} ||a_i||^2$，我们有

引理5. 鞅收敛。设$\{a_r\}_{r\ge0}, a_r \ge 0$, $\{e_r\}_{r\ge0}, e_r \ge 0$是满足以下条件的序列

$$a_{r+1} \leq (1-\mu \eta_r)a_r - \eta_r e_r A + \eta_r^2 B + \eta_r^3 C,$$

对于$\eta_r = \frac{4}{\mu(a+r)}$和常数$A > 0, B,C \ge 0, \mu > 0, a > 1$。那么

$$\frac{A}{S_R} \sum_{r=1}^{R-1} p_r e_r \leq \frac{\mu a^3}{4 S_R} a_0 + \frac{2 R(R+2 a)}{\mu S_R} B + \frac{16 R}{\mu^2 S_R} C,$$

对于$p_r = (a + r)^2$和$S_R \triangleq \sum_{r=0}^{R-1} p_r = \frac{R}{6} (2R^2 + 6aR - 3R + 6a^2 - 6a + 1) \ge \frac{1}{3} R^3$。

定理1的证明。基于引理3和引理4，我们得到

$$\begin{aligned} \begin{aligned} \mathbb{E}\|\bar{w}_{r+1}-w^*\|^2 \leq & (1-\mu \eta_r) \mathbb{E}\|\bar{w}_r-w^*\|^2+\frac{\sigma^2}{K} \eta_r^2 \\ & -\frac{1}{2} \eta_r \mathbb{E}(f(\bar{w}_r)-f^*)+8 G^2 H^2 \beta \eta_r^3. \end{aligned} \end{aligned}$$

根据引理5，取$a_r = E||\bar{w}_r - w^*||^2$, $e_r = E(f(\bar{w}_r) - f^*)$, $A = \frac{1}{2}$, $B = \frac{\sigma^2}{K}$, $C = 8G^2H^2\beta$，我们有

$$ \frac{1}{2 S_{R}} \sum_{r=1}^{R} p_{r} \mathbb{E}\left(f\left(\bar{w}_{r}\right)-f^{*}\right) \leq \frac{\mu a^{3}}{4 S_{R}} \mathbb{E}\left\|\bar{w}_{0}-w^{*}\right\|^{2}+\frac{2 R \sigma^{2}(R+2 a)}{\mu K S_{R}}+\frac{128 G^{2} H^{2} R}{\mu^{2} S_{R}} . $$

根据$f$的凸性，我们有

$$\mathbb{E} f\left(\hat{w}_R\right)-f^* \leq \frac{1}{S_R} \sum_{r=1}^R \mathbb{E}\left(f\left(\bar{w}_r\right)-f^*\right)$$

通过公式20和21，我们得到结果

$$\begin{aligned} \begin{aligned} \mathbb{E} f\left(\hat{w}_R\right)-f^* & \leq \frac{\mu a^3}{2 S_R}\left\|w_0-w^*\right\|^2 \\ & +\frac{4 R(R+2 a)}{\mu K S_R} \sigma^2+\frac{256 R}{\mu^2 S_R} G^2 H^2 \beta, \end{aligned} \end{aligned}$$

这完成了定理1的证明。