An Efficient Split Fine-tuning Framework for Edge and Cloud Collaborative Learning

Existing popular PLMs in NLP are mainly built on the Transformer [17] architecture. As shown in Fig. 1(a), a PLM consists of a series of transformer layers (or blocks), each of which contains an attention layer and a feed-forward layer (FFN) with residual connections. Since each layer requires a residual connection (Add & Norm), the input size and the output size of the layer should be identical, which means the shapes of all input and output tensors within the transformer layer or between transformer layers are identical. This is an important feature causing extensive communication traffics in SL for fine-tuning such models. See below for more details.

Split learning (SL) [10, 18] has provided an emerging solution for edge and cloud collaborative learning without exposing the private data on the edge while enjoying the powerful computational resources from the cloud servers. As shown in Fig. 1(b), a DNN model is split into two parts at a particular layer. One part (say $net_{1}$) is stored on the edge and the other part (say $net_{2}$) is stored on the cloud. The number of layers on each part may be different, and more layers typically take higher computational costs. Thus, in SL, the data located on the edge is loaded for feed-forward computation on $net_{1}$, and its activation output is transferred (upload) to the cloud. We assume the activation output is a tensor ${\bf\it a}$ with shape $(B,M,N)$, where $B$ is the mini-batch size and $M$ and $N$ are two dimensions of a matrix of the hidden feature representation. Note that $M$ and $N$ may be the explicit dimensions of the hidden feature as different types of layers (e.g., FFN, attention, etc.) have different structures, but they all can be organized as a matrix. For the transformer-based architecture, as we analyzed in the above subsection, the shape is the same no matter which layer is chosen as the split layer in SL. During the backpropagation pass, the gradient w.r.t. ${\bf\it a}$, which is denoted as ${\bf\it\delta}$ and has the same shape with ${\bf\it a}$, is calculated on the cloud side and transferred to the edge for calculating the gradients w.r.t. the model parameters for updating the model.

It is seen that the computations of training have been partially ($net_{2}$) offloaded to the server to utilize the powerful server to reduce the training time for edge devices. However, the exchange data (i.e., ${\bf\it a}$ and ${\bf\it\delta}$) has a large number of elements, which requires significant communication time at each training iteration and becomes the performance bottleneck. In this work, we will present how to reduce the communicating volume of ${\bf\it a}$ and ${\bf\it\delta}$ without sacrificing model accuracy.

Some works in reducing ${\bf\it a}$ and ${\bf\it\delta}$ are activation compression techniques [19, 20, 21, 16, 22]. These studies try to compress the activation outputs for each layer to save computational costs and memory for storing the temporary data in activation, which can be classified as the quantization method and the pruning method with tensor decomposition. Particularly, Evans et al., [20] propose the AC-GC lossy compression algorithm to dynamically compress the activation via quantization and an optimization objective of maximizing the compression ratio while minimizing the accuracy loss. While the quantization technique can maximally reduce the storage of ${\bf\it a}$ (with 1-bit) by 32 times compared to the 32-bit counterpart, too aggressive compression easily introduces an accuracy loss in model training [20] in practice. The tensor decomposition approaches [14, 16, 21] are to decompose the parameter tensors to small-size approximate tensors, thus reducing the computational and memory costs. Particularly, MPO [14, 16] decomposes an original parameter matrix (say ${\bf\it w}$) into $n$ ($n\geq 2$) small matrices (i.e., ${\bf\it w}\rightarrow[{\bf\it w}_{1},{\bf\it w}_{2},...,{\bf\it w}_{n}]$). Even though MPO achieves high compression ratios with slight accuracy loss, it requires particular tricks like dimension squeezing and takes extra computational costs for finding the layer with the least reconstruction error, which makes the optimization procedure difficult to be applied in edge and cloud collaborative learning.

The architecture overview of our proposed SFT is shown in Fig. 1(c). The key idea is two folds: 1) decomposing the FFN layer into three smaller FFN layers (FFN-1, FFN-2, and FFN-3) after loading the pre-trained parameters, and 2) the residual connection in the original FFN layer is eliminated, but the model convergence is guaranteed. The communication volume between the edge and cloud in the split layer of FFN-1 becomes much smaller than the original FFN output, thus significantly reducing the communication time in collaborative learning. The details are provided as follows.

