FedSpaLLM: Federated Pruning of Large Language Models

Introduction

Large Language Models (LLMs) have revolutionized natural language processing but impose significant computational and memory demands. Model pruning—the process of removing redundant weights—is a standard technique for compression. However, state-of-the-art pruning methods typically require access to a calibration dataset to identify weight importance. In many real-world scenarios, such data is private (e.g., healthcare or finance) and cannot be centralized.

FedSpaLLM introduces a Federated Learning (FL) approach to LLM pruning. It allows multiple clients to collaboratively prune a global model using their local, private data. This framework overcomes two primary hurdles: the privacy risks associated with data sharing and the system heterogeneity inherent in distributed networks where clients possess varying hardware capabilities.

Methodology

FedSpaLLM incorporates three primary technical innovations to optimize federated pruning for LLMs.

1. Resource-Aware Layer Sampling

To accommodate clients with limited hardware, FedSpaLLM employs a layer sampling strategy. Rather than requiring every client to process the entire model, the server assigns a subset of layers to each client based on its computational resources. This significantly reduces local memory requirements and communication overhead.

Expand

Source, p. 4

2. $\ell_0$-norm Aggregation

Standard federated averaging can be detrimental to pruned models. If a weight is deemed important by one client but pruned (set to zero) by others, simple averaging dilutes the parameter's magnitude. FedSpaLLM uses $\ell_0$-norm aggregation, which averages only the non-zero values for each parameter across clients, thereby preserving the signal of important weights.

3. Adaptive Mask Expansion

Different local datasets lead to different pruning decisions (masks). To ensure the global model achieves the desired target sparsity while retaining the most critical weights, FedSpaLLM uses a voting-based mask expansion. The server aggregates local masks and prioritizes keeping weights that are consistently identified as important across the client pool.

Expand

Source, p. 5

Experimental Evaluation

The framework was evaluated using OPT (125M, 1.3B) and Llama-2 (7B) models across standard benchmarks including WikiText-2, PTB, and C4.

Key Results

• Performance Gains: FedSpaLLM consistently outperformed standalone pruning. By leveraging the "collective intelligence" of multiple clients, the global model achieves lower perplexity than models pruned on isolated local datasets.
• Resilience to High Sparsity: The framework maintains robust performance even at 80% sparsity. While traditional methods often see a sharp decline in generation quality at these levels, FedSpaLLM's aggregation strategies preserve model utility.
• Scalability: Experiments with varying client counts (up to 16) demonstrated that the method scales effectively, benefiting from increased data diversity in larger federated networks.

Conclusion

FedSpaLLM represents a significant step forward in making LLM compression accessible and privacy-preserving. By combining layer sampling with specialized aggregation techniques, it enables heterogeneous devices to contribute to a high-quality, sparse global model. Future work will explore scaling these techniques to frontier models exceeding 70 billion parameters and addressing extreme data distribution shifts.