Optimal Distributed Training With Co-Adaptive Data Parallelism in Heterogeneous Environments

1 Introduction

The rapid expansion of Internet applications has led to an exponential increase in global data, with recent data generation surpassing that of the entire previous history [Yi et al., 2020]. The surge of big data stimulates the rapid expansion of deep learning models, ranging from simple neural networks to Large Language Models (LLMs) containing hundreds of billions of parameters, such as the GPT family [Brown et al., 2020] and Transformer [Wang et al., 2022], thereby placing an increasingly urgent demand on efficient training of largescale learning models.

∗Corresponding authors: Pan Li, Liqi Yan.

In order to cope with this challenge, distributed training has emerged. Its core concept is to leverage distributed nodes to decompose and distribute computational tasks for collaborative completion, effectively overcoming the limitations of single-machine processing power. Data parallelism, a key form of distributed computing, is widely used in training large-scale learning models [Li et al., 2020; Verbraeken et al., 2020]. It efficiently trains complex models by splitting training data and distributing it across computing nodes. In Computer Vision (CV), data parallelism enables efficient training for tasks such as image classification and object detection, and improves model accuracy by extracting rich visual features. In Natural Language Processing (NLP), data parallelism ensures the processing of large-scale text data, helping models capture intricate semantic and syntactic relationships, which drives continuous advancements in NLP technologies.

Despite the considerable benefits of data parallelism in enhancing computational efficiency, its practical implementation encounters numerous challenges. First, distributed computing could introduce large communication overhead. In the data distribution phase, the transmission of large volumes of training data is limited by network bandwidth and latency, while in the training process, frequent synchronization of gradients or parameters further increases the communication cost. Second, computing device heterogeneity can lead to ineffective data distribution and uneven loading due to uniform data distribution assumed by popular distributed training mechanisms like Distributed Data Parallelism (DDP) [Li et al., 2020]. Although frameworks such as DeepSpeed [Ren et al., 2021; Rasley et al., 2020] have improved memory management, they are still deficient in balanced data allocation. Third, most distributed training frameworks such as DDP and DeepSpeed are primarily designed for co-located clusters with homogeneous network conditions, and hence cannot account for geo-distributed clusters with heterogeneous network environment over the Interne.

To address these challenges, we develop a new CoAdaptive Data Parallelism (C-ADP) framework by optimizing data parallel scheduling while taking into account computing and communication heterogeneity in a geo-distributed cluster. Specifically, C-ADP enables more efficient and robust distributed training through the following designs. (1) Delayed parameter synchronization. By employing local training and delayed parameter synchronization at the conclusion of each epoch instead of each batch, we substantially reduce communication overhead. (2) Data parallel scheduling formulation. We formulate a data parallel scheduling optimization problem to minimize the training time by optimizing data distribution. (3) An efficient algorithm. We develop a very efficient algorithm to solve the optimization problem, and formally prove that the obtained solution is optimal in the asymptotic sense. Through strategic data distribution, we ensure balanced loading across computing devices with heterogeneous capabilities and significantly reduced communication overhead, thus accelerating the training process and enhancing overall efficiency.

In this paper, our main contributions are threefold:

• We propose a new data parallel based distributed training framework, named Co-Adaptive Data Parallelism (CADP), for a geo-distributed cluster with heterogeneous computing and communication environments.
• We formulate data parallel scheduling as an optimization problem and develop an efficient algorithm that is able to efficiently and optimally determine the amount of data samples to be allocated to different computing devices.
• Experiments on the ImageNet100 dataset show that, compared with DDP and DeepSpeed, our C-ADP framework achieves 21.6 times and 26.3 times improvements in FLOPS, respectively, and a reduction in trainning time of about 72% and 47%, respectively.

2 RelatedWork