As the high communication cost is caused by the large dimension of ${\bf\it a}$ and ${\bf\it\delta}$ at each iteration, we propose to decompose ${\bf\it w}$ to multiple small tensors, so that the generated activation output is with small dimensions. Specifically, given a transformer-based model, we first load the pre-trained parameters to the original model architecture. In SFT, we only need to compress the layer that needs to be communicated between the edge and the cloud. According to our observation that decomposition does not affect the model accuracy (§IV-B), the weights in FFN layers in fine-tuning PLMs are mostly low-rank matrices. We choose the FFN layer as the split layer in SL. Formally, for an FFN layer at layer $l$ with weight matrix ${\bf\it w}_{l}$ and the input ${\bf\it a}_{l-1}$ (the output of its previous layer), its output can be represented as

Assuming that SL splits the DNN into two parts at layer $l$, ${\bf\it a}_{l}\in\mathbb{R}^{M\times N}$ should be communicated from the edge to the cloud in the forward pass, and its gradient ${\bf\it\delta}_{l}\in\mathbb{R}^{M\times N}$ should be communicated from the cloud to the edge in the backward pass. Note that in fine-tuning tasks, ${\bf\it w}_{l}$ has been initialized with pre-trained parameters. Our goal is to compress ${\bf\it a}_{l}$ and ${\bf\it\delta}_{l}$ such that the communication volume is small enough to eliminate the communication bottleneck in SL.

In SFT, we decompose the split layer (i.e., layer $l$) whose weights are denoted by ${\bf\it w}\in\mathbb{R}^{N\times H}$ to three matrices via singular value decomposition (SVD), i.e.,

where ${\bf\it u}\in\mathbb{R}^{N\times R}$, $\Sigma\in\mathbb{R}^{R\times R}$ is a diagonal matrix whose diagonal entries are singular values of ${\bf\it w}$, and ${\bf\it v}\in\mathbb{R}^{R\times H}$. $R\leq\min\{N,H\}$ is the rank of ${\bf\it w}$. With the decomposed matrices, Eq. (1) becomes

Let $\hat{{\bf\it a}}_{l}={\bf\it a}_{l-1}{\bf\it u}$ and $\hat{{\bf\it a}}_{l}\in\mathbb{R}^{M\times R}$. Note that $\hat{{\bf\it a}}_{l}={\bf\it a}_{l-1}{\bf\it u}$ is equivalent to constructing an FFN layer whose weight is ${\bf\it u}$. Similarly, ${\bf\it a}_{l}=\hat{{\bf\it a}}_{l}(\Sigma{\bf\it v})$ is equivalent to constructing two FFN layers whose weights are $\Sigma$ and ${\bf\it v}$ respectively. Instead of splitting the DNN in the original architecture at layer $l$, SFT splits its decomposed form at ${\bf\it u}$. It means that the activation output on the edge side is $\hat{{\bf\it a}}_{l}\in\mathbb{R}^{M\times R}$, which should be communicated to the server side in the forward pass. The corresponding gradient is $\hat{{\bf\it\delta}}_{l}\in\mathbb{R}^{M\times R}$ in the backward pass. Thus, the communication volume is reduced from $M\times N$ to $M\times R$, i.e., the communication time is shortened by $N/R$ times theoretically.

Note that SFT decomposes a single FFN layer to three smaller FFN layers after loading the pre-trained parameters but before fine-tuning. The three constructed FFN layers are initialized with the SVD decomposition and they are tuned every iteration in fine-tuning. Due to the low-rank feature of the weight matrix of the FFN layer during fine-tuning, we can choose an extremely small value of $R$ without affecting the convergence performance in fine-tuning. In our experiments, setting $R=1$ or $R=8$ can almost preserve the model accuracy (§IV). In summary, SFL can reduce the communication traffic by $N/R$ times over SL.

The SFT algorithm is shown in Algorithm 1, where some comments illustrate that some code is only executed on the edge and some are only on the cloud. The inputs are the neural network architecture $net$, the split layer $l$, and the number of iterations $I$ for fine-tuning. In Algorithm 1, lines 1-3 are splitting the model, loading pre-trained models, and reconstructing the split layer based on SVD. The for loop in lines 4-14 is the training procedure in both the edge and cloud sides. For each iteration, the edge first loads the training data (line 5) and then performs the feed-forward on the first part of the network (line 6), whose results ($\hat{{\bf\it a}}$) are sent to the cloud (line 7). After the cloud receives the activation from the client, it performs feed-forward on the second part of the network (line 8), followed by the backward computation with loss (lines 9-10). The gradient ($\hat{{\bf\it\delta}}$) w.r.t. the activation is sent from the cloud to the edge (line 11), followed by back-propagation on the edge side (line 12). After that, the models are updated simultaneously on the edge and the cloud (lines 13-14). Our system enables users with existing training scripts to be able to explore the cloud to fine-tune their models without exposing the data to the server.

