Pipeline Parallelism for Inference on Heterogeneous Edge Computing

Pipeline Parallelism. Pipeline parallelism partitions a neural network model into multiple stages, where each stage consists of a consecutive set of layers in the original model Yang et al. (2021); Narayanan et al. (2019). Each stage in the pipeline is assigned to a worker to achieve the parallel execution of model training or inference. The input minibatch is split into multiple chunks of equal size, which are called microbatches Huang et al. (2019). The microbatch size affects the pipeline performance, with the optimal size depending on multiple factors including the characteristics of the model and the number of pipeline stages Narayanan et al. (2021b). A worker in a system that has pipeline parallelism need only send its output data to a single worker, which avoids the collective communication to synchronize results with all workers. Pipelining can also overlap computation and communication to improve the performance Narayanan et al. (2019).

Transformer-based Models. The transformer model was proposed to improve the effectiveness in learning dependencies between distant positions for sequence modeling tasks Vaswani et al. (2017). A transformer encoder includes multiple transformer layers with identical structures. Every transformer layer is composed of a multi-head self-attention layer, a multi-layer perceptron (MLP), two layer normalization operations, and residual connections. The multi-head self-attention layer calculates the attention score of the input sequence $A=(a_{1},...,a_{n})$ through dot product operations and generates the output representation $B=(b_{1},..,b_{n})$ with the same dimension. Outputs $b_{i}$ and $b_{j}$ (where $i\neq j$) are operationally independent, which provides the possibility of parallel execution for both training and inference Park et al. (2020b). Attention-based models have recently been extended to replace conventional convolutional neural networks (CNNs) Dosovitskiy et al. (2020); Kundu and Sundaresan (2021) in performing complex CV tasks. In particular, the ViT models enjoy superior representation ability d’Ascoli et al. (2021), and also suffer less from positional invariance issues which are prevalent in conventional CNNs Su et al. (2019).

Large-scale model inference is a challenging task for resource-limited devices. Classical model compression techniques, including pruning Zhang et al. (2018); Kundu et al. (2021), quantization Yao et al. (2021), low-rank approximation Chin et al. (2020), and knowledge distillation Hinton et al. (2015) can shrink neural network model sizes to potentially accelerate deep neural networks. However, these techniques often require iterative retraining and a full-precision pre-trained model to avoid significant accuracy-drop. More importantly, these methods generally focus on a single compute node. Distributed edge computing scenarios, in contrast, often include a large number of resource-limited devices, e.g., in vehicular edge computing (VEC) for internet of vehicles Liu et al. (2021), wireless-connected AI-enabled sensors Sharma et al. (2021), and smart home systems Isyanto et al. (2020); Cheng and Kunz (2009).

Pipeline parallelism has proven to be effective for distributed training on accelerators in data centers where devices are relatively homogeneous and the network bandwidth is generally high Huang et al. (2019); Yang et al. (2021); He et al. (2021). However, edge computing has unique characteristics:

• •

  Small Memory Capacity. Compared with data-center servers, edge devices usually have smaller memory capacities, ranging from tens of MB to several GB.

• •

  Heterogeneous Devices. Edge devices have diverse computational performance and memory capacities.

• •

  Limited Bandwidth. Unlike communication within data centers, edge computing systems often rely on wireless communication with limited bandwidth.

• •

  Heterogeneous Network Link Capacities. The bandwidth between different pairs of devices depends on physical distance and channel interference, and could range from tens of Kbps to hundreds of Mbps.

Model partition methods for homogeneous clusters will therefore perform poorly in heterogeneous edge environments. A new pipeline parallelism framework is needed to overcome these challenges. In the next section, we introduce a framework for distributed edge clusters that uses heterogeneity-aware pipeline parallelism to improve inference performance and enable running larger—and more accurate—models.

