Unity is Power: Semi-Asynchronous Collaborative Training of Large-Scale Models with Structured Pruning in Resource-Limited Clients

In this section, we focus on the balanced structured pruning of the globally large models at both depth and width dimensions, as shown in Fig.3. For the heterogeneous clients with weak computing power, we obtain the overall pruned rate $R_{n}$ according to the available resources, which consists of the width pruned rate $R_{n}^{width}$ and the depth pruned rate $R_{n}^{depth}$. The pruned rate is defined as the ratio of the structured-pruned model size of the full model size. For client $n$, the server first generates the depth-pruned submodel with $\lfloor L\times R_{n}^{depth}\rfloor$ consecutive trainable blocks by freezing the shallow blocks bottom-to-top and pruning the deep blocks top-to-bottom in a rolling fashion.

Then in order to make the submodels fit the heterogeneous data distributions, the clients train the width-based structured masks based on the local data distributions. A Transformer-like model mainly consists of linear layers and Multi-head Self-Attention(MSA) layers. Consequently, we design trainable segment-wise masks for two types of layers, where the segment represents part parameters of linear layer or head in MSA layer. We denote the weight matrix and its binary mask by $w_{n}^{l_{i}}\in\mathbb{R}^{d_{out}\times d_{in}}$ and $M_{n}^{l_{i}}\in\{0,1\}^{d}$ respectively, where $l_{i}$ represents the $i$-th layer. For the linear layer, each value of $M_{n}^{l_{i}}\in\{0,1\}^{d_{out}}$ indicates whether the part of the layer is pruned or not. For the MSA layer, each value of $M_{n}^{l_{i}}\in\{0,1\}^{h}$ indicates whether the head is pruned or not and $h$ represents the number of heads.

However, considering that the segment-wise mask $M_{n}^{l_{i}}$ is a binary tensor, directly applying the existing back-propagation methods to model training would lead to gradients backward errors. Inspired by [44, 13], we introduce the corresponding segment-wise importance mask $I_{n}^{l_{i}}\in\mathbb{R}^{d}$ and segment-wise probability mask $P_{n}^{l_{i}}\in[0,1]^{d}$. To enable back-propagation, we apply the sigmoid function scaling to the importance mask, and further use the Bernoulli distribution [7] to generate the binary mask $M_{n}^{l_{i}}\sim Bern(\frac{1}{1+e^{-I_{n}^{l_{i}}}})$. Then we use the inverse of the sigmoid function and record $P_{n}^{l_{i}}$ with the straight-through estimator [2] for back-propagation.

Furthermore, the binary matrix mask $\hat{M}_{n}^{l_{i}}\in\mathbb{R}^{d_{out}\times d_{in}}$ is introduced as an extension of $M_{n}^{l_{i}}$ to produce Hadamard product with the weight matrix of the model. For training the personalized mask $M_{n}^{l_{i}}$ efficiently, we freeze the corresponding weights of the depth-pruned submodels. Considering the resource-limited problem, we use the submodel size to reflect computation and communication costs. The loss function is as follows:

where $\mathcal{L}_{CE}(\hat{y},y)$ represents the cross entropy loss for target class $y$, $S$ represents the set of the linear layers and the MSA layers. $\lambda_{1}$ is the hyperparameter that indicates the extent to which the submodel adapts to the limited resource. For each client, $\frac{\sum_{l_{i}\in S}sum(\hat{M}_{n}^{l_{i}})}{\sum_{l_{i}\in S}size(\hat{M}_{n}^{l_{i}})}$ stays as close as possible to $R_{n}^{width}$ to ensure a balanced learning capability of submodel.

Finally, we perform Hadamard product for the depth-pruned submodel and binary matrix to obtain the structured-pruned submodel. The structured pruning strategy significantly reduces the capacity of model, saves storage and computation resources, and accelerates directly on resource-limited clients.

To mitigate the knowledge loss problem in resource-limited clients, $Co\text{-}S^{2}P$ introduces lightweight self-distillation mechanism [38] to implement cross-block knowledge transfer. For the structured-pruned submodel, we place multiple classifiers at specific depths $h$, where $h\in\{h^{frozen}+L\times{R_{i}^{depth}}|{R_{i}^{depth}}\leq{R_{n}^{depth}},1\leq i\leq N\}$, and $h^{frozen}$ is the number of the frozen blocks and $L$ is the number of blocks in the global model. We use self-distillation by treating the deepest classifier as the teacher and other classifiers as the students. The loss function for the width pruning is as follows:

where $c_{n}$ is the number of classifiers in client $n$, $\mathcal{L}_{CE}(\hat{y}_{i},y)$ represents the cross entropy loss of the $i$-th classifier for target class $y$. $\mathcal{L}_{KL}(\hat{y}_{i},\hat{y}_{c_{n}})=sum(\sigma(\hat{y}_{i}/t)\ln{\frac{\sigma(\hat{y}_{i}/t)}{\sigma(\hat{y}_{c_{n}}/t)}})t^{2}$, temperature t is a hyperparameter that controls the information capacity of data distributions provided by the teacher, $\lambda_{2}$ indicates the balance the cross entropy loss and self-distillation. The self-distillation mechanism enables the shallow blocks to be equipped with high-level knowledge without extra forward overhead, because the blocks share feature maps.

Due to the stragglers, we design a semi-asynchronous aggregation strategy with minimum ratio $\mu$ of received clients and waiting interval $T_{clk}$ to balance the performance of the collaborative training and the resource utilization of all clients, which is shown in Algorithm 1. After the server receives $\mu N$ client updates, it turns on the clock timing and waits for $T_{clk}$ seconds before aggregation.

Considering that the updates may be too stale and deviate significantly from the latest global model, for received gradients, we compute a segment importance score on each trained segment $i$ as follows:

where $\Delta_{q,n}^{i}$ is the gradients of the client $n$ on segment $i$ in $q$-th round that indicates the extent of updating, $w_{q}^{i}$ represents the weights of the global model on segment $i$ in $q$-th round. $w_{q-\tau_{q,n}}^{i}$ represents the weights of the global model on segment $i$ in the $(q-\tau_{q,n})$-th round, $\tau_{q,n}$ is the delay round that the client $n$ has not taken part in global aggregation. $\|w_{q}^{i}-w_{q-\tau_{q,n}}^{i}\|_{1}$ indicates the difference of the received model by the client $n$ and the latest model.

Finally, the server aggregates the gradients and updates segment $i$ by normalizing the segment important scores on segment $i$ as follows:

where $N_{q}^{i}=\{n:M_{n,q}^{i}=1\}$ represents the set of clients training segment $i$ in the $q$-th round, $\eta$ is the learning rate for training the masked submodel. It is worth noting that the aggregation strategy also transfers the high-order knowledge from the deep blocks to shallow blocks in all clients.

Testbed Implementation. To effectively demonstrate the real-world applicability of the proposed collaborative training framework, we utilize a GeForce RTX 3090 GPU as the server for aggregation and 4 types of 16 Jetson developer kits as the resource-limited clients to build a real-world testbed, as shown in Fig.4. The details are shown in Tab.5 of Appendix.D.

Baselines. To fully evaluate the performance and efficiency of the proposed framework ${Co\text{-}S}^{2}{P}$, we choose the classical FL algorithm FedAvg [21]. Additionally, we combine FedAvg with existing standalone training for sparsifying transformer-like models: PLATON [39], AdaViT [22] and Q-ViT [18], named $FedAvg+$. Moreover, we evaluate ${Co\text{-}S}^{2}{P}$ against several FL algorithms designed for resource-limited scenarios, including FedPM [13], PruneFL [14], RAM-Fed [31] and FedRolex [1].

Backbones and datasets. We use three ViT [9] variants with different capabilities, as shown in Tab.4. We conduct experiments based on ImageNet200 and ImageNet300 sampled from ImageNet1K [27] and design a evaluation metric for the real-world testbed, namely Resource Utilization Rate (RU).

We deploy ${Co\text{-}S}^{2}{P}$ and all baselines on the testbed based on the general FL framework NVFlare [28]. Considering that the unstructured baselines fail to work in the resource-limited clients, we replace the unstructured pruning in the baselines with the structured pruning. We conduct three groups of real-world experiments, and the detailed configurations are shown in Tab.6 and Tab.7. Due to limited space, please refer to Appendix D for additional details of experimental implementation.