SFL has two main goals: 1) making the low-memory edge device possible for fine-tuning when the edge cannot store the whole model, and 2) enabling edge devices to explore the cloud servers to fine-tune a model more efficiently. The first goal is obviously our advantage using SFT when low-memory devices cannot store the whole model. For the second goal, however, we should consider whether SFT can reduce the overall fine-tuning time. Assume that the original iteration time on the edge to fine-tune a model $net$ is denoted by $t_{naive}=t_{edge}(net)$. When we use SFT, the fine-tuning time per iteration is

where $t_{cloud}(net_{2})$ is the computation time on the cloud and $t_{comm}$ is the communication time between the edge and the cloud. Thus, SFT can be used if $t_{sft}<t_{naive}$. As the cloud server typically has much higher computational power than the edge device, $t_{edge}(net_{1})+t_{cloud}(net_{2})$ should become much smaller than $t_{naive}$. However, $t_{comm}$ should not be too large so that one can enjoy the efficiency of SFT.

$t_{edge}(net_{1})$ and $t_{cloud}(net_{2})$ depend on the split layer. A lower split layer indicates that more computational workloads are uploaded to the server for calculation, which would make $t_{edge}(net_{1})$ smaller and $t_{cloud}(net_{2})$ larger, and vice visa. As the cloud is much more powerful than the edge, we expect the split layer to be as low as possible, but the lower layer may sacrifice the model accuracy (§IV-B). Therefore, we have a trade-off between accuracy and efficiency in SFT. $t_{comm}$ depends on the rank we used for decomposition. A larger rank makes $t_{comm}$ higher. Therefore, the two parameters (i.e., split layer and rank) should be well-tuned for achieving better end-to-end training performance in SFT. In this work, we do not develop a strategy for tuning them, but we present some observations from the experiments for helping tune them.

To enable users to easily use our SFT framework, we implement a distributed optimizer named “SFTOptimizer” atop the PyTorch optimizer to integrate the layer decomposition of the model and the communication between the edge and the cloud. To use SFT, users only need two more lines of code to extend their original training scripts. An example is shown in Listing 1, where lines 2-3 are inserted into the original training code to use SFT.

Testbed. We use two Nvidia Tesla V100 GPUs in a single node as an emulation environment. One GPU is used to simulate an edge device, and one GPU is used to work as a cloud server.

DNN and datasets. We use the popular language model, BERT_BASE [2], which has 12 layers and 110M parameters, as our experimental neural network. The dimension of the split layer is $(M,N)=(3072,768)$. Its corresponding pre-trained model is downloaded from the well-trained parameters at HuggingFace²²2https://huggingface.co/bert-base-uncased. We fine-tune the model on 9 datasets from GLUE [8] and SQuAD [9] for different downstream tasks including named entity recognition, textual entailment, co-reference resolution, etc.

SFT w/ the residual connection. We first demonstrate the convergence results of the SVD decomposition for replacing the large FFN with three small FFNs for fine-tuning, which shows the low-rank feature of FFN weights. We use SST-2 and QNLI datasets to study the convergence property of our SFT as shown in Fig. 2. Due to the page limit, we do not show the results of other datasets as they have similar patterns. The results show that the weight matrix in the split layer can be decomposed as a rank-1 matrix without sacrificing the model accuracy. In some scenarios, rank-1 decomposition is better than the baseline. For example, decomposing at layers 2 or 3 (11 or 12)³³3The smaller layer index means the lower layer is much better than the baseline in SST-2 (QNLI). How to determine which layer should be split to achieve the highest efficiency in a given environment is also of importance to explore, and we will leave it as our future work.