Figure 2 presents the EdgePipe system design, which consists of three major components: the partitioning algorithm for heterogeneous clusters, the data loader, and the runtime framework that implements pipeline parallelism. First, the configuration of the original transformer model and partition constraints are sent to the partition algorithm to generate an optimal partition method. The partition constraints include available edge devices, computation and memory capabilities of these devices, and the bandwidth between devices. The specific partition problem will be introduced in Section 3.2. To fit into every edge device and improve the throughput, input data should be split into small chunks, called a microbatch. At runtime, selected devices are only responsible for the inference of one part of the original model. After finishing the inference of one microbatch, the edge device transmits intermediate outputs to the device in the next pipeline stage. The device in the final stage produces the final result, which could be transmitted to another host or stored locally.

To construct a pipeline with $k$ stages, the transformer model should be partitioned into $k$ parts $T_{1},T_{2},...,T_{k}$. Part $T_{i}$ includes $l_{i}$ layers, and $\sum_{i=1}^{k}l_{i}=L$, where $L$ is the number of layers. Part $T_{i}$ is assigned to the $i$-th device to construct the $i$-th pipeline stage. For one microbatch input $\mathcal{X}$, the process of inference may be denoted as $\mathcal{\hat{Y}}=T_{k}(T_{k-1}...(T_{2}(T_{1}(\mathcal{X}))))$, where $\mathcal{\hat{Y}}$ is the final output of the inference. The intermediate output of the $i$-th device is sent to the $(i+1)$-th device to continue the computation. The number of pipeline inference stages in EdgePipe is equal to the number of devices participating in inference.

Fully heterogeneous clusters include heterogeneity in both the devices and the communication networks. It is common in edge computing for devices have different computation and memory capabilities. In addition, the network bandwidth between devices may be different. It is therefore challenging to decide the partition method for the cluster.

We define a transformer model $\mathbb{T}$ with $L$ layers, inter-layer transmission data size $P_{j}$ for the $j$-th layer, and a list of heterogeneous devices $\mathbb{D}$ ($|\mathbb{D}|=D$) with different memory, computation, and communication capabilities. In heterogeneous communication, the bandwidth between a pair of devices $u$ and $v$ may be different than the bandwidth between a different pair of devices $u^{\prime}$ and $v^{\prime}$: $b_{u,v}\neq b_{u^{\prime},v^{\prime}}$. The optimal strategy $\mathbb{R}$ partitions the model $\mathbb{T}$ into $S$ parts and allocates them to the selected devices $\mathbb{S}\subseteq\mathbb{D},|\mathbb{S}|=S\leq D$ to achieve maximal throughput and conform to the memory limitations of the selected devices.

We denote $T_{comp}(l,u)$ as the execution time for the set of layers $l$ on a device $u$. $T_{comm}(u,v,P_{j})$ is the time to communicate data $P_{j}$ between devices $u$ and $v$, which is computed by Equation (1), where $b_{u,v}$ is the bandwidth between devices $u$ and $v$.

We assume the pipeline system supports asynchronous communication, and the computation time and communication time are perfectly overlapped. Thus, the maximum latency for the single device $u$ can be calculated as:

For the selected devices, achieving the maximal throughput is equivalent to minimizing the execution time $T_{\text{opt}}$, which is determined by the slowest stage under the given strategy and is equal to the largest $T_{\text{period}}(l,u,v,P_{j})$. The problem of pipeline partitioning can itself be partitioned. The optimal solution for partitioning the whole pipeline on given set of devices can be constructed from the optimal partitioning result for the sub-problem, which could be solved by DP methods.

To tackle this partition problem for fully heterogeneous clusters, we design a three-dimensional DP algorithm recording the state of processed layers, used devices, and the device in the last pipeline stage. Let $h(i,\mathbb{S},u)$ denote the minimum time to process the first $i$ layers with the set of used devices $\mathbb{S}$, and the device $u$ is the next device to be used. $h(i,\mathbb{S},u)$ is the optimal solution of the subproblem for $i$ layers and $\mathbb{S}$ devices. The final optimal solution of this partition problem is the minimum $T_{\text{opt}}=h(L,\mathbb{S},\varnothing)$ with $\mathbb{S}\subseteq\mathbb{D}$.