Distributed Training Architectures: There have been a few distributed training frameworks such as DDP [Li et al., 2020], DeepSpeed [Rasley et al., 2020], Megatron-LM [Shoeybi et al., 2020]. Data parallelism [Zhao et al., 2023], pipeline parallelism [Narayanan et al., 2019], tensor parallelism [Shazeer et al., 2018] are key methods for parallelizing operations during the training process. In data parallelism, a global model is replicated across a cluster of computing devices, each of which is allocated a subset of data, and is collaboratively trained by them. Synchronous and asynchronous Stochastic Gradient Descent (SGD) are commonly used for convergence and scalability, leveraging parameter server or ring-allreduce [Yu et al., 2022] architectures. The parameter server architecture uses a centralized server for global parameters, while ring-allreduce employs a decentralized approach. Pipeline parallelism divides the model layers into stages and distributes them across different computing devices [Narayanan et al., 2019]. Tensor parallelism shards the model's parameters and computations along tensor dimensions and distributes the computations of a single layer, e.g., matrix multiplications across devices. Frameworks like DDP [Li et al., 2020], Horovod [Sergeev and Balso, 2018] use pure data parallelism. Those like DAPPLE [Fan et al.,

2021], PipeDream [Narayanan et al., 2019] combines data and pipeline parallelism. and Megatron-LM [Shoeybi et al., 2020], AMP [Li et al., 2022] integrate all three types of parallelism. In this paper, we focus on data parallelism since it has excellent scalability and usually moderate communication cost compared with pipeline and tensor parallelism.

Scheduling Optimization in Heterogeneous Environments: A few mechanisms [Jayaram Subramanya et al., 2023; Mo et al., 2024; Unger et al., 2022] have been developed to optimize cluster-level scheduling. In particular, [Jayaram Subramanya et al., 2023; Chen et al., 2022] optimize the scheduling of multiple jobs, e.g., by increasing data parallelism, on heterogeneous GPU clusters. [Yang et al., 2024] intelligently assigns different computational tasks in model training within heterogeneous environments based on the characteristics of the NICs to which the GPU devices are connected. [Yuan et al., 2022] minimizes communication costs in heterogeneous environments. [Liu et al., 2023] designs a reinforcement learning-based scheduling method to assign each layer to an appropriate type of computing resource, minimizing cost while ensuring throughput limits based on a provisioning method. Previous scheduling optimization mechanisms mostly only considers either computing or communication heterogeneity. In this paper, we address both computing and communication heterogeneity in data parallelism.

3 SystemModel and ProblemFormulation

We consider a geo-distributed system with a data owner and parameter server (PS) and N computing devices as shown in Figure 1. After finding the optimal distributed training schedule, with will be delated later, the PS sends the data subset Di (1 ≤ i ≤ N) to each participating computing devices i. Each device i performs local training on the received data subset Di using stochastic gradient descent (SGD), and then uploads the locally updated model parameters wt+1 i based on the current global model wt to the central server. The PS performs weighted averaging based on the local model parameters of all devices to update the global model wt+1 distributes it to all devices, i.e.,

where D is the total amount of data samples. Then, we have

where Di ∈ N0, we assume that all data samples have the same size of b bits.

In the realm of distributed training, we focus solely on data parallelism, where our primary objective is to minimize the total time consumed by the entire training process encompassing multiple epochs. This total time is composed of three crucial components: data allocation time, computation time and communication time. Some important notations are summarized in Table 1.

3.1 Data Allocation Time

Denote by B0,i the bandwidth between device 0 and device i. The communication time for device i to receive its allocated Di data samples, denoted by Tsend,i, is

3.2 Computation Time

During distributed training, the computational time of each device is related to the amount of data allocated to it, i.e., the larger the dataset, the longer the computation time. Denote the maximum parallelism of device i as Pmax(i), i.e., the maximum number of samples that a device can process concurrently. Then, we can get

which needs to be no less than the batchsize at device i, denoted by batchsize(i), i.e.,

Then, the computation time at device i during one epoch, denoted by Tcomp,i, can be obtained by

where Ti batch denotes the time required by device i to process a single batch of data, including both forward propagation and backward propagation computations.

3.3 Communication Time

Each computing device i receives the current global model parameters from the PS, updates its local model using its allocated training data, and sends the updated model parameters to the PS at the end of each epoch. We denote by Mm is the size of the model parameters. Thus, the communication cost of device i for data parallelism in one epoch, denoted by Tcomm,i can be calculated as follows:

3.4 Data Parallel Scheduling Optimization

