FedSpaLLM: Federated Pruning of Large Language Models

TL;DR

FedSpaLLM introduces a federated learning framework for pruning large language models locally on clients' private data, effectively reducing model size while preserving performance and addressing privacy, heterogeneity, and communication challenges.

Contribution

It is the first framework to enable privacy-preserving, federated pruning of LLMs with novel aggregation, adaptive mask expansion, and layer sampling techniques.

Findings

Improves pruning performance across diverse federated settings.

Reduces communication overhead through layer sampling.

Maintains model accuracy while significantly reducing size.

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.