The calculation of $h(j,\mathbb{S}\cup\{u\},v)$ needs to use the optimal subproblem property: it is determined by the previous state $h(i,\mathbb{S},u),0\leq i<j\leq L$, or the calculation time $T_{\text{period}}(\{i\rightarrow j\},u,v,P_{j})$ from $i$-th layer to $j$-th layer on the current device $u$. We further analyze these two situations:

• •

  $h(j,\mathbb{S}\cup\{u\},v)\leftarrow h(i,\mathbb{S},u)$, the slowest pipeline stage for $j$ transformer layers is determined by the previous stage $h(i,\mathbb{S},u)$. Since device $u$ implements the current stage from the $i$-th to the $j$-th layer in the current pipeline, the used devices set for the next state $h(j,\mathbb{S}\cup\{u\},v)$ should include the device $u$, noted as $\mathbb{S}\cup\{u\}$. Parameters $i$ and $u\in\mathbb{D}\setminus\mathbb{S}$ will be enumerated to find the optimal solution of the subproblem.

• •

  $h(j,\mathbb{S}\cup\{u\},v)\leftarrow T_{\text{period}}(\{i\rightarrow j\},u,v,P_{j})$, the device $u$ that constitutes the slowest stage of the current pipeline and limits the performance of the system. The $T_{\text{period}}(\{i\rightarrow j\},u,v,P_{j})$ is calculated from Equation 2. Similarly, device $u$ and first $i$ layers are enumerated to obtain the minimum value.

Thus, the state transition equation can be formulated as:

where the first term inside the max is the minimum time for the first $i$ layers with the set of devices $\mathbb{S}$ and the next used devices $u$; the second term is communication time of transferring $P_{j}$ data from device $u$ to device $v$; the third term is the computation time for the last $j-i$ layers on the device $u$. For initialization, $h(0,\emptyset,\varnothing)$ is set to $0$.

Equation 3 calculates the optimal pipeline execution time. However, we need to obtain the selected devices and their order in the pipeline for the optimal strategy. Algorithm 1 describes the memoization technique pseudo-code to find the optimal time $T_{opt}$ and the corresponding pipelining strategy.

The computational complexity of the proposed algorithm is $O(2^{D}\times L^{2}\times D^{2})$, where $D$ is the number of available devices and $L$ is the number of layers. The $2^{D}$ factor is due to the assumption that all devices are distinct. As a comparison, for the naive brute force solution, the search space is $\sum_{i=1}^{\min\{D,L\}}\frac{D!}{(D-i)!}\binom{L-1}{i-1}\gg D!\gg 2^{D}$, which has a much higher complexity. Moreover, in most scenarios, there should exist identical devices with the same computation and communication capabilities. Therefore, the number of devices $D$ could be divided into $N$ categories, where the $i$-th category has $n_{i}$ devices ($\sum_{i=1}^{N}n_{i}=D$). The search state of DP can be further reduced, hence the computation complexity could be reduced to $O(\prod_{i=1}^{N}(n_{i}+1)\times L^{2}\times N^{2})$. For instance, consider the case where there are three types of devices, $N=3$, and each type has the same number of devices, i.e., $n_{1}=n_{2}=n_{3}=n$. Then the actual computational complexity is $O((n+1)^{3}\times L^{2}\times N^{2})=O(9\times(n+1)^{3}\times L^{2})$. For example, given $N=3$ device types, where each type has $n=3$ devices, we measure the execution time for these three methods using the ViT-Base model on a 1.6 GHz Intel Core i5 CPU and present the results in Table 2.

We first evaluate EdgePipe’s performance on the 2 GB MinnowBoard and the 8 GB RCC-VE Network Board devices in homogeneous clusters. Figure 3 presents throughput on these clusters for up to 16 stages (devices).