To minimize the total time for the above data parallel training process, we aim to optimize the data allocation Di (1 ≤ i ≤ N). Consequently, the data parallel scheduling optimization problem can be formulated as:

4 An Efficient Data Parallel Scheduling

4.1 Optimality Analysis

To simplify notations, we let x = {D1,D2, ...,DN}. We rewrite the optimization problem formulated in equation (7) into the following:

where

Note that is C independent of the variable x. It dependents on communication bandwidth B and model size Mm, and hence it is a constant.

Lemma 1. x ≥ 0, we have lim p→∞ ∥x∥p = max x.

Proof. Let ˆx =max x. Then, we have

Therefore, we can find lim p→∞ ∥x∥p = lim p→∞( N∑ k=1 xpk ) 1p = max x.

In fact, when p = 100, the error between ∥x∥p and max x is already sufficiently small, as shown in Figure 2.

Moreover, we can also prove that ∥x∥p has the same subadditivity and homogeneity as max x.

Proof. According to the Minkowski inequality, we get

Therefore, we get

and

Consequently, according to Lemma 1, when p is large, the objective function in equation (8) can be rewritten into

whose feasible domain is

We then define a relaxed function

whose feasible domain is

We arrive at the following theorem, the proofs of which can be found in the appendices.

Theorem 1. The relaxed function equation (10) is a convex function on its feasible domain.

Lemma 3. For any x ∈ RN and x ≥ 0, we have ∥⌈x⌉ - x∥p < N 1p .

Proof. For any element xk ∈ x, we have |⌈xk⌉ -xk| < 1. Thus, we can obtain

Theorem 2. Given x ∈ χ′ , for any ˆx ∈ χ that satisfies

where γ ≥ N 1p , we have

According to Theorem 2, we first find the optimal solution ˜x∗ to the relaxed equation (10). Then, we can find ˆx ∈ χ, in the neighborhood of ˜x∗ , i.e., U(˜x∗ , γ), γ ≥ N 1p , such that |Φ(ˆx) - ˜Φ(˜x∗ )| ≤ Mγ. Note that Mγ is independent of the total amount of data samples D. Therefore, for the optimal solution ˆx∗ to the original equation (9), we have ˜Φ(˜x∗ ) ≤ Φ(ˆx∗ ) ≤ Φ(ˆx), and hence, Φ(ˆx) - Φ(ˆx∗ ) ≤ Mγ. Since Φ(ˆx∗ ) is positively correlated to D, we can have

i.e., we can find a solution that is optimal in the asymptotic sense when there are a sufficiently large number of training data samples.

Algorithm 1 Our Data Parallel Scheduling Algorithm

4.2 AlgorithmDesign

Based on the analysis in Section 4.1, we design an algorithm that first solves for the optimal solution ˜x∗ to equation (10), and then finds the solution to the original equation (9), i.e., ˆx∗ by rounding each element of ˜x∗ to an integer. According to Theorem 1, we are able to find a globally optimal solution to the relaxed equation (10). So we construct an augmented Lagrangian function as follows:

Where λ is the Lagrange multiplier for the reciprocal constraint, ρ > 0 is the coefficient of the quadratic penalty term, and µ > 0 is the coefficient of the logarithmic barrier used to ensure that xi ≥ 0. We detail our data parallel scheduling algorithm in Algorithm 1.

4.3 Computational Complexity Analysis

The inner loop of Algorithm 1 involves computing the gradients and diagonal matrix multiplication. Since the inner loop is executed at most ζ times, the computational complexity is O(ζN). The outer loop is executed at most ψ times. Therefore, the computational complexity of our data parallel scheduling algorithm is O(ψζN), where N is the number of computing devices in the system.

5 Experiment Results

5.1 Dataset and LearningModel

Dataset: The experiments were carried out using the ImageNet100 dataset, a subset of the ImageNet dataset. ImageNet100 consists of 100 classes, particularly 100,000 training samples and 10,000 test samples, making it a balanced and diverse dataset for evaluating model performance.

Models: To evaluate C-ADP, Wide - ResNet101 was chosen and validated on ImageNet100. Wide-ResNet101, engineered to process intricate image features, incurs a computational cost of 22.8 GFLOPs.

