FedSpaLLM: Federated Pruning of Large Language Models

TL;DR

This work tackles the challenge of pruning Large Language Models in privacy-sensitive federated settings. It introduces FedSpaLLM, a framework that enables clients to prune locally on private data and aggregates updates with an ell0-norm based scheme, along with adaptive mask expansion and a layer-sampling strategy to accommodate heterogeneous resources and reduce communication. The method maintains target global sparsity while preserving important weights and achieves unbiased global pruning through layer sampling. Extensive experiments on OPT and LlaMA-2 demonstrate that FedSpaLLM outperforms isolated pruning and random baselines, especially at higher sparsity, while keeping data private. The approach holds practical potential for deploying sparse LLMs in resource-constrained, privacy-aware environments with scalable federated coordination.

Abstract

Large Language Models (LLMs) achieve state-of-the-art performance but are challenging to deploy due to their high computational and storage demands. Pruning can reduce model size, yet existing methods assume public access to calibration data, which is impractical for privacy-sensitive applications. To address the challenge of pruning LLMs in privacy-preserving settings, we propose FedSpaLLM, the first federated learning framework designed specifically for pruning LLMs. FedSpaLLM enables clients to prune their models locally based on private data while accounting for system heterogeneity and maintaining communication efficiency. Our framework introduces several key innovations: (1) a novel $\ell_0$-norm aggregation function that ensures only non-zero weights are averaged across clients, preserving important model parameters; (2) an adaptive mask expansion technique that meets global sparsity targets while accommodating client-specific pruning decisions; and (3) a layer sampling strategy that reduces communication overhead and personalizes the pruning process based on client resources. Extensive experiments show that FedSpaLLM improves pruning performance in diverse federated settings.

Figures (2)

• Figure 1: Visualization of the proposed FedSpaLLM framework. Instead of transmitting the full model at each communication round, the server samples a subset of layers based on each client’s computational resources. Clients prune only the sampled layers and retain the rest from their cached pre-trained dense model. After local pruning, clients only send their pruned layers to the server, which aggregates the pruned layers using a novel $\ell_0$-norm aggregation function that averages only the non-zero parameters. This approach ensures that important weights are preserved while reducing communication overhead. The layer sampling strategy enables personalized pruning tailored to client heterogeneity, reducing resource usage without compromising overall model performance.
• Figure 2: Visualization of the proposed Aggregation Function of FedSpaLLM to handle heterogeneous sparsified parameters. After clients prune their local models, the server aggregates the pruned layers by using the $\ell_0$-norm aggregation function. This method avoids diluting the effect of unpruned weights by excluding zeros from the averaging process, thus preserving important parameters. To achieve the target global sparsity, an adaptive mask expansion is applied: the server counts the number of times each weight has been pruned across clients and uses this information to expand the pruning mask. The mask expansion prioritizes pruning weights that are most commonly pruned across clients, balancing individual client pruning decisions with the global sparsity goal.

Theorems & Definitions (4)

• Corollary 4.1: Sparsity Guarantee
• Theorem 4.2: Unbiased Estimator
• proof
• proof