For MinnowBoard devices, we achieve $0.63$ images per second throughput with $4$ devices using the ViT-Base model, which is $1.98\times$ faster compared to the single-device performance. The ViT-Large and ViT-Huge models cannot fit in memory on a single MinnowBoard device. Hence, we use 2-stage and 4-stage throughput as the speedup baselines for the ViT-Large and ViT-Huge models, respectively. With 16 MinnowBoard devices, EdgePipe achieves $1.95$ images per second throughput, which is a $7.48\times$ speedup over the 2-stage baseline (where the optimal speedup is 16/2=8). For the ViT-Huge model, EdgePipe achieves $0.77$ images per second throughput with $16$ MinnowBoard devices, which gives a $3.93\times$ speedup over the 4-stage baseline (where the optimal speedup is 16/4=4).

We achieve similar scalability on the RCC-VE Network Board devices. With the ViT-Base model, EdgePipe achieves $0.82$ images per second throughput with $1.99\times$ speedup with four devices. Compared with the single-device baseline, EdgePipe achieves $2.43$ and $1.01$ images per second throughput for the ViT-Large and ViT-Huge models with $10.59\times$ and $11.88\times$ speedup using $16$ devices, respectively.

The sub-linear performance for the ViT-Base model is primarily due to uneven execution times of different stages. We observe the performance difference for the slowest and fastest stages is about $25\%$ for the 2-stage pipeline. With $4$ stages, this time difference increases to $80\%$ and causes a more serious performance loss. Figure 4 illustrates the issue by quantifying each ViT-Base layer’s execution time on the MinnowBoard. The second dense layer in the $11$-th transformer layer, which only includes the linear transformation with the matrix multiplication operation and thus cannot be further partitioned with pipeline parallelism, requires considerably more execution time than other layers. The variation in execution time is due to differences in the sparsity of weights. We observe the similar performance behavior on the RCC-VE Network Board devices. In contrast, each transformer layer in the ViT-Large and ViT-Huge models have similar inference times, so EdgePipe scales better on the ViT-Large and ViT-Huge models. In addition, we observe compressed models reduce stage imbalances and mitigate this challenge. We report on compressed models in subsection 5.5.

EdgePipe achieves nearly linear performance improvements with the ViT-Large and ViT-Huge models when pipelining with both edge device types. EdgePipe achieves $3.93\times$ speedup (over a 4-node baseline) on 16 MinnowBoard devices and $11.88\times$ speedup on 16 RCC-VE Network Board devices with the ViT-Huge model. These results demonstrate EdgePipe’s effectiveness for large-scale models, including ones that otherwise cannot fit on single devices.

Compared with the data center, edge devices are more heterogeneous in computation and communication capabilities. To better simulate heterogeneous resource-constrained devices, we throttle the CPU usage of the RCC-VE Network Boards to emulate diminished inference performance. We also vary the maximum bandwidth for both MinnowBoard and RCC-VE Network Board devices to emulate different network link capacities. We compare EdgePipe with GPipe and PipeDream on six clusters with increasing heterogeneity. Because GPipe and PipeDream do not specify the device mapping order, we test them with $10$ random device orders and measure average performance and variance. Device configuration details are presented in Table 4, and experimental results are shown in Figure 5.

By comparing Cases 1 and 2, we show the effect of introducing heterogeneous computing capabilities on system performance. For the ViT-Base model, EdgePipe achieves the best performance of $0.82$ images per second in both cases. GPipe and PipeDream show a significant variance with different device orders for the ViT-Base model. In Case 1, GPipe’s throughput ranges from $0.57$ to $0.76$ images per second, and PipeDream’s throughput ranges from $0.64$ to $0.82$ images per second. In Case 2, which introduces more heterogeneity, GPipe and PipeDream show a larger variance of $0.47$ to $0.75$ and $0.50$ to $0.82$ for the ViT-Base model. For the ViT-Large and ViT-Huge models, both GPipe and PipeDream adopt the even partitioning strategy and obtain the same performance. EdgePipe achieves $2.23$ and $1.69$ images per second for the ViT-Large model and $0.88$ and $0.67$ images per second for the ViT-Huge model in Cases 1 and 2, respectively. For the same two cases, GPipe and PipeDream only achieve $1.99$ and $0.76$ images per second for the ViT-Large model and $0.81$ and $0.31$ images per second for the ViT-Huge model. EdgePipe demonstrates both better performance and robustness when compute capabilities are more heterogeneous.