5.2 Experimental Environment andMetrics

Computing Devices: To account for heterogeneous environments, the hardware setup is as follows. In the Zhongwei area, rank 0 as a data owner and parameter server handles data transfer and parameter synchronization. Ranks 1 - 4 use RTX3090 GPUs with 24GB VRAM. Ranks 5 - 6 are 4 core CPUs with 8GB memory. Rank7 is a 16 core CPU with 32GB VRAM, and rank8 is an 8 core CPU with 16GB VRAM. To emulate real-world constraints, we simulate network heterogeneity among these devices by controlling the available bandwidth between devices.

Hyperparameter Configuration: For all GPUs, the batch size is set to 64 due to its high computational complexity, to ensure stable training within the GPU's memory capacity. For all CPUs, the batch size is set to 4, considering their limited computing power. All settings are designed to keep operations within the device's memory and computational limits. The maximum number of epochs is set to 50. The initial learning rate is 0.01. To improve convergence, we reduce the learning rate by half after every 10 epochs.

Baselines: To ensure comprehensive performance comparison, two widely-used and efficient frameworks are selected as baselines, as they represent state-of-the-art solutions to distributed deep learning training:

• Deepspeed: an open source library developed by Microsoft to accelerate the training of deep learning models. It provides an efficient training framework that supports distributed training, model parallelism, and data parallelism.
• DDP (Distributed Data Parallel): a widely-adopted feature in PyTorch. DDP provides robust support for distributed training by replicating models across multiple devices and synchronizing gradients.

Metrics: The following metrics are adopted to evaluate the performance of C-ADP and the baselines comprehensively:

• Accuracy Metrics: standard metrics including Accuracy
• (Acc), F1-Score (F1), and Precision to assess the quality of the trained model.
• Tsend: the data transmission time required for computing devices to receive datasets from the rank 0 data owner and parameter server.
• Ttraining: the end-to-end training time, which includes the forward propagation, backward propagation, and parameter synchronization.
• Ttotal: The total time for completing a fixed number of epochs on a device, calculated as Ttotal = Tsend+Ttraining.

5.3 Model Convergence in Heterogeneous Environments

Verification of AccuracyMetrics: We trains the computationally intensive WideResNet101 on ImageNet100. To address the frequent out-of-memory issues when using a single 3090 card, we employ a more powerful A800 card for singlecard training and compare its performance with the proposed C-ADP framework. Training accuracy, validation loss, and final metrics are shown in Figure 3 and Table 2.

Loss Convergence Validation: To validate the convergence properties of C-ADP in a heterogeneous environment, we show the training loss curves of all computing devices in Figure 4. As can be clearly observed, C-ADP achieves successful convergence across all 8 devices, with the loss function steadily decreasing and low loss values reached. Additionally, a detailed examination reveals that the scheduling method employed effectively allocates data based on the performance disparities among devices, enabling efficient handling of data processing tasks across devices with varying capabilities, without requiring the same number of iterations.

Result Analysis: Early in training, single-card training outperforms C-ADP. But as the number of training rounds grows, C-ADP takes the lead. At the 30th round, C-ADP accuracy reached 70.56%, exceeding the single-card's accuracy of 69.46%, and increases to 73.67% at the 50th round. F1-score, precision, and recall follows similar trends. The test loss is lower in the case of single-card training initially, but lower for C-ADP later, suggesting better late-stage performance. Despite slow convergence due to low-frequency parameter synchronization, C-ADP achieves convergence across all devices. It also outperforms single-card training in key metrics like accuracy and F1-score later on, proving effectiveness for improving model accuracy. Its data-parallel strategy can leverage distributed nodes, compensating for singledevice limitations, which is crucial for large-scale training.

5.4 Comparison with Other Data Parallel Frameworks

Training Time: Figure 5 compares the training processes of the three frameworks, measured in time. C-ADP optimizes data allocation, improves the throughput of large-scale models, and completes the training in the shortest time. DDP fails to converge for a long time, and its accuracy decreases in the later stage. Deepspeed, due to frequent parameter updates, incurs high communication costs during the training of largescale models, which slows down the training speed. Through comprehensive comparison, C-ADP outperforms the dataparallel frameworks DDP and Deepspeed in terms of training speed. It can converge smoothly and quickly and complete the model training.

