CELLM: An Efficient Communication in Large Language Models Training for Federated Learning

TL;DR

This paper introduces CELLM, a method combining low-rank adaptation and sparse communication to efficiently train large language models in federated learning, significantly reducing communication costs while maintaining high utility.

Contribution

CELLM is the first approach to effectively integrate LoRA and sparse updates for federated LLM training, addressing communication and computation bottlenecks.

Findings

Reduces communication costs by up to 10x compared to vanilla LoRA.

Achieves up to 5x reduction over complex sparse LoRA baselines.

Maintains or improves model utility with optimized sparsity and rank configurations.

Abstract

Federated Learning (FL) is a recent model training paradigm in which client devices collaboratively train a model without ever aggregating their data. Crucially, this scheme offers users potential privacy and security benefits by only ever communicating updates to the model weights to a central server as opposed to traditional machine learning (ML) training which directly communicates and aggregates data. However, FL training suffers from statistical heterogeneity as clients may have differing local data distributions. Large language models (LLMs) offer a potential solution to this issue of heterogeneity given that they have consistently been shown to be able to learn on vast amounts of noisy data. While LLMs are a promising development for resolving the consistent issue of non-I.I.D. Clients in federated settings exacerbate two other bottlenecks in FL: limited local computing and expensive communication. This thesis aims to develop efficient training methods for LLMs in FL. To this end, we employ two critical techniques in enabling efficient training. First, we use low-rank adaptation (LoRA) to reduce the computational load of local model training. Second, we communicate sparse updates throughout training to significantly cut down on communication costs. Taken together, our method reduces communication costs by up to 10x over vanilla LoRA and up to 5x over more complex sparse LoRA baselines while achieving greater utility. We emphasize the importance of carefully applying sparsity and picking effective rank and sparsity configurations for federated LLM training.