Case 3 has the same compute resources as Case 1, but with less communication bandwidth. We further vary the bandwidth in Case 4 with the same compute resources. For the ViT-Base model, EdgePipe achieves the best throughput of $0.78$ and $0.63$ images per second in both Cases 3 and 4. GPipe and PipeDream show a significant performance degradation due to the limited bandwidth. In Case 3 for the ViT-Base model, GPipe’s throughput ranges from $0.29$ to $0.32$ images per second, and PipeDream’s throughput ranges from $0.32$ to $0.38$ images per second. In Case 4 for the same model, GPipe’s throughput ranges from $0.16$ to $0.20$ images per second, and PipeDream’s throughput ranges from $0.18$ to $0.32$ images per second. For the ViT-Large and ViT-Huge models, EdgePipe achieves the best performance with fewer devices than GPipe and PipeDream for Cases 3 and 4. For the ViT-Large and ViT-Huge models on Case 3, EdgePipe selects $8$ devices with $40$ Mbps bandwidth and one device with $10$ Mbps bandwidth as the last stage and achieves throughput of $1.37$ and $0.53$ images per second, respectively. GPipe and PipeDream use $16$ devices and performance degrades with $1.02$ and $0.44$ images per second for the ViT-Large and ViT-Huge models in Case 3. In Case 4, EdgePipe selects $7$ devices for the ViT-Large model and $9$ devices for the ViT-Huge model to achieve throughput of $1.05$ and $0.51$ images per second, respectively. GPipe and PipeDream achieve $0.55$ and $0.31$ for the ViT-Large and ViT-Huge models. In Case 4, EdgePipe shows $1.90\times$ and $1.54\times$ speedup compared to PipeDream for the ViT-Large and ViT-Huge models, respectively. These two cases demonstrate the effectiveness of EdgePipe’s partitioning strategy for heterogeneous network.

In Cases 5 and 6, we mix the heterogeneity of devices and networks. In Case 5, we added $8$ extremely resource-constrained devices with CPUs at $10\%$ capacity and $20$ Mbps bandwidth. EdgePipe achieves the best throughput with $0.73$, $0.99$, and $0.39$ images per second using $4$, $7$, and $7$ devices for the ViT-Base, ViT-Large, and ViT-Huge models. For the ViT-Base model in Case 5, GPipe’s throughput ranges from $0.22$ to $0.66$ images per second, and PipeDream’s throughput ranges from $0.26$ to $0.69$ images per second. For the ViT-Large and ViT-Huge models in Case 5, GPipe and PipeDream achieve $0.26$ and $0.10$ images per second. In Case 5, EdgePipe shows $1.55\times$, $3.75\times$, $3.84\times$ speedup relative to PipeDream’s average throughput for the ViT-Base, ViT-Large, and ViT-Huge models, respectively. Case 6 shows a scenario with $6$ types of devices weighted toward devices with medium performance. In this case, EdgePipe uses $4$, $12$, and $14$ devices to achieve $0.80$, $1.33$, and $0.57$ images per second for these three ViT models. For the ViT-Base model in Case 6, GPipe’s throughput ranges from $0.18$ to $0.48$, and PipeDream’s throughput ranges from $0.22$ to $0.54$ images per second. GPipe and PipeDream achieve $0.33$ and $0.14$ images per second for the ViT-Large and ViT-Huge models, respectively. EdgePipe achieves speedup of $1.98\times$, $3.98\times$, and $4.16\times$ for the ViT-Base, ViT-Large, and ViT-Huge models compared to PipeDream’s average throughput. Cases 5 and 6 demonstrate EdgePipe’s ability to schedule around low-performance devices and map the task reasonably to achieve the best throughput.