Comparison of FLOPS: As illustrated in Figure 6, during the evaluation, the device FLOPS under the C-ADP training method consistently stabilizes at 8.95 TFLOPS, achieving 21.6 times improvement over DDP and 26.3 times improvement over DeepSpeed. In contrast, other methods exhibites significant performance fluctuations due to the impact of heterogeneous computing and communication capabilities of all computing devices, leading to reduced resource utilization efficiency. By optimizing the data allocation strategy and reducing the parameter synchronization frequency, C-ADP effectively mitigates the impact of communication bottlenecks on training performance. This approach not only significantly improves computational efficiency but also ensures system robustness and resource utilization, demonstrating its advantages in large-scale distributed training.

Quantization Result Comparison: Table 3 compares the performance of the three distributed training frameworks, including DDP, Deepspeed, and C-ADP. C-ADP exhibits distinct advantages. Due to its adaptive data allocation strategy, C-ADP only requires a data transfer time of 431.82s, much less than the 1411.68s by DDP and 1398.46s by Deepspeed. The C-ADP framework achieves even loading across all devices by fully addressing heterogeneous computing and communication constraints in the system.

Result Analysis: Our C-ADP framework focuses on finding the optimal data distribution to all the computing devices in distributed training. As shown in Table 3, C-ADP converges after 12915.28s and acheives an accuracy of 74.36%, while other two frameworks. As a result, C-ADP can achieve convergence more rapidly during the training process. In summary, C-ADP significantly reduces the total time for distributed training and achieves a reduction in trainning time of about 72% and 47%, respectively, compared with DDP and DeepSpeed, making it a more efficient and reliable option for large-scale deep learning model training.

6 Conclusion

In this paper, we explore the opportunity to train deep learning models in a distributed manner with computing devices connected in a heterogeneous environment. We have designed an efficient and effective data parallel training framework called C-ADP that can find optimal data distribution in the asymptotic sense. Experiments on Wide-ResNet101 with the ImageNet100 dataset show that C-ADP achieves excellent convergence performance in heterogeneous environments, reaching 73.67% accuracy at the 50th round of training and outperforming single-device training in key metrics such as F1- score. Compared with DDP and DeepSpeed, C-ADP achieves 21.6 times and 26.3 times improvements in FLOPS, respectively, and a reduction in training time of about 72% and 47%, respectively.

A Proof of Theorem 1

Proof. Consider an optimization domain DN = {(x1, x2, ..., xN)| ∑Ni=1 xi ≤ D, xi ≥ 0}. Obviously, this domain is convex. Thus, according to Lemma 2, for any x, y ∈ DN , λ ∈ (0, 1), the first component on the right-hand side of equation (10) is

which is a convex function. Similarly, we can easily show that the second component on the right-hand side of equation (10) is a convex function as well. Therefore, we can have that ˜Φ(λx + (1 - λ)y) ≤ λ˜Φ(x) + (1 - λ)˜Φ(y), and the function ˜Φ(x) is convex on the domain DN . Consequently, ˜Φ(x) is also a convex function on the boundary ∂DN = χ′ = {(x1, x2, ..., xN)| ∑Ni=1 xi = D, xi ≥ 0} of this domain.

B Proof of Theorem 2

Proof. Let δ = ˆx- ˜x, where ˆx ∈ χ, and ˜x ∈ χ′ . Then, we have ∥δ∥p ≤ γ. According to equations (9) and (10), we can get

For the first component on the right-hand side of the above inequality, we have

where (a) is due to the sub-multiplication property [Horn and Johnson, 2012], and hence

For the second component on the right-hand side of the inequality in (11), letting ˆy = B-1 ˆx, ˜y = B-1 ˜x, we can have

where (b) has been proved in [Weisstein, 2023]. Combining equations (12) and (13) and letting Mγ = bγ · ∥∥V-1∥∥p + num · ∥T∥p (N 1p + γ ∥∥B-1∥∥p ), we have