Exploring Federated Pruning for Large Language Models

TL;DR

The paper tackles privacy-preserving compression of large language models by proposing FedPrLLM, a federated pruning framework where clients locally compute pruning masks from private calibration data and share only masks for server-side aggregation. Through extensive experiments across multiple LLMs, sparsity levels, and datasets, it finds that layer-based comparison and one-shot pruning without weight scaling achieve the best performance, while weight scaling and iterative pruning offer little to no benefit and incur higher communication costs. The study also shows FedPrLLM robustly outperforms local pruning across varying numbers of clients and calibration samples, underscoring its practicality in privacy-sensitive domains. Overall, FedPrLLM provides a principled, scalable approach to privacy-preserving LLM compression with clear guidance for deployment in real-world settings.

Abstract

LLM pruning has emerged as a promising technology for compressing LLMs, enabling their deployment on resource-limited devices. However, current methodologies typically require access to public calibration samples, which can be challenging to obtain in privacy-sensitive domains. To address this issue, we introduce FedPrLLM, a comprehensive federated pruning framework designed for the privacy-preserving compression of LLMs. In FedPrLLM, each client only needs to calculate a pruning mask matrix based on its local calibration data and share it with the server to prune the global model. This approach allows for collaborative pruning of the global model with the knowledge of each client while maintaining local data privacy. Additionally, we conduct extensive experiments to explore various possibilities within the FedPrLLM framework, including different comparison groups, pruning strategies, and the decision to scale weights. Our extensive evaluation reveals that one-shot pruning with layer comparison and no weight scaling is the optimal choice within the FedPrLLM framework. We hope our work will help guide future efforts in pruning LLMs in privacy-sensitive fields. Our code is available at https://github.com/Pengxin-Guo/FedPrLLM.

Figures (3)

• Figure 1: Top). Research questions alongside the corresponding findings and experimental scenarios. Bottom). The FedPrLLM framework. ① Each client calculates a pruning mask matrix $\mathbf{M}_i$ using its calibration dataset $\mathcal{D}_i$. ② Clients send the mask matrices $\mathbf{M}_i$ to the server. ③ The server aggregates these mask matrices $\mathbf{M}_i$ to obtain an aggregated mask matrix $\mathbf{\hat{M}} = \sum_{i=1}^{m}\mathbf{M}_i$. ④ Top-k values are selected from the aggregated mask matrix $\mathbf{\hat{W}}$ to derive the final mask matrix $\mathbf{M}$. ⑤ Prune the global model $\mathbf{W}$ using the mask matrix $\mathbf{M}$ as follows: $\mathbf{\hat{W}} = \mathbf{W} \odot (1 - \mathbf{M})$, where $\odot$ denotes element-wise multiplication. ⑥ Scale the model weights of the retained parameters using the aggregated mask matrix $\mathbf{\hat{M}}$ as follows: $\mathbf{\hat{W}} \odot \frac{(m - \mathbf{\hat{M}})}{m}$ (if needed). ⑦ The server broadcasts the mask matrix $\mathbf{M}$ to each client (for iterative pruning). The dashed arrow indicates that this operation is optional; step ⑥ is used for weight scaling, while ⑦ is used for iterative pruning. Note that this visualization is primarily for one-shot pruning, which requires only one communication round. For iterative pruning, multiple communication rounds will occur between steps ② and ⑦, and the layer index is omitted here.
• Figure 2: The effect of different client numbers on federated pruning LLMs.
• Figure 3: The effect of the number of calibration samples on federated pruning LLMs.