EdgePipe performs significantly better than the GPipe and PipeDream partition methods on all six heterogeneous clusters. Unlike GPipe and PipeDream, EdgePipe successfully avoids the lowest-performing devices by considering multiple factors in exploiting pipelining to improve performance.

Edge computing often has more limited bandwidth between devices compared to the data center. To evaluate the relationship between system performance and bandwidth, we vary the bandwidth between all devices from 120 Mbps to 5 Mbps. We test with 4 pipeline stages using the ViT-Base model and 16 pipeline stages using the ViT-Large and ViT-Huge models. As shown in Figure 6, performance does not decrease significantly until bandwidth drops below 30 Mbps for all three ViT models, demonstrating the feasibility of our system for practical edge applications. Upon further reduction in bandwidth from 30 Mbps to 5 Mbps, although the performance shows a linear decline, the throughput still shows some improvements compared to the single-device baseline. For the ViT-Base model with $15$ Mbps bandwidth, MinnowBoards and RCC-VE Network Boards show about $1.48\times$ and $1.15\times$ speedup compared to the single-device baseline. For the ViT-Large and ViT-Huge models with 5 Mbps, EdgePipe still achieves $2.69\times$ and $3.67\times$ speedup compared to the single device baseline on RCC-VE Network Board devices. These results demonstrate that EdgePipe is still effective for large-scale models under limited network conditions.

As in other pipeline frameworks, performance is affected by microbatch size. We demonstrate the relationship between throughput and microbatch size in Figure 7. For a $2$-stage pipeline, the maximum throughput of the even partitioning method in GPipe is around $0.34$ images per second with a microbatch size $12$. The throughput shows a significant increase before microbatch size $14$, which is mainly due to the continuous increase in CPU utilization as the microbatch size increases. Beyond this size, the throughput begins to slowly decline because larger microbatch sizes reduce the efficiency of the pipeline parallelism. There is a similar pattern for microbatch size and the system throughput for EdgePipe. However, it achieves a significant performance improvement with a maximum throughput of $0.48$ images per second with a microbatch size of $12$. The fine-grained partitioning method in EdgePipe achieves more efficient CPU utilization than even partitioning. Other pipeline stages show similar behavior. To give a fair comparison of throughput, we choose the optimal microbatch and its related performance for all methods in the evaluations.

Model compression techniques shrink the model size to reduce compute cost and potentially accelerate inference Han et al. (2015); Hinton et al. (2015) and training Kundu et al. (2020). These approaches can be considered as important complementary strategies to EdgePipe to improve the performance on resource-constrained platforms. For example, compared to the ViT-Base model, DeiT-Tiny and DeiT-Small use distillation to achieve similar ImageNet top-1 accuracy with compressed models of up to $17.2\times$ and $3.9\times$, respectively Touvron et al. (2021). To demonstrate the efficacy of using compressed models with EdgePipe, we evaluate DeiT-Base, Small, and Tiny on up to $4$ RCC-VE Network Board devices. We also provide the throughput of the ViT-Base model on same devices for comparison. Figure 8 presents the throughput results.

DeiT-Base, which has identical model structure as ViT-Base, achieves $0.62$ images per second throughput on a single RCC-VE Network Board. With $4$ devices, the DeiT-Base model achieves $0.95$ images per second and outperforms ViT-Base’s $0.82$ images per second. DeiT-Small and DeiT-Tiny demonstrate more significant improvements. In particular, DeiT-Small and DeiT-Tiny can provide throughput of up to $5.55$ and $17.23$ images per second, respectively. Interestingly, we observe DeiT-Small and DeiT-Tiny models ease the uneven execution times seen with ViT-Base and demonstrate better scalability with EdgePipe. The model compression technique, as an orthogonal method, can potentially further improve the performance of EdgePipe, making this combined approach a promising solution for large scale model inference at the edge.