Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zhenheng Tang, Zichen Tang, Junlin Huang, Xinglin Pan, Rudan Yan, Yuxin Wang, Amelie Chi Zhou, Shaohuai Shi, Xiaowen Chu, Bo Li

TL;DR

DreamDDP addresses the bottleneck of geo-distributed LLM training by replacing full synchronization in Local SGD with layer-wise partial synchronization across $L$ layers and a synchronization period $H$, enabling overlap between communication and computation without extra memory. It pairs a profiling-guided time model with a DFS-based scheduler that exploits three properties (optimal hiding, delayed CO assignment, and at-least-one assignment) to prune the search space to $\mathcal{O}(2^{\min(L-H,H)})$, while also filling idle bandwidth (bubble time) with additional communications. The approach preserves the convergence rate of S-SGD and delivers substantial improvements in iteration time and wall-clock convergence, achieving up to $3.91\times$ speedups over ASC-WFBP and competitive convergence across ResNet-18/50, GPT-2, and Llama-2 on 32 GPUs. These results indicate DreamDDP can enable faster, privacy-preserving geo-distributed training of large models by efficiently hiding communication costs.

Abstract

The growth of large language models (LLMs) increases challenges of accelerating distributed training across multiple GPUs in different data centers. Moreover, concerns about data privacy and data exhaustion have heightened interest in geo-distributed data centers. Communication in geo-distributed data parallel training (DDP) with stochastic gradient descent (S-SGD) is the main bottleneck in low-bandwidth environments. Local SGD mitigates communication overhead by reducing synchronization frequency, and recent studies have successfully applied it to geo-distributedly pre-train LLMs. However, we identify that its model synchronization mechanism prevents overlapping communication and computation, which makes the system lose opportunities to overlap communication and computation. To overcome this limitation, we expand the design space of local SGD by layer-wisely decoupling model synchronization. In each iteration, only some layers are synchronized instead of the entire model after a specific number of iterations. Leveraging this methodology, we introduce DreamDDP, a training framework to accelerate low-bandwidth distributed training with three key innovations: (1) partial local SGD with theoretical assurances of convergence rates comparable to S-SGD; (2) overlapping parameter synchronization with computation without extra GPU memory occupation; (3) identifying and exploiting three properties to schedule the communication and computation to reduce the training time based on fine-grained profiling of layer-wise communication and computation time. Empirical evaluations conducted on 32 GPUs using prominent deep learning models, including ResNet-18, ResNet-50, GPT-2, and Llama-2, demonstrate that DreamDDP enhances the convergence properties of Local SGD (and Adam) and achieves speedups ranging from $1.49\times$ to $3.91\times$ over leading baseline methods.

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

TL;DR

DreamDDP addresses the bottleneck of geo-distributed LLM training by replacing full synchronization in Local SGD with layer-wise partial synchronization across layers and a synchronization period , enabling overlap between communication and computation without extra memory. It pairs a profiling-guided time model with a DFS-based scheduler that exploits three properties (optimal hiding, delayed CO assignment, and at-least-one assignment) to prune the search space to , while also filling idle bandwidth (bubble time) with additional communications. The approach preserves the convergence rate of S-SGD and delivers substantial improvements in iteration time and wall-clock convergence, achieving up to speedups over ASC-WFBP and competitive convergence across ResNet-18/50, GPT-2, and Llama-2 on 32 GPUs. These results indicate DreamDDP can enable faster, privacy-preserving geo-distributed training of large models by efficiently hiding communication costs.

Abstract

The growth of large language models (LLMs) increases challenges of accelerating distributed training across multiple GPUs in different data centers. Moreover, concerns about data privacy and data exhaustion have heightened interest in geo-distributed data centers. Communication in geo-distributed data parallel training (DDP) with stochastic gradient descent (S-SGD) is the main bottleneck in low-bandwidth environments. Local SGD mitigates communication overhead by reducing synchronization frequency, and recent studies have successfully applied it to geo-distributedly pre-train LLMs. However, we identify that its model synchronization mechanism prevents overlapping communication and computation, which makes the system lose opportunities to overlap communication and computation. To overcome this limitation, we expand the design space of local SGD by layer-wisely decoupling model synchronization. In each iteration, only some layers are synchronized instead of the entire model after a specific number of iterations. Leveraging this methodology, we introduce DreamDDP, a training framework to accelerate low-bandwidth distributed training with three key innovations: (1) partial local SGD with theoretical assurances of convergence rates comparable to S-SGD; (2) overlapping parameter synchronization with computation without extra GPU memory occupation; (3) identifying and exploiting three properties to schedule the communication and computation to reduce the training time based on fine-grained profiling of layer-wise communication and computation time. Empirical evaluations conducted on 32 GPUs using prominent deep learning models, including ResNet-18, ResNet-50, GPT-2, and Llama-2, demonstrate that DreamDDP enhances the convergence properties of Local SGD (and Adam) and achieves speedups ranging from to over leading baseline methods.

Paper Structure

This paper contains 21 sections, 5 theorems, 24 equations, 16 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Data-Parallel Distributed Training
  3. Mini-batch SGD
  4. Distributed Synchronous SGD
  5. Local SGD
  6. Challenges of Improving LSGD
  7. DreamDDP Design
  8. Partial Synchronization: Soft Model Synchronization Point
  9. Overlapping Communication and Computation
  10. Scheduling and Optimization
  11. Filling the Bubble Time
  12. Experimental Studies
  13. Experimental Settings
  14. Wall-clock Iteration Time
  15. Convergence w.r.t. Epoch
  16. ...and 6 more sections

Key Result

Theorem 3.1

Let $f$ be $\beta$-smooth and $\mu$-strongly convex, $\mathbb{E}_{\xi_r^k}||\nabla f(w_r^k;\xi_r^k)-\nabla f(w_r^k)||^2\leq \sigma^2$, $\mathbb{E}_{\xi_r^k}||\nabla f(w_r^k;\xi_r^k)||^2\leq G^2$, for $r=0,1,...,R-1$, where $\{w_r^k\}_{r=0}^R$ for $k\in[K]$ are generated according to (eq:partial-LSGD where $\hat{w}_R=\frac{1}{KS_R}\sum_{k=1}^K\sum_{r=0}^{R-1}p_rw_r^k$, for $p_r=(a+r)^2$ and $S_R=\s

Figures (16)

  • Figure 1: Comparison of communication and computation times under varying bandwidths for GPT-2 and Llama-2.
  • Figure 2: Dissecting training time of GPT-2 and LLaMA-2 with Local SGD and S-SGD under different bandwidth.
  • Figure 3: Timelines of different distributed training algorithms.
  • Figure 4: An overview of how DreamDDP schedules and launches the overlapping of backward propagation and parameter synchronization. In different iterations, DreamDDP synchronizes model parameters according to the scheduled iteration-layer mappings.
  • Figure 5: Training ResNet-18 on 32 workers.
  • ...and 11 more figures

Theorems & Definitions (9)

  • Definition 3.1
  • Example 3.1
  • Theorem 3.1
  • Corollary 3.1
  • Definition 3.2
  • Definition 3.3
  • Lemma 1
  • Lemma 2
  • Lemma 3