Server’s Performance. As shown in Tab.1 and Tab.2, ${Co\text{-}S}^{2}{P}$ outperforms all baselines in terms of global performance across three groups of experiments: (1) ${Co\text{-}S}^{2}{P}$ improves Top1 accuracy of global model by 8.8%/6.4%/8.7% compared to the closest baseline FedRolex in three groups of experiments respectively, demonstrating that our algorithm achieves better generalization in resource-limited scenarios. (2) For Top1 accuracy, the $FedAvg+$ baselines significantly decrease $2.0\%\sim 14.5\%$ compared to the federated resource-adaptive methods, indicating that the mixed heterogeneity of collaborative training scenarios has severe implications for model training. (3) Considering Top1 accuracy, ${Co\text{-}S}^{2}{P}$ achieves 5.4%/3.0%/3.5% higher improvement compared to the fully asynchronous algorithm on three groups of experiments respectively. This highlights that ${Co\text{-}S}^{2}{P}$ effectively balances resource utilization and model performance.

Clients’ Average Performance. As shown in Tab.1 and Tab.2, ${Co\text{-}S}^{2}{P}$ outperforms all baselines in terms of weighted average performance of the clients in three groups of experiments: (1) For Top1 accuracy, ${Co\text{-}S}^{2}{P}$ achieves 11.1%/8.8%/10.1% higher personalized submodel performance compared to the closest baseline FedRolex in three groups of experiments respectively, demonstrating that our framework obtains a better personalized mask for each client in resource-limited scenarios. (2) Comparing the difference between global and local performance, ${Co\text{-}S}^{2}{P}$ has closer gap than all baselines. This is because the trainable masks better adapt to local data distributions and the depth-pruned submodels in resource-limited clients. (3) Considering Top1 accuracy, ${Co\text{-}S}^{2}{P}$ achieves 5.1%/3.1%/4.4% higher personalized model performance compared to the designed asynchronous methods on three groups of experiments respectively, which demonstrates that ${Co\text{-}S}^{2}{P}$ achieves a better balance between resource utilization and personalized submodel performance.

Efficiency Comparison. As shown in Tab.1, Tab.2 and Fig.5, ${Co\text{-}S}^{2}{P}$ outperforms all baselines in terms of efficiency in three groups of experiments: (1) Compared to FedRolex, ${Co\text{-}S}^{2}{P}$ reduces memory consumption by about 22% and training time per round by about 24% on all resource-limited devices, demonstrating the superiority of the structured pruning in memory and computation efficiency. (2) For resource utilization rate, ${Co\text{-}S}^{2}{P}$ achieves 1.4$\times$, 1.5$\times$, and 1.5$\times$ improvements compared to FedAvg and 1.2$\times$, 1.1$\times$, and 1.2$\times$ improvements compared to FedRolex across the three groups of experiments. This demonstrates that the proposed semi-asynchronous aggregation strategy can effectively improve the resource utilization during large-scale model training.

Effectiveness of trainable mask policies. We validate that the trainable mask strategy is effective in improving global performance and generating personalized masks in resource-limited scenarios. We replace the trainable masking strategies with randomly mask strategies and report the results in Tab.3. We observe that replacing any set of trainable masks with randomized ones leads to a performance decrease. With MSA/Linear/Both Random, $Co\text{-}S^{2}P$ improves global model accuracy by 5.3%/6.6%/10.6%, and the average accuracy of submodels on the client by 6.5%/7.5%/12.9% in Top1 accuracy. These confirm the effectiveness of the trainable mask strategy.

Effectiveness of self-distillation mechanism. We validate that the self-distillation mechanism helps shallow submodels on resource-limited clients to learn high-level and complex knowledge. As shown in Tab.3, $Co\text{-}S^{2}P$ improves Top1 accuracy by 4.9% for the global performance and by 9.5% for the average performance of submodels. This demonstrates the effectiveness of self-distillation in enabling shallow models on resource-limited clients to obtain more complex and high-level knowledge.

Impact of minimum received ratio $\mu$ and waiting interval $T_{clk}$. We study the impact of $\mu$ and $T_{clk}$ on model performance and the resource utilization of the system. As shown in Fig.6, increasing $\mu$ or $T_{clk}$, improves the performance of both the global model and the submodels, but reduces the resource utilization rate due to requiring a longer wait for the other clients. In addition, we provide more hyperparameter analysis details of weight term $\lambda_{1}$ of $\mathcal{L}_{n,mask}$, balance term $\lambda_{2}$ of $\mathcal{L}_{n,weight}$ and the ratio between of $R_{n}^{width}$ and $R_{n}^{depth}$ in Appendix F.2.