SFT w/o the residual connection. Since we also need to eliminate the residual connection in the original transformer layer in SFT, we conduct experiments to verify the convergence performance by eliminating the residual connection and decomposing the FFN with SVD. The results are shown in Fig. 3, which shows that the model accuracy tends to decrease when the split layer is chosen at lower layers. However, it can still preserve the model accuracy when splitting at higher layers. Thus, it is possible to use SFT to enjoy the powerful cloud server to accelerate the fine-tuning tasks for edge devices without sharing the data.

Model accuracy in all datasets. Based on the convergence results in Fig. 3, we use a rank of 8, 16, and 32 for decomposition at layer 11 in SFT (denoted as SFT(l=11,R=8), SFT(l=11,R=12), and SFT(l=11,R=32) respectively). Therefore, the communication cost can be reduced by 768/8=96 times in the chosen BERT model if R=8. The model accuracy of validation sets in all chosen 9 datasets is shown in Table I. The results show that different ranks have little differences and a higher rank does not guarantee a higher accuracy. For example, SFT(l=11,R=8) achieves higher accuracy than SFT(l=11,R=16) on QQP and RTE datasets. In most cases, SFT(l=11,R=8) preserves the model accuracy compared to the baseline. In the particular case of RTE, SFT(l=11,R=8) achieves much lower accuracy than the baseline due to the extremely small size (only around 2,500 training samples) of the dataset. The results show that SFT (using a rank of 8) makes the fine-tuning tasks possible to train in a collaborative learning environment so that the edge does not need to expose the data to the server.

To show the end-to-end performance of our SFT compared with SL on the BERT_BASE model, we benchmark the time performance on a V100 GPU using SST-2 (other datasets have similar patterns). Let $t_{bert}(gpu,nlayers)$ denote the wall-clock time per iteration training the BERT_BASE model on a particular $gpu$ using $nlayers$ layers of the model. The full 12 layer BERT_BASE on SST-2 takes

with a mini-batch size of 32 and a sequence length of 66. Since the 12 layers of BERT_BASE have an identical architecture, each layer has the same computational workloads and thus has the same computation time. We can estimate the iteration time of one layer as

Assume that the edge side is a relatively new Nvidia edge device, XAVIER-NX, with 21 TOPs AI performance, and the cloud side is a V100 GPU with 130 TOPs AI performance, which means the cloud server is around 6 times faster than the edge device. Therefore, on a XAVIER-NX GPU, we have

and

The communication traffic with SL is $32\times 3076\times 768\times 4=288$MB, while it is $32\times 3076\times 8\times 4=36$MB with our SFT. Assume that the typical bandwidth between the edge to the cloud is 1000Mbps Ethernet⁴⁴4We also assume the bandwidth can be fully utilized, while it cannot in practice. The main purpose of the emulation is to demonstrate the potential benefits of our SFT in real-world environments, and the performance also varies with different configurations., the communication time of SL and SFT is 2,300ms and 24ms respectively. Thus, the iteration times with local training (on an edge device), with SL (on an edge and a cloud), and with SFT (on an edge and a cloud) are

and

respectively. The results show that the introduced communication time in SL is very high, making it even slower than the local training. Our SFT, on the other hand, can achieve faster training time (14% reduction in the end-to-end training time) by reducing the communication volume between the cloud server and the edge device.

Our experimental results conclude two folds. First, when the edge devices cannot conduct fine-tuning tasks due to their memory constraint, it is possible to enjoy our SFT to fine-tune the model. Second, even if the edge clients can fine-tune the model locally, our SFT can help accelerate the training by exploring the more powerful cloud servers without requiring the data from the edge devices. However, there are still two important problems that should be further studied to make SFT more practical. First, could it be possible to split the layer in the lower layer of the model such that more computational workloads can be offloaded to the server? In our existing results, splitting the lower layer may introduce some accuracy loss. Thus, it is a trade-off between the model accuracy and training efficiency. Enjoying higher training speed may sacrifice some accuracy. Second, as keeping the residual connection can preserve the model accuracy (as shown in Fig. 2) even using rank-1 decomposition, could it be possible to keep the residual connection without introducing significant communication costs between the edge and the cloud? As SFT decomposes the FFN to smaller FFNs which are distributed to the client and server, the residual connection should require the activation data to be transferred between the edge and the cloud, which makes the communication extremely heavy. However, eliminating the residual connection only allows higher layers to be split layers, which means only a small proportion of layers can be uploaded to the server if we want to preserve the model’